Infrared-ai-research-lab/src/diagnose.py

"""
故障诊断模块
基于检测结果和温度特征进行故障诊断
"""

import json
from typing import Dict, List, Optional, Union
from pathlib import Path
import numpy as np

from .extract_temp import TempFeatureExtractor
from .detect import InfraredDetector


class FaultDiagnoser:
    """故障诊断器"""

    def __init__(self,
                 model_path: Optional[str] = None,
                 temp_threshold: Optional[Dict] = None):
        """
        初始化诊断器

        Args:
            model_path: YOLO模型路径
            temp_threshold: 温度阈值配置
        """
        self.detector = InfraredDetector(model_path) if model_path else None
        self.extractor = TempFeatureExtractor(model_path)
        self.temp_threshold = temp_threshold or {}

    def diagnose(self, image: Union[str, np.ndarray],
                 device: Optional[str] = None) -> Dict:
        """
        对图像进行故障诊断

        Args:
            image: 图像路径或数组
            device: 运行设备，例如 '0' 表示GPU 0，'cpu' 表示CPU

        Returns:
            诊断结果
        """
        result = {
            'image': str(image) if isinstance(image, str) else 'array',
            'detections': [],
            'faults': [],
            'summary': {}
        }

        if self.detector and self.detector.model:
            detections = self.detector.detect(image, device=device)
            result['detections'] = detections

            import cv2
            if isinstance(image, str):
                img = cv2.imread(image, cv2.IMREAD_GRAYSCALE)
            else:
                img = image

            for det in detections:
                bbox = det['bbox']
                x1, y1, x2, y2 = map(int, bbox)

                roi = img[y1:y2, x1:x2]
                temp_features = self.extractor.extract_from_roi(roi)

                fault_status, fault_type = self._analyze_fault(temp_features)
                temp_features['fault_status'] = fault_status
                temp_features['fault_type'] = fault_type

                result['faults'].append(temp_features)
        else:
            raise ValueError("Detector not initialized with model")

        result['summary'] = self._generate_summary(result['faults'])
        return result

    def _analyze_fault(self, temp_features: Dict) -> tuple:
        """
        分析温度特征，判断故障状态和类型

        Args:
            temp_features: 温度特征

        Returns:
            (fault_status, fault_type) 元组
        """
        max_temp = temp_features.get('temp_max', 0)
        std_temp = temp_features.get('temp_std', 0)
        temp_range = max_temp - temp_features.get('temp_mean', 0)

        max_threshold = self.temp_threshold.get('max_temp', 200)
        std_threshold = self.temp_threshold.get('std_temp', 50)

        # 判断状态：两项都不超标才为 normal
        if max_temp <= max_threshold and std_temp <= std_threshold:
            return ('normal', 'none')

        # 分类故障类型
        if std_temp > 30:
            return ('abnormal', 'hotspot')
        elif temp_range > 50:
            return ('abnormal', 'local_overheating')
        elif max_temp > 180:
            return ('abnormal', 'overheat')
        else:
            return ('abnormal', 'suspect')

    def _generate_summary(self, faults: List[Dict]) -> Dict:
        """
        生成诊断摘要

        Args:
            faults: 故障列表

        Returns:
            摘要信息
        """
        if not faults:
            return {'status': 'no_detection', 'detection_count': 0}

        abnormal_count = sum(1 for f in faults if f.get('fault_status') == 'abnormal')
        fault_types = [f.get('fault_type', 'unknown') for f in faults]

        return {
            'status': 'fault_detected' if abnormal_count > 0 else 'normal',
            'detection_count': len(faults),
            'abnormal_count': abnormal_count,
            'fault_types': list(set(fault_types))
        }

    def batch_diagnose(self, image_dir: Union[str, Path],
                       output_path: Union[str, Path],
                       device: Optional[str] = None) -> None:
        """
        批量诊断

        Args:
            image_dir: 图像目录
            output_path: 输出JSON路径
            device: 运行设备，例如 '0' 表示GPU 0，'cpu' 表示CPU
        """
        image_dir = Path(image_dir)
        results = []

        for img_path in sorted(image_dir.glob('*')):
            if img_path.suffix.lower() not in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']:
                continue

            try:
                result = self.diagnose(str(img_path), device=device)
                result['image'] = img_path.name
                results.append(result)
            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(results, f, indent=2, ensure_ascii=False)

        print(f"Batch diagnosis complete. {len(results)} images processed.")
        print(f"Results saved to: {output_path}")


def load_threshold_config(config_path: str) -> Dict:
    """加载温度阈值配置"""
    with open(config_path, 'r', encoding='utf-8') as f:
        return json.load(f)


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(description='Fault Diagnosis')
    parser.add_argument('--model', type=str, required=True, help='YOLO model path')
    parser.add_argument('--image-dir', type=str, required=True, help='Image directory')
    parser.add_argument('--output', type=str, required=True, help='Output JSON path')
    parser.add_argument('--threshold', type=str, help='Threshold config JSON')
    parser.add_argument('--device', type=str, default=None, help='Device (e.g., "0" for GPU, "cpu" for CPU)')

    args = parser.parse_args()

    threshold = None
    if args.threshold:
        threshold = load_threshold_config(args.threshold)

    diagnoser = FaultDiagnoser(args.model, threshold)
    diagnoser.batch_diagnose(args.image_dir, args.output, device=args.device)