#!/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.detect import InfraredDetector 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 ): """ 分析验证集，提取ROI温度特征并绘制直方图 Args: image_dir: 验证集图像目录 model_path: YOLO模型路径 conf: 检测置信度阈值 output_dir: 输出目录 """ 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) # 初始化检测器和温度特征提取器 detector = InfraredDetector(model_path, conf=conf) 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)) # 进行检测（复用已加载的图像） detections = detector.detect(gray_img, conf=conf) for det in detections: bbox = det['bbox'] x1, y1, x2, y2 = map(int, bbox) # 确保边界在图像范围内 x1, y1 = max(0, x1), max(0, y1) x2, y2 = min(gray_img.shape[1], x2), min(gray_img.shape[0], y2) if x2 - x1 <= 0 or y2 - y1 <= 0: continue roi = gray_img[y1:y2, x1:x2] if roi.size == 0: continue # 提取温度特征 features = extractor.extract_from_roi(roi) # 收集全局特征 all_features['temp_max'].append(features['temp_max']) all_features['temp_min'].append(features['temp_min']) all_features['temp_mean'].append(features['temp_mean']) all_features['temp_std'].append(features['temp_std']) all_features['temp_median'].append(features['temp_median']) all_features['gradient_mean'].append(features['gradient_mean']) all_features['gradient_std'].append(features['gradient_std']) # 计算方差（标准差的平方） variance = features['temp_std'] ** 2 all_features['temp_variance'].append(variance) # 按类别收集特征 class_name = det['class_name'] class_features[class_name]['temp_max'].append(features['temp_max']) class_features[class_name]['temp_mean'].append(features['temp_mean']) class_features[class_name]['temp_std'].append(features['temp_std']) class_features[class_name]['temp_variance'].append(variance) except Exception as e: print(f"处理 {img_path.name} 时出错: {e}") continue # 打印统计信息 print(f"\n总共提取了 {len(all_features['temp_max'])} 个ROI的温度特征") # 绘制全局直方图 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='输出目录' ) args = parser.parse_args() analyze_validation_set( image_dir=args.image_dir, model_path=args.model, conf=args.conf, output_dir=args.output )