SkyLi
b4308d6cdc
- 移除对InfraredDetector的依赖，改用TempFeatureExtractor.detect_and_extract() - 添加--device参数支持GPU/CPU切换 - 简化特征提取逻辑，直接使用detect_and_extract()返回的特征 - 添加空结果警告提示
315 lines
10 KiB
Python
315 lines
10 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
验证集目标检测与ROI温度特征直方图分析
|
||
对验证集图像进行YOLO检测，提取ROI温度特征，绘制直方图
|
||
"""
|
||
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
# 添加项目根目录到路径
|
||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||
|
||
import cv2
|
||
import numpy as np
|
||
import matplotlib.pyplot as plt
|
||
from tqdm import tqdm
|
||
from collections import defaultdict
|
||
|
||
from src.extract_temp import TempFeatureExtractor, load_infrared_image
|
||
|
||
|
||
def analyze_validation_set(
|
||
image_dir: str,
|
||
model_path: str,
|
||
conf: float = 0.25,
|
||
output_dir: str = None,
|
||
device: str = None
|
||
):
|
||
"""
|
||
分析验证集，提取ROI温度特征并绘制直方图
|
||
|
||
Args:
|
||
image_dir: 验证集图像目录
|
||
model_path: YOLO模型路径
|
||
conf: 检测置信度阈值
|
||
output_dir: 输出目录
|
||
device: 运行设备，例如 '0' 表示GPU 0，'cpu' 表示CPU
|
||
"""
|
||
image_dir = Path(image_dir)
|
||
if output_dir:
|
||
output_dir = Path(output_dir)
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
else:
|
||
output_dir = Path("output/temp_analysis")
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
# 初始化温度特征提取器（包含检测功能）
|
||
extractor = TempFeatureExtractor(model_path)
|
||
|
||
# 获取所有图像
|
||
image_paths = sorted([
|
||
p for p in image_dir.glob('*')
|
||
if p.suffix.lower() in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']
|
||
])
|
||
|
||
print(f"找到 {len(image_paths)} 张图像")
|
||
|
||
# 收集所有温度特征
|
||
all_features = defaultdict(list)
|
||
class_features = defaultdict(lambda: defaultdict(list))
|
||
|
||
for img_path in tqdm(image_paths, desc="处理图像"):
|
||
try:
|
||
# 加载灰度图像（用于温度提取）
|
||
gray_img = load_infrared_image(str(img_path))
|
||
|
||
# 使用detect_and_extract一次性完成检测和特征提取
|
||
detections = extractor.detect_and_extract(gray_img, conf=conf, device=device)
|
||
|
||
for det in detections:
|
||
# 收集全局特征
|
||
all_features['temp_max'].append(det['temp_max'])
|
||
all_features['temp_min'].append(det['temp_min'])
|
||
all_features['temp_mean'].append(det['temp_mean'])
|
||
all_features['temp_std'].append(det['temp_std'])
|
||
all_features['temp_median'].append(det['temp_median'])
|
||
all_features['gradient_mean'].append(det['gradient_mean'])
|
||
all_features['gradient_std'].append(det['gradient_std'])
|
||
|
||
# 计算方差（标准差的平方）
|
||
variance = det['temp_std'] ** 2
|
||
all_features['temp_variance'].append(variance)
|
||
|
||
# 按类别收集特征
|
||
class_name = det['class_name']
|
||
class_features[class_name]['temp_max'].append(det['temp_max'])
|
||
class_features[class_name]['temp_mean'].append(det['temp_mean'])
|
||
class_features[class_name]['temp_std'].append(det['temp_std'])
|
||
class_features[class_name]['temp_variance'].append(variance)
|
||
|
||
except Exception as e:
|
||
print(f"处理 {img_path.name} 时出错: {e}")
|
||
continue
|
||
|
||
# 打印统计信息
|
||
total_rois = len(all_features['temp_max'])
|
||
print(f"\n总共提取了 {total_rois} 个ROI的温度特征")
|
||
|
||
if total_rois == 0:
|
||
print("警告: 未提取到任何ROI特征，请检查:")
|
||
print(" 1. 模型文件是否存在: ", model_path)
|
||
print(" 2. 图像目录是否正确: ", image_dir)
|
||
print(" 3. 图像中是否包含可检测的目标")
|
||
return
|
||
|
||
# 绘制全局直方图
|
||
draw_global_histograms(all_features, output_dir)
|
||
|
||
# 绘制各类别的温度STD直方图
|
||
draw_class_histograms(class_features, output_dir)
|
||
|
||
# 绘制温度散点图
|
||
draw_temp_scatter(all_features, output_dir)
|
||
|
||
# 保存特征统计到CSV
|
||
save_stats_csv(all_features, class_features, output_dir)
|
||
|
||
print(f"\n分析完成! 结果保存至: {output_dir}")
|
||
|
||
|
||
def draw_global_histograms(features: dict, output_dir: Path):
|
||
"""绘制全局温度特征直方图"""
|
||
fig, axes = plt.subplots(2, 4, figsize=(16, 8))
|
||
fig.suptitle('ROI Temperature Features Distribution (All Classes)', fontsize=14)
|
||
|
||
plot_configs = [
|
||
('temp_max', 'Max Temperature', 'red'),
|
||
('temp_min', 'Min Temperature', 'blue'),
|
||
('temp_mean', 'Mean Temperature', 'green'),
|
||
('temp_std', 'Temperature Std Dev', 'orange'),
|
||
('temp_median', 'Median Temperature', 'purple'),
|
||
('temp_variance', 'Temperature Variance', 'brown'),
|
||
('gradient_mean', 'Gradient Mean', 'pink'),
|
||
('gradient_std', 'Gradient Std', 'gray'),
|
||
]
|
||
|
||
for idx, (key, title, color) in enumerate(plot_configs):
|
||
ax = axes[idx // 4, idx % 4]
|
||
data = features.get(key, [])
|
||
|
||
if data:
|
||
ax.hist(data, bins=30, color=color, alpha=0.7, edgecolor='black')
|
||
ax.axvline(np.mean(data), color='red', linestyle='--',
|
||
linewidth=2, label=f'Mean: {np.mean(data):.2f}')
|
||
ax.axvline(np.median(data), color='blue', linestyle=':',
|
||
linewidth=2, label=f'Median: {np.median(data):.2f}')
|
||
|
||
ax.set_title(title)
|
||
ax.set_xlabel('Value')
|
||
ax.set_ylabel('Count')
|
||
ax.legend(fontsize=8)
|
||
|
||
plt.tight_layout()
|
||
plt.savefig(output_dir / 'global_temp_histograms.png', dpi=150)
|
||
plt.close()
|
||
print(f" 保存: global_temp_histograms.png")
|
||
|
||
|
||
def draw_class_histograms(class_features: dict, output_dir: Path):
|
||
"""绘制各类别的温度标准差直方图"""
|
||
if not class_features:
|
||
return
|
||
|
||
n_classes = len(class_features)
|
||
n_cols = 3
|
||
n_rows = (n_classes + n_cols - 1) // n_cols
|
||
|
||
fig, axes = plt.subplots(n_rows, n_cols, figsize=(15, 4 * n_rows))
|
||
fig.suptitle('Temperature Std Dev Distribution by Class', fontsize=14)
|
||
axes = axes.flatten() if n_classes > 1 else [axes]
|
||
|
||
for idx, (class_name, features) in enumerate(sorted(class_features.items())):
|
||
ax = axes[idx]
|
||
std_data = features.get('temp_std', [])
|
||
|
||
if std_data:
|
||
ax.hist(std_data, bins=20, color='steelblue', alpha=0.7, edgecolor='black')
|
||
ax.axvline(np.mean(std_data), color='red', linestyle='--',
|
||
linewidth=2, label=f'Mean: {np.mean(std_data):.2f}')
|
||
|
||
ax.set_title(f'{class_name} (n={len(std_data)})')
|
||
ax.set_xlabel('Temperature Std Dev')
|
||
ax.set_ylabel('Count')
|
||
ax.legend(fontsize=8)
|
||
|
||
# 隐藏多余的子图
|
||
for idx in range(n_classes, len(axes)):
|
||
axes[idx].axis('off')
|
||
|
||
plt.tight_layout()
|
||
plt.savefig(output_dir / 'class_temp_std_histograms.png', dpi=150)
|
||
plt.close()
|
||
print(f" 保存: class_temp_std_histograms.png")
|
||
|
||
|
||
def draw_temp_scatter(features: dict, output_dir: Path):
|
||
"""绘制温度特征散点图"""
|
||
fig, axes = plt.subplots(1, 2, figsize=(12, 5))
|
||
|
||
# temp_max vs temp_std
|
||
ax = axes[0]
|
||
ax.scatter(features['temp_max'], features['temp_std'],
|
||
alpha=0.3, s=10, c='red')
|
||
ax.set_xlabel('Max Temperature')
|
||
ax.set_ylabel('Temperature Std Dev')
|
||
ax.set_title('Max Temperature vs Std Dev')
|
||
ax.grid(True, alpha=0.3)
|
||
|
||
# temp_mean vs temp_std
|
||
ax = axes[1]
|
||
ax.scatter(features['temp_mean'], features['temp_std'],
|
||
alpha=0.3, s=10, c='blue')
|
||
ax.set_xlabel('Mean Temperature')
|
||
ax.set_ylabel('Temperature Std Dev')
|
||
ax.set_title('Mean Temperature vs Std Dev')
|
||
ax.grid(True, alpha=0.3)
|
||
|
||
plt.tight_layout()
|
||
plt.savefig(output_dir / 'temp_scatter_plots.png', dpi=150)
|
||
plt.close()
|
||
print(f" 保存: temp_scatter_plots.png")
|
||
|
||
|
||
def save_stats_csv(all_features: dict, class_features: dict, output_dir: Path):
|
||
"""保存特征统计到CSV"""
|
||
import csv
|
||
|
||
# 全局统计
|
||
stats_file = output_dir / 'temp_feature_stats.csv'
|
||
with open(stats_file, 'w', newline='', encoding='utf-8') as f:
|
||
writer = csv.writer(f)
|
||
writer.writerow(['Feature', 'Count', 'Mean', 'Std', 'Min', 'Max', 'Median'])
|
||
|
||
for key in ['temp_max', 'temp_min', 'temp_mean', 'temp_std',
|
||
'temp_median', 'temp_variance', 'gradient_mean', 'gradient_std']:
|
||
data = all_features.get(key, [])
|
||
if data:
|
||
writer.writerow([
|
||
key,
|
||
len(data),
|
||
f'{np.mean(data):.2f}',
|
||
f'{np.std(data):.2f}',
|
||
f'{np.min(data):.2f}',
|
||
f'{np.max(data):.2f}',
|
||
f'{np.median(data):.2f}'
|
||
])
|
||
|
||
print(f" 保存: temp_feature_stats.csv")
|
||
|
||
# 各类别统计
|
||
class_stats_file = output_dir / 'class_temp_stats.csv'
|
||
with open(class_stats_file, 'w', newline='', encoding='utf-8') as f:
|
||
writer = csv.writer(f)
|
||
writer.writerow(['Class', 'Count', 'Mean_Max', 'Mean_Mean', 'Mean_Std', 'Mean_Variance'])
|
||
|
||
for class_name in sorted(class_features.keys()):
|
||
features = class_features[class_name]
|
||
writer.writerow([
|
||
class_name,
|
||
len(features['temp_std']),
|
||
f'{np.mean(features["temp_max"]):.2f}',
|
||
f'{np.mean(features["temp_mean"]):.2f}',
|
||
f'{np.mean(features["temp_std"]):.2f}',
|
||
f'{np.mean(features["temp_variance"]):.2f}'
|
||
])
|
||
|
||
print(f" 保存: class_temp_stats.csv")
|
||
|
||
|
||
if __name__ == '__main__':
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser(
|
||
description='验证集ROI温度特征提取与直方图分析'
|
||
)
|
||
parser.add_argument(
|
||
'--image-dir',
|
||
type=str,
|
||
default='data/valid/images',
|
||
help='验证集图像目录'
|
||
)
|
||
parser.add_argument(
|
||
'--model',
|
||
type=str,
|
||
default='weights/best.pt',
|
||
help='YOLO模型路径'
|
||
)
|
||
parser.add_argument(
|
||
'--conf',
|
||
type=float,
|
||
default=0.25,
|
||
help='检测置信度阈值'
|
||
)
|
||
parser.add_argument(
|
||
'--output',
|
||
type=str,
|
||
default='output/temp_analysis',
|
||
help='输出目录'
|
||
)
|
||
parser.add_argument(
|
||
'--device',
|
||
type=str,
|
||
default=None,
|
||
help='运行设备，例如 "0" 表示GPU 0，"cpu" 表示CPU'
|
||
)
|
||
|
||
args = parser.parse_args()
|
||
|
||
analyze_validation_set(
|
||
image_dir=args.image_dir,
|
||
model_path=args.model,
|
||
conf=args.conf,
|
||
output_dir=args.output,
|
||
device=args.device
|
||
) |