docs: remove obsolete dataset download scripts from README

.gitignore

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

pyproject.toml

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

scripts/usage_example.py

@@ -8,7 +8,9 @@ from pathlib import Path
 # 添加项目根目录到路径
 sys.path.insert(0, str(Path(__file__).parent.parent))
 
-from src.infrared_data import InfraredDetector, TempFeatureExtractor, FaultDiagnoser
+from src.detect import InfraredDetector
+from src.extract_temp import TempFeatureExtractor
+from src.diagnose import FaultDiagnoser
 import cv2
 import json
 
@@ -46,7 +48,7 @@ def example3_temperature_extraction():
     """示例3：温度特征提取"""
     print("\n【示例3】温度特征提取")
 
-    extractor = TempFeatureExtractor()
+    extractor = TempFeatureExtractor("weights/best.pt")
     detector = InfraredDetector("weights/best.pt")
 
     # 读取图像

src/__init__.py

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

src/detect.py

@@ -3,6 +3,7 @@ YOLO目标检测脚本
 用于红外图像中的目标检测
 """
 
+import os
 from pathlib import Path
 from typing import Optional, List, Dict, Union
 import numpy as np
@@ -116,6 +117,66 @@ 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
 
@@ -130,9 +191,23 @@ 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/diagnose.py

@@ -26,7 +26,7 @@ class FaultDiagnoser:
             temp_threshold: 温度阈值配置
         """
         self.detector = InfraredDetector(model_path) if model_path else None
-        self.extractor = TempFeatureExtractor()
+        self.extractor = TempFeatureExtractor(model_path)
         self.temp_threshold = temp_threshold or {}
 
     def diagnose(self, image: Union[str, np.ndarray],

src/extract_temp.py

@@ -0,0 +1,189 @@
+"""
+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

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

src/infrared_data/extract_temp.py

@@ -1,128 +0,0 @@
-"""
-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}")