红外设备检测系统 (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+

安装依赖

# 使用uv安装依赖
uv sync

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

GPU环境测试

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张

---

模型训练

完整训练流程

训练模型

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

---

使用方法

快速开始

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

1. 目标检测

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

检测结果格式:

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

2. 检测并绘制结果

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

3. 温度特征提取

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. 故障诊断

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

诊断输出格式:

{
  "image": "20230907095711003_jpg.rf.0e14e8b4344dce21114683dc3f3e7e2e.jpg",
  "detections": [
    {
      "bbox": [x1, y1, x2, y2],
      "confidence": 0.995,
      "class_id": 0,
      "class_name": "bushing"
    }
  ],
  "faults": [
    {
      "bbox": [x1, y1, x2, y2],
      "confidence": 0.995,
      "temp_max": 255.0,
      "temp_mean": 97.49,
      "temp_std": 45.2,
      "fault_status": "abnormal",
      "fault_type": "hotspot"
    }
  ],
  "summary": {
    "status": "fault_detected",
    "fault_count": 1,
    "abnormal_count": 1,
    "fault_types": ["hotspot"]
  }
}

故障类型分类规则:

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

5. 批量诊断

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

6. 自定义温度阈值

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

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

---

核心模块API

src/detect.py - 目标检测

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 - 温度提取

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 - 故障诊断

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 - 温度工具

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

---

命令行工具

训练模型

uv run python scripts/train.py train \
    --data data/dataset.yaml \
    --model-size n \
    --epochs 50 \
    --device 0

验证模型

uv run python scripts/train.py validate \
    --model weights/best.pt \
    --data data/dataset.yaml

导出模型

uv run python scripts/train.py export \
    --model weights/best.pt \
    --format onnx

下载数据集

uv run python scripts/download_dataset.py \
    --urls-file data/dataurls.txt \
    --output-dir data/raw/images

重试下载失败文件

uv run python scripts/retry_download.py \
    --urls-file data/dataurls.txt \
    --output-dir data/raw/images \
    --wait 60

转换数据格式

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

[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)

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:

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