chore: adjust pyproject.toml comment positions

- Remove detect_and_extract() method from TempFeatureExtractor
  (moved to FaultDiagnoser pipeline)
- Remove YOLO model loading from TempFeatureExtractor.__init__
- Remove duplicate train_yolo_model/validate_model from detect.py
  (already exist in scripts/train.py)
- Fix usage_example.py to use new import path and TempFeatureExtractor API

.gitignore

@@ -5,6 +5,7 @@ __pycache__/
 dist/
 build/
 *.egg
+*.whl
 
 # PyTorch
 *.pth

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "infrared-data"
-version = "0.1.0"
+version = "0.2.0"
 description = "Infrared Equipment Detection based on YOLO - New Dataset"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -18,6 +18,10 @@ dependencies = [
     "ultralytics>=8.4.27",
 ]
 
+
+[tool.setuptools.packages]
+find = {where = ["src"]}
+
 # PyTorch CUDA版本配置
 [tool.uv]
 # 指定PyTorch使用CUDA索引

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")
 
     # 读取图像

src/__init__.py

@@ -1 +0,0 @@
-# Infrared Fault Detection Package

src/extract_temp.py

@@ -1,189 +0,0 @@
-"""
-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 .utils.temp_utils import (
-    extract_temp_stats,
-    extract_temp_histogram,
-    calculate_temp_gradient,
-    normalize_roi
-)
-
-
-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:
-        """
-        从指定ROI提取温度特征
-
-        Args:
-            roi: ROI区域图像
-            include_hist: 是否包含直方图特征
-            include_gradient: 是否包含梯度特征
-
-        Returns:
-            温度特征字典
-        """
-        features = extract_temp_stats(roi)
-
-        if include_hist:
-            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:
-        """
-        从给定边界框提取温度特征
-
-        Args:
-            image: 输入图像
-            bbox: [x1, y1, x2, y2] 边界框坐标
-
-        Returns:
-            温度特征字典
-        """
-        x1, y1, x2, y2 = bbox
-        roi = image[y1:y2, x1:x2]
-        return self.extract_from_roi(roi)
-
-
-def load_infrared_image(image_path: str) -> np.ndarray:
-    """
-    加载红外图像
-
-    Args:
-        image_path: 图像路径
-
-    Returns:
-        图像数组
-    """
-    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
-    if image is None:
-        raise ValueError(f"Failed to load image: {image_path}")
-    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:
-    """
-    批量处理图像并提取温度特征
-
-    Args:
-        image_dir: 图像目录路径
-        output_path: 输出JSON文件路径
-        model_path: YOLO模型路径
-        conf: 检测置信度阈值
-        device: 运行设备，例如 '0' 表示GPU 0，'cpu' 表示CPU
-    """
-    image_dir = Path(image_dir)
-    extractor = TempFeatureExtractor(model_path)
-
-    all_results = []
-    for img_path in sorted(image_dir.glob('*')):
-        if img_path.suffix.lower() not in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']:
-            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
-            })
-        except Exception as e:
-            print(f"Error processing {img_path.name}: {e}")
-            continue
-
-    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)

src/infrared_data/__init__.py

@@ -0,0 +1,6 @@
+# Infrared Fault Detection Package
+from .diagnose import FaultDiagnoser
+from .detect import InfraredDetector
+from .extract_temp import TempFeatureExtractor
+
+__all__ = ['FaultDiagnoser', 'InfraredDetector', 'TempFeatureExtractor']

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}")

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],

src/infrared_data/extract_temp.py

@@ -0,0 +1,128 @@
+"""
+ROI温度特征提取模块
+基于YOLO检测结果，提取ROI区域的温度特征
+"""
+
+import json
+from pathlib import Path
+from typing import Dict, List, Union
+
+import cv2
+import numpy as np
+
+from .utils.temp_utils import (
+    calculate_temp_gradient,
+    extract_temp_histogram,
+    extract_temp_stats,
+)
+
+
+class TempFeatureExtractor:
+    """红外图像ROI温度特征提取器"""
+
+    def extract_from_roi(
+        self, roi: np.ndarray, include_hist: bool = True, include_gradient: bool = True
+    ) -> Dict:
+        """
+        从指定ROI提取温度特征
+
+        Args:
+            roi: ROI区域图像
+            include_hist: 是否包含直方图特征
+            include_gradient: 是否包含梯度特征
+
+        Returns:
+            温度特征字典
+        """
+        features = extract_temp_stats(roi)
+
+        if include_hist:
+            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:
+        """
+        从给定边界框提取温度特征
+
+        Args:
+            image: 输入图像
+            bbox: [x1, y1, x2, y2] 边界框坐标
+
+        Returns:
+            温度特征字典
+        """
+        x1, y1, x2, y2 = bbox
+        roi = image[y1:y2, x1:x2]
+        return self.extract_from_roi(roi)
+
+
+def load_infrared_image(image_path: str) -> np.ndarray:
+    """
+    加载红外图像
+
+    Args:
+        image_path: 图像路径
+
+    Returns:
+        图像数组
+    """
+    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
+    if image is None:
+        raise ValueError(f"Failed to load image: {image_path}")
+    return image
+
+
+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文件路径
+    """
+    image_dir = Path(image_dir)
+    extractor = TempFeatureExtractor()
+
+    all_results = []
+    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))
+            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_name}: {e}")
+            continue
+
+    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}")