diff --git a/scripts/usage_example.py b/scripts/usage_example.py index c5bb052..60ef3e5 100644 --- a/scripts/usage_example.py +++ b/scripts/usage_example.py @@ -8,9 +8,7 @@ from pathlib import Path # 添加项目根目录到路径 sys.path.insert(0, str(Path(__file__).parent.parent)) -from src.detect import InfraredDetector -from src.extract_temp import TempFeatureExtractor -from src.diagnose import FaultDiagnoser +from src.infrared_data import InfraredDetector, TempFeatureExtractor, FaultDiagnoser import cv2 import json @@ -48,7 +46,7 @@ def example3_temperature_extraction(): """示例3：温度特征提取""" print("\n【示例3】温度特征提取") - extractor = TempFeatureExtractor("weights/best.pt") + extractor = TempFeatureExtractor() detector = InfraredDetector("weights/best.pt") # 读取图像 diff --git a/src/infrared_data/detect.py b/src/infrared_data/detect.py index a0f0ae8..2e9c376 100644 --- a/src/infrared_data/detect.py +++ b/src/infrared_data/detect.py @@ -3,7 +3,6 @@ YOLO目标检测脚本 用于红外图像中的目标检测 """ -import os from pathlib import Path from typing import Optional, List, Dict, Union import numpy as np @@ -117,66 +116,6 @@ class InfraredDetector: return img, len(detections) -def train_yolo_model(data_yaml: str, - model_size: str = 'n', - epochs: int = 100, - imgsz: int = 640, - batch: int = 16, - output_dir: str = 'models/train') -> str: - """ - 训练YOLO模型 - - Args: - data_yaml: 数据集配置文件路径 - model_size: 模型大小 ('n', 's', 'm', 'l', 'x') - epochs: 训练轮数 - imgsz: 输入图像大小 - batch: 批量大小 - output_dir: 输出目录 - - Returns: - 最佳模型路径 - """ - model = YOLO(f'yolov8{model_size}.pt') - - results = model.train( - data=data_yaml, - epochs=epochs, - imgsz=imgsz, - batch=batch, - project=output_dir, - name='infrared_fault', - exist_ok=True, - patience=20, - save=True, - plots=True - ) - - best_model_path = f"{output_dir}/infrared_fault/weights/best.pt" - return best_model_path - - -def validate_model(model_path: str, data_yaml: str) -> Dict: - """ - 验证模型性能 - - Args: - model_path: 模型路径 - data_yaml: 数据集配置路径 - - Returns: - 验证指标 - """ - model = YOLO(model_path) - metrics = model.val(data=data_yaml) - return { - 'mAP50': float(metrics.box.map50), - 'mAP50-95': float(metrics.box.map), - 'precision': float(metrics.box.mp), - 'recall': float(metrics.box.mr) - } - - if __name__ == '__main__': import argparse @@ -191,23 +130,9 @@ if __name__ == '__main__': detect_parser.add_argument('--conf', type=float, default=0.25, help='Confidence threshold') detect_parser.add_argument('--device', type=str, default=None, help='Device (e.g., "0" for GPU, "cpu" for CPU)') - # Train command - train_parser = subparsers.add_parser('train', help='Train YOLO model') - train_parser.add_argument('--data', type=str, required=True, help='Data yaml path') - train_parser.add_argument('--model-size', type=str, default='n', choices=['n', 's', 'm', 'l', 'x']) - train_parser.add_argument('--epochs', type=int, default=100) - train_parser.add_argument('--imgsz', type=int, default=640) - train_parser.add_argument('--batch', type=int, default=16) - args = parser.parse_args() if args.command == 'detect': detector = InfraredDetector(args.model, args.conf) img, n_detections = detector.detect_and_draw(args.image, args.output, device=args.device) print(f"Detected {n_detections} objects") - - elif args.command == 'train': - best_path = train_yolo_model( - args.data, args.model_size, args.epochs, args.imgsz, args.batch - ) - print(f"Training complete. Best model: {best_path}") diff --git a/src/infrared_data/diagnose.py b/src/infrared_data/diagnose.py index 9f8e187..524d09d 100644 --- a/src/infrared_data/diagnose.py +++ b/src/infrared_data/diagnose.py @@ -26,7 +26,7 @@ class FaultDiagnoser: temp_threshold: 温度阈值配置 """ self.detector = InfraredDetector(model_path) if model_path else None - self.extractor = TempFeatureExtractor(model_path) + self.extractor = TempFeatureExtractor() self.temp_threshold = temp_threshold or {} def diagnose(self, image: Union[str, np.ndarray], diff --git a/src/infrared_data/extract_temp.py b/src/infrared_data/extract_temp.py index 98ff5b0..6cfb17c 100644 --- a/src/infrared_data/extract_temp.py +++ b/src/infrared_data/extract_temp.py @@ -3,83 +3,26 @@ ROI温度特征提取模块 基于YOLO检测结果，提取ROI区域的温度特征 """ -import numpy as np -import cv2 import json from pathlib import Path -from typing import Dict, List, Optional, Union -from ultralytics import YOLO +from typing import Dict, List, Union + +import cv2 +import numpy as np from .utils.temp_utils import ( - extract_temp_stats, - extract_temp_histogram, calculate_temp_gradient, - normalize_roi + extract_temp_histogram, + extract_temp_stats, ) class TempFeatureExtractor: """红外图像ROI温度特征提取器""" - def __init__(self, model_path: Optional[str] = None): - """ - 初始化温度特征提取器 - - Args: - model_path: YOLO模型路径，如果为None则只提取温度特征 - """ - self.model = None - if model_path and Path(model_path).exists(): - self.model = YOLO(model_path) - - def detect_and_extract(self, image: np.ndarray, - conf: float = 0.25, - device: Optional[str] = None) -> List[Dict]: - """ - 对图像进行检测并提取ROI温度特征 - - Args: - image: 输入红外图像 - conf: 检测置信度阈值 - device: 运行设备，例如 '0' 表示GPU 0，'cpu' 表示CPU - - Returns: - 检测结果列表，每个元素包含检测框和温度特征 - """ - if self.model is None: - raise ValueError("Model not loaded. Please provide a valid model path.") - - results = self.model(image, conf=conf, device=device) - detections = [] - - for result in results: - boxes = result.boxes - if boxes is None: - continue - - for box in boxes: - xyxy = box.xyxy[0].cpu().numpy() - conf_score = float(box.conf[0]) - cls_id = int(box.cls[0]) - - x1, y1, x2, y2 = map(int, xyxy) - roi = image[y1:y2, x1:x2] - - temp_features = extract_temp_stats(roi) - temp_features.update({ - 'bbox': [int(x1), int(y1), int(x2), int(y2)], - 'confidence': conf_score, - 'class_id': cls_id, - 'class_name': result.names[cls_id] - }) - - detections.append(temp_features) - - return detections - - def extract_from_roi(self, roi: np.ndarray, - include_hist: bool = True, - include_gradient: bool = True) -> Dict: + def extract_from_roi( + self, roi: np.ndarray, include_hist: bool = True, include_gradient: bool = True + ) -> Dict: """ 从指定ROI提取温度特征 @@ -94,15 +37,14 @@ class TempFeatureExtractor: features = extract_temp_stats(roi) if include_hist: - features['histogram'] = extract_temp_histogram(roi).tolist() + features["histogram"] = extract_temp_histogram(roi).tolist() if include_gradient: features.update(calculate_temp_gradient(roi)) return features - def extract_from_bbox(self, image: np.ndarray, - bbox: List[int]) -> Dict: + def extract_from_bbox(self, image: np.ndarray, bbox: List[int]) -> Dict: """ 从给定边界框提取温度特征 @@ -134,56 +76,53 @@ def load_infrared_image(image_path: str) -> np.ndarray: return image -def batch_extract_features(image_dir: Union[str, Path], - output_path: Union[str, Path], - model_path: Optional[str] = None, - conf: float = 0.25, - device: Optional[str] = None) -> None: +def batch_extract_from_detections( + image_dir: Union[str, Path], + detections_list: List[Dict], + output_path: Union[str, Path], +) -> None: """ - 批量处理图像并提取温度特征 + 批量处理已检测的结果，提取温度特征 Args: image_dir: 图像目录路径 + detections_list: 检测结果列表 output_path: 输出JSON文件路径 - model_path: YOLO模型路径 - conf: 检测置信度阈值 - device: 运行设备，例如 '0' 表示GPU 0，'cpu' 表示CPU """ image_dir = Path(image_dir) - extractor = TempFeatureExtractor(model_path) + extractor = TempFeatureExtractor() all_results = [] - for img_path in sorted(image_dir.glob('*')): - if img_path.suffix.lower() not in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']: + for det_result in detections_list: + img_name = det_result.get("image", "") + img_path = image_dir / img_name + if not img_path.exists(): continue try: image = load_infrared_image(str(img_path)) - detections = extractor.detect_and_extract(image, conf=conf, device=device) - - all_results.append({ - 'image': img_path.name, - 'detections': detections - }) + for det in det_result.get("detections", []): + bbox = det.get("bbox") + if bbox is None: + continue + x1, y1, x2, y2 = bbox + roi = image[y1:y2, x1:x2] + temp_features = extractor.extract_from_roi(roi) + temp_features.update( + { + "bbox": bbox, + "class_name": det.get("class_name"), + "confidence": det.get("confidence"), + } + ) + all_results.append( + {"image": img_name, "detections": det_result.get("detections", [])} + ) except Exception as e: - print(f"Error processing {img_path.name}: {e}") + print(f"Error processing {img_name}: {e}") continue - with open(output_path, 'w', encoding='utf-8') as f: + with open(output_path, "w", encoding="utf-8") as f: json.dump(all_results, f, indent=2, ensure_ascii=False) - print(f"Processed {len(all_results)} images, results saved to {output_path}") - - -if __name__ == '__main__': - import argparse - - parser = argparse.ArgumentParser(description='Extract temperature features from infrared images') - parser.add_argument('--image-dir', type=str, required=True, help='Input image directory') - parser.add_argument('--output', type=str, required=True, help='Output JSON path') - parser.add_argument('--model', type=str, default=None, help='YOLO model path') - parser.add_argument('--conf', type=float, default=0.25, help='Detection confidence threshold') - parser.add_argument('--device', type=str, default=None, help='Device (e.g., "0" for GPU, "cpu" for CPU)') - - args = parser.parse_args() - batch_extract_features(args.image_dir, args.output, args.model, args.conf, args.device) + print(f"Processed {len(all_results)} images, results saved to {output_path}") \ No newline at end of file