HVC-Tunnel-Multimodal/Infrared-ai-research-lab
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 2 Wiki Activity

chore: initial commit for Infrared Equipment Detection System

Browse Source 
- Add YOLO-based infrared thermal image detection system
- Add dataset configuration (dataset.yaml) with 12 equipment classes
- Add training, validation and detection scripts
- Add temperature feature extraction and fault diagnosis modules
- Update README and CLAUDE.md with correct dataset paths
- Optimize .gitignore to track dataset.yaml while ignoring large data files
This commit is contained in:
SkyLi
2026-03-26 21:36:08 +08:00
commit fd1a5f4050
19 changed files with 2920 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+572
View File
@@ -0,0 +1,572 @@
# 红外设备检测系统 (Infrared Equipment Detection System)
基于YOLO的红外热图像目标检测系统，用于识别电力设备。
## 项目概述
本项目使用YOLOv8实现红外热图像中的设备目标检测。
### 主要功能
- **YOLO目标检测**: 识别红外图像中的电力设备
- **ROI温度提取**: 从检测区域提取温度统计特征
- **故障诊断**: 基于温度特征判断设备状态
### 设备类别
请根据你的数据集配置类别（修改dataset.yaml中的names字段）
---
## 项目结构
```
1BaseOnYOLO/
├── src/ # 核心源代码
│ ├── detect.py # YOLO目标检测模块
│ ├── extract_temp.py # ROI温度特征提取
│ ├── diagnose.py # 故障诊断Pipeline
│ └── utils/
│ └── temp_utils.py # 温度工具函数
├── scripts/ # 工具脚本
│ ├── train.py # 模型训练/验证入口
│ ├── test_cuda.py # CUDA环境测试
│ ├── download_dataset.py # 数据集下载
│ ├── convert_to_yolo.py # 数据集格式转换
│ ├── retry_download.py # 下载失败重试
│ └── usage_example.py # 使用示例
├── data/ # 数据目录
│ ├── dataset.yaml # YOLO数据集配置
│ ├── train/ # 训练集
│ │ ├── images/ # 训练图像
│ │ └── labels/ # 训练标签
│ ├── valid/ # 验证集
│ │ ├── images/ # 验证图像
│ │ └── labels/ # 验证标签
│ ├── test/ # 测试集
│ │ ├── images/ # 测试图像
│ │ └── labels/ # 测试标签
│ └── 标签.txt # 类别标签说明
├── weights/ # 模型权重
│ └── infrared_yolov8n.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
```
---
## 数据集
### 数据来源
SCIDB红外电力设备数据集: https://www.scidb.cn/en/detail?dataSetId=e416c488169f484485ad7575dcfc43ce
### 数据规模
| 类别 | 图像数量 |
|------|----------|
| Wave Traps | 153 |
| Power Transformers | 178 |
| Circuit Breakers | 162 |
| Surge Arresters | 181 |
| Disconnectors | 89 |
| **总计** | **763** |
### 数据划分
- 训练集: 610张 (80%)
- 验证集: 153张 (20%)
---
## 模型训练
### 完整训练流程
**1. 下载数据集**
```bash
uv run python scripts/download_dataset.py \
--urls-file data/dataurls.txt \
--output-dir data/raw/images
```
**2. 转换数据格式**
```bash
uv run python scripts/convert_to_yolo.py \
--input-dir data/raw/images \
--output-dir data/infrared \
--whole-image
```
**3. 训练模型**
```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 | **99.5%** |
| mAP@0.5:0.95 | **99.5%** |
| Precision | **99.7%** |
| Recall | **99.4%** |
**各类别性能**:
| 类别 | Precision | Recall | mAP@0.5 |
|------|----------|--------|---------|
| Wave Traps | 0.998 | 1.000 | 0.995 |
| Power Transformers | 0.993 | 1.000 | 0.995 |
| Circuit Breakers | 0.998 | 1.000 | 0.995 |
| Surge Arresters | 1.000 | 0.971 | 0.995 |
| Disconnectors | 0.998 | 1.000 | 0.995 |
---
## 使用方法
### 快速开始
```bash
# 运行使用示例
uv run python scripts/usage_example.py
```
### 1. 目标检测
```python
from src.detect import InfraredDetector
# 初始化检测器
detector = InfraredDetector("weights/infrared_yolov8n.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': 'Wave Traps'
}
```
### 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/infrared_yolov8n.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/infrared_yolov8n.pt")
result = diagnoser.diagnose("path/to/image.jpg")
print(f"诊断状态: {result['summary']['status']}")
print(f"故障类型: {result['summary']['fault_types']}")
```
**诊断输出格式**:
```json
{
"image": "FLIR0150.jpg",
"detections": [
{
"bbox": [x1, y1, x2, y2],
"confidence": 0.995,
"class_id": 0,
"class_name": "Wave Traps"
}
],
"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. 批量诊断
```python
diagnoser.batch_diagnose(
"data/valid/images",
"/tmp/batch_results.json"
)
```
### 6. 自定义温度阈值
```python
threshold = {
"max_temp": 200, # 最高温度阈值
"std_temp": 40 # 标准差阈值
}
diagnoser = FaultDiagnoser(
"weights/infrared_yolov8n.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/infrared_yolov8n.pt \
--data data/dataset.yaml
```
### 导出模型
```bash
uv run python scripts/train.py export \
--model weights/infrared_yolov8n.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/infrared_yolov8n.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