# 红外设备检测系统 (Infrared Equipment Detection System)

基于YOLO的红外热图像目标检测系统，用于识别电力设备。

## 项目概述

本项目使用YOLOv8实现红外热图像中的设备目标检测。

### 主要功能

- **YOLO目标检测**: 识别红外图像中的电力设备
- **ROI温度提取**: 从检测区域提取温度统计特征
- **故障诊断**: 基于温度特征判断设备状态

### 设备类别

请根据你的数据集配置类别（修改dataset.yaml中的names字段）

---

## 项目结构

```
2BaseOnYOLO_InfraredData/
├── src/                        # 核心源代码
│   ├── detect.py              # YOLO目标检测模块
│   ├── extract_temp.py        # ROI温度特征提取
│   ├── diagnose.py             # 故障诊断Pipeline
│   └── utils/
│       └── temp_utils.py      # 温度工具函数
├── scripts/                    # 工具脚本
│   ├── train.py              # 模型训练/验证入口
│   ├── test_cuda.py          # CUDA环境测试
│   ├── convert_to_yolo.py     # 数据集格式转换
│   ├── usage_example.py       # 使用示例
│   └── analyze_temp_features.py # ROI温度特征分析
├── data/                      # 数据目录
│   ├── dataset.yaml          # YOLO数据集配置
│   ├── train/                # 训练集
│   │   ├── images/           # 训练图像
│   │   └── labels/           # 训练标签
│   ├── valid/                # 验证集
│   │   ├── images/           # 验证图像
│   │   └── labels/           # 验证标签
│   ├── test/                 # 测试集
│   │   ├── images/           # 测试图像
│   │   └── labels/           # 测试标签
│   └── 标签.txt              # 类别标签说明
├── weights/                   # 模型权重
│   └── best.pt              # 训练好的模型
├── pyproject.toml            # uv项目配置
└── README.md                  # 项目文档
```

---

## 环境配置

### 系统要求

- Python 3.10+
- NVIDIA GPU (RTX 4060 8GB)
- CUDA 13.0+
- cuDNN 9.0+

### 安装依赖

```bash
# 使用uv安装依赖
uv sync

# 或手动安装
uv pip install -e .
```

### GPU环境测试

```bash
uv run python scripts/test_cuda.py
```

**预期输出**:
```
CUDA available: True
GPU: NVIDIA GeForce RTX 4060 Laptop GPU
cuDNN enabled: True
cuDNN version: 9.1.0
```

---

## 数据集

### 数据来源

电力设备红外图像数据集（12类）

### 数据规模

| 类别 | 类别ID |
|------|--------|
| bushing | 0 |
| arrester | 1 |
| breaker | 2 |
| clamp | 3 |
| conservator | 4 |
| current-transformer | 5 |
| disconnector | 6 |
| disconnector2 | 7 |
| disconnector3 | 8 |
| heat-sink | 9 |
| insulator | 10 |
| transformer | 11 |
| **总计** | **12类** |

### 数据划分

- 训练集: 6735张
- 验证集: 765张

---

## 模型训练

### 完整训练流程

**训练模型**
```bash
uv run python scripts/train.py train \
    --data data/dataset.yaml \
    --model-size n \
    --epochs 50 \
    --device 0
```

### 训练参数说明

| 参数 | 说明 | 默认值 |
|------|------|--------|
| `--data` | 数据集配置文件路径 | 必填 |
| `--model-size` | 模型大小 (n/s/m/l/x) | n |
| `--epochs` | 训练轮数 | 100 |
| `--imgsz` | 输入图像尺寸 | 640 |
| `--batch` | 批量大小 | 16 |
| `--device` | 训练设备 (0=cuda, cpu=cpu) | 0 |
| `--output` | 输出目录 | models/train |

### 模型性能

训练完成的YOLOv8n模型性能 (46 epochs):

| 指标 | 值 |
|------|-----|
| mAP@0.5 | **76.42%** |
| mAP@0.5:0.95 | **61.64%** |
| Precision | **66.58%** |
| Recall | **87.93%** |

**各类别性能**:

| 类别 | Precision | Recall | mAP@0.5 |
|------|----------|--------|---------|
| bushing | 0.925 | 0.990 | 0.994 |
| arrester | 0.930 | 0.980 | 0.991 |
| breaker | 0.797 | 0.910 | 0.898 |
| clamp | 0.781 | 0.973 | 0.918 |
| conservator | 0.546 | 0.823 | 0.718 |
| current-transformer | 0.856 | 0.977 | 0.951 |
| disconnector | 0.651 | 0.975 | 0.813 |
| disconnector2 | 0.306 | 0.767 | 0.286 |
| disconnector3 | 0.566 | 0.940 | 0.886 |
| heat-sink | 0.612 | 0.429 | 0.546 |
| insulator | 0.731 | 0.966 | 0.880 |
| transformer | 0.289 | 0.822 | 0.290 |

---

## 使用方法

### 快速开始

```bash
# 运行使用示例
uv run python scripts/usage_example.py
```

### 1. 目标检测

```python
from src.detect import InfraredDetector

# 初始化检测器
detector = InfraredDetector("weights/best.pt", conf=0.25)

# 检测图像
detections = detector.detect("path/to/image.jpg")

# 打印检测结果
for det in detections:
    print(f"{det['class_name']}: conf={det['confidence']:.3f}")
```

**检测结果格式**:
```python
{
    'bbox': [x1, y1, x2, y2],           # 边界框坐标
    'bbox_xywh': [x_center, y_center, w, h],  # 中心格式
    'confidence': 0.995,
    'class_id': 0,
    'class_name': 'bushing'
}
```

### 2. 检测并绘制结果

```python
detector.detect_and_draw(
    "path/to/image.jpg",
    output_path="result.jpg"
)
```

### 3. 温度特征提取

```python
from src.extract_temp import TempFeatureExtractor
import cv2

extractor = TempFeatureExtractor("weights/best.pt")

# 读取红外图像
gray_img = cv2.imread("path/to/image.jpg", cv2.IMREAD_GRAYSCALE)

# 提取ROI温度特征
roi = gray_img[y1:y2, x1:x2]  # 从检测框获取ROI
features = extractor.extract_from_roi(roi)

print(f"最高温度: {features['temp_max']}")
print(f"平均温度: {features['temp_mean']}")
print(f"温度标准差: {features['temp_std']}")
```

**温度特征说明**:

| 特征 | 说明 |
|------|------|
| temp_max | 最高温度值 |
| temp_min | 最低温度值 |
| temp_mean | 平均温度值 |
| temp_std | 温度标准差 |
| temp_median | 温度中位数 |
| gradient_mean | 温度梯度均值 |
| gradient_std | 温度梯度标准差 |
| histogram | 温度分布直方图 (32 bins) |

### 4. 故障诊断

```python
from src.diagnose import FaultDiagnoser

diagnoser = FaultDiagnoser("weights/best.pt")
result = diagnoser.diagnose("path/to/image.jpg")

print(f"诊断状态: {result['summary']['status']}")
print(f"故障类型: {result['summary']['fault_types']}")
```

**诊断输出格式**:
```json
{
  "image": "20230907095711003_jpg.rf.0e14e8b4344dce21114683dc3f3e7e2e.jpg",
  "detections": [
    {
      "bbox": [x1, y1, x2, y2],
      "confidence": 0.995,
      "class_id": 0,
      "class_name": "bushing"
    }
  ],
  "faults": [
    {
      "temp_max": 255.0,
      "temp_mean": 97.49,
      "temp_std": 45.2,
      "fault_status": "abnormal",
      "fault_type": "hotspot"
    }
  ],
  "summary": {
    "status": "fault_detected",
    "detection_count": 1,
    "abnormal_count": 1,
    "fault_types": ["hotspot"]
  }
}
```

**故障类型分类规则**:

| 故障类型 | 条件 | 说明 |
|----------|------|------|
| hotspot | std_temp > 30 | 热点故障 |
| local_overheating | temp_range > 50 | 局部过热 |
| overheat | max_temp > 180 | 过热 |
| suspect | 其他 | 可疑 |

### 5. 批量诊断

```python
diagnoser.batch_diagnose(
    "data/valid/images",
    "/tmp/batch_results.json"
)
```

### 6. 自定义温度阈值

```python
threshold = {
    "max_temp": 200,   # 最高温度阈值
    "std_temp": 40     # 标准差阈值
}

diagnoser = FaultDiagnoser(
    "weights/best.pt",
    temp_threshold=threshold
)
```

---

## 核心模块API

### src/detect.py - 目标检测

```python
class InfraredDetector:
    def __init__(self, model_path: str, conf: float = 0.25)
    def load_model(self, model_path: str) -> None
    def detect(self, image: Union[str, np.ndarray], conf: float) -> List[Dict]
    def detect_and_draw(self, image, output_path: str, conf: float) -> np.ndarray

def train_yolo_model(data_yaml, model_size='n', epochs=100, imgsz=640, batch=16, output_dir='models/train') -> str
def validate_model(model_path: str, data_yaml: str) -> Dict
```

### src/extract_temp.py - 温度提取

```python
class TempFeatureExtractor:
    def __init__(self, model_path: Optional[str] = None)
    def detect_and_extract(self, image: np.ndarray, conf: float) -> List[Dict]
    def extract_from_roi(self, roi: np.ndarray, include_hist=True, include_gradient=True) -> Dict
    def extract_from_bbox(self, image: np.ndarray, bbox: List[int]) -> Dict

def load_infrared_image(image_path: str) -> np.ndarray
def batch_extract_features(image_dir, output_path, model_path=None, conf=0.25) -> None
```

### src/diagnose.py - 故障诊断

```python
class FaultDiagnoser:
    def __init__(self, model_path: Optional[str] = None, temp_threshold: Optional[Dict] = None)
    def diagnose(self, image: Union[str, np.ndarray]) -> Dict
    def batch_diagnose(self, image_dir: Union[str, Path], output_path: Union[str, Path]) -> None

def load_threshold_config(config_path: str) -> Dict
```

### src/utils/temp_utils.py - 温度工具

```python
def extract_temp_stats(roi: np.ndarray) -> Dict[str, float]
def extract_temp_histogram(roi: np.ndarray, bins: int = 32) -> np.ndarray
def calculate_temp_gradient(roi: np.ndarray) -> Dict[str, float]
def extract_temp_features(roi: np.ndarray, include_hist=True, include_gradient=True) -> Dict
def normalize_roi(roi: np.ndarray, method: str = 'minmax') -> np.ndarray
```

---

## 命令行工具

### 训练模型
```bash
uv run python scripts/train.py train \
    --data data/dataset.yaml \
    --model-size n \
    --epochs 50 \
    --device 0
```

### 验证模型
```bash
uv run python scripts/train.py validate \
    --model weights/best.pt \
    --data data/dataset.yaml
```

### 导出模型
```bash
uv run python scripts/train.py export \
    --model weights/best.pt \
    --format onnx
```

### 下载数据集
```bash
uv run python scripts/download_dataset.py \
    --urls-file data/dataurls.txt \
    --output-dir data/raw/images
```

### 重试下载失败文件
```bash
uv run python scripts/retry_download.py \
    --urls-file data/dataurls.txt \
    --output-dir data/raw/images \
    --wait 60
```

### 转换数据格式
```bash
uv run python scripts/convert_to_yolo.py \
    --input-dir data/raw/images \
    --output-dir data/infrared \
    --whole-image
```

---

## 开发流程

### 阶段一：项目初始化与环境配置

1. 使用uv初始化项目
2. 配置PyTorch CUDA环境
3. 安装ultralytics YOLO框架
4. 验证GPU环境可用性

### 阶段二：数据集准备

1. 下载SCIDB红外图像数据集
2. 按类别组织图像目录
3. 转换为YOLO格式
4. 创建dataset.yaml配置文件

### 阶段三：模型训练

1. 选择YOLOv8n作为基础模型
2. 配置训练参数
3. 启动GPU训练
4. 监控训练指标 (mAP, Precision, Recall)
5. EarlyStopping防止过拟合

### 阶段四：ROI温度特征提取

1. 实现温度统计特征 (max/min/mean/std/median)
2. 实现温度直方图特征
3. 实现温度梯度特征

### 阶段五：故障诊断Pipeline

1. 整合检测+温度提取
2. 定义故障分类规则
3. 实现批量诊断功能
4. 输出结构化JSON结果

### 阶段六：测试与验证

1. 端到端Pipeline测试
2. 批量图像诊断测试
3. 性能指标验证

---

## 技术栈

| 组件 | 技术/版本 |
|------|----------|
| 深度学习框架 | PyTorch 2.11.0+cu130 |
| 目标检测 | YOLOv8 (ultralytics 8.4.27) |
| 图像处理 | OpenCV 4.13.0 |
| 数值计算 | NumPy 2.2.6 |
| 包管理 | uv |
| GPU | NVIDIA RTX 4060 (8GB) |
| CUDA | 13.0 |
| cuDNN | 9.19.0 |

---

## 项目配置

### pyproject.toml

```toml
[project]
name = "infrared-fault"
version = "0.1.0"
requires-python = ">=3.10"

dependencies = [
    "torch>=2.5.1",
    "torchvision>=0.20.1",
    "ultralytics>=8.4.27",
    "opencv-python>=4.13.0.92",
    "numpy>=2.2.6",
    "pillow>=12.1.1",
    "matplotlib>=3.10.8",
    "pandas>=2.3.3",
    "scikit-learn>=1.7.2",
    "seaborn>=0.13.2",
    "tqdm>=4.67.3",
]
```

### 数据集配置 (dataset.yaml)

```yaml
path: /path/to/data
train: train/images
val: valid/images

nc: 12
names: ['bushing', 'arrester', 'breaker', 'clamp', 'conservator',
        'current-transformer', 'disconnector', 'disconnector2',
        'disconnector3', 'heat-sink', 'insulator', 'transformer']
```

---

## 常见问题

### Q: 训练时GPU内存不足
A: 减小batch大小或使用更小的模型 (--batch 8 --model-size n)

### Q: 下载失败429错误
A: 使用retry_download.py脚本，设置更长的等待时间 (--wait 120)

### Q: 如何使用自己的数据集
A:
1. 将图像按类别放入不同子目录
2. 运行convert_to_yolo.py转换格式
3. 更新dataset.yaml中的类别名称

### Q: 如何调整故障诊断阈值
A:
```python
threshold = {"max_temp": 200, "std_temp": 40}
diagnoser = FaultDiagnoser("weights/best.pt", temp_threshold=threshold)
```

---

## 更新日志

| 日期 | 版本 | 更新内容 |
|------|------|----------|
| 2026-03-26 | v0.1.0 | 初始版本，完成基础功能 |
| 2026-03-26 | v0.1.1 | 修复数据集类别转换错误，5类别独立训练 |
| 2026-03-26 | v0.1.2 | 完成端到端测试，性能达标 |

---

## 许可证

MIT License