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]) -> Dict:
        """
        对图像进行故障诊断

        Args:
            image: 图像路径或数组

        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)
            result['detections'] = detections

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

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

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

                fault_status = self._check_fault(temp_features)
                temp_features['fault_status'] = fault_status
                temp_features['fault_type'] = self._classify_fault(temp_features)

                result['faults'].append({
                    'bbox': bbox,
                    'confidence': det['confidence'],
                    **temp_features
                })
        else:
            raise ValueError("Detector not initialized with model")

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

    def _check_fault(self, temp_features: Dict) -> str:
        """
        根据温度特征判断是否有故障

        Args:
            temp_features: 温度特征

        Returns:
            'normal' 或 'abnormal'
        """
        max_temp = temp_features.get('temp_max', 0)
        mean_temp = temp_features.get('temp_mean', 0)
        std_temp = temp_features.get('temp_std', 0)

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

        if max_temp > max_threshold or std_temp > std_threshold:
            return 'abnormal'
        return 'normal'

    def _classify_fault(self, temp_features: Dict) -> str:
        """
        故障分类

        Args:
            temp_features: 温度特征

        Returns:
            故障类型
        """
        max_temp = temp_features.get('temp_max', 0)
        mean_temp = temp_features.get('temp_mean', 0)
        std_temp = temp_features.get('temp_std', 0)

        temp_range = max_temp - mean_temp

        if std_temp > 30:
            return 'hotspot'
        elif temp_range > 50:
            return 'local_overheating'
        elif max_temp > 180:
            return 'overheat'
        else:
            return 'suspect'

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

        Args:
            faults: 故障列表

        Returns:
            摘要信息
        """
        if not faults:
            return {'status': 'no_detection', 'fault_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',
            'fault_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]) -> None:
        """
        批量诊断

        Args:
            image_dir: 图像目录
            output_path: 输出JSON路径
        """
        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))
                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')

    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)