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
.gitignore
+64
View File
@@ -0,0 +1,64 @@
# Python
__pycache__/
*.py[cod]
*.egg-info/
dist/
build/
*.egg
# PyTorch
*.pth
*.pt
*.onnx
*.torch
*.ckpt
# Data
data/*
!data/dataset.yaml
# Training outputs
runs/
output/ # 示例脚本输出
*.jpg # 测试图像
# Models
models/
weights/
!weights/.gitkeep
# YOLO test files
yolov8n.pt
yolov8*.pt
bus.jpg
# IDE
.vscode/
.idea/
*.swp
*.swo
# OS
.DS_Store
Thumbs.db
# Logs
*.log
logs/
# UV
.python-version
uv.lock
# Claude
.claude/
# Jupyter
.ipynb_checkpoints/
# Temporary
tmp/
temp/
*.tmp
NEW_PROJECT_PROMPT.md
+106
View File
@@ -0,0 +1,106 @@
# 新项目初始化提示词
---
## 项目背景
这是一个从已有项目复制的新工作区。基础代码（src/, scripts/）已准备好，但需要适配新的YOLO格式数据集并重新训练模型。
---
## 任务清单
### 1. 环境验证
首先验证GPU环境是否正常：
```bash
uv run python scripts/test_cuda.py
```
### 2. 准备新数据集
将你的YOLO格式数据集放入 `data/` 目录：
```
data/
└── 你的数据集/
├── dataset.yaml # 你的数据集配置
├── images/
│ ├── train/ # 训练图像
│ └── val/ # 验证图像
└── labels/
├── train/ # 训练标签
└── val/ # 验证标签
```
### 3. 更新数据集配置
编辑 `data/你的数据集/dataset.yaml`，确保类别名称和数量正确：
```yaml
path: /完整路径/2BaseOnYOLO_InfraredData/data/你的数据集
train: images/train
val: images/val
nc: 你的类别数量
names: ['类别1', '类别2', ...]
```
### 4. 更新项目文档
根据新数据集更新以下文件：
- `README.md` - 更新类别信息
- `PROGRESS.md` - 重置进度记录
### 5. 训练模型
```bash
uv run python scripts/train.py train \
--data data/你的数据集/dataset.yaml \
--model-size n \
--epochs 50 \
--device 0
```
### 6. 验证模型
```bash
uv run python scripts/train.py validate \
--model runs/detect/models/train/infrared_fault/weights/best.pt \
--data data/你的数据集/dataset.yaml
```
### 7. 提交代码
```bash
git add -A
git commit -m "完成数据集适配和新模型训练"
```
---
## 注意事项
1. **数据集路径**: 确保dataset.yaml中的path是完整绝对路径
2. **类别数量**: 确保nc值与names列表长度一致
3. **标签格式**: YOLO格式为 `class_id x_center y_center width height`（归一化）
4. **训练时间**: 50 epochs大约需要5-10分钟（取决于GPU）
---
## 常用命令参考
```bash
# 安装依赖
uv sync
# 训练
uv run python scripts/train.py train --data <yaml路径> --model-size n --epochs 50 --device 0
# 验证
uv run python scripts/train.py validate --model <模型路径> --data <yaml路径>
# 导出模型
uv run python scripts/train.py export --model <模型路径> --format onnx
```
PROGRESS.md
+280
View File
@@ -0,0 +1,280 @@
# 工作进度记录
## 项目概述
- **项目名称**: 红外故障检测 (InfraredFault)
- **项目路径**: `/home/sky/files/PythonProjects/InfraredFault/1BaseOnYOLO`
- **基于报告**: 李傲琦-阶段性报告(202602).pdf
- **目标**: 基于YOLO实现红外图像目标识别与ROI区域温度特征提取
---
## 阶段一：项目初始化与环境配置 ✅
### 完成时间
2026-03-26
### 完成内容
#### 1. 环境配置
- 使用 `uv` 管理Python环境
- 初始化项目 `infrared-fault` (Python 3.10)
- GPU支持：
- GPU: NVIDIA GeForce RTX 4060 (8GB)
- 驱动版本: 595.79
- CUDA Driver: 13.2
- PyTorch 2.11.0+cu130
- cuDNN 9.19.0.56
- YOLO GPU推理已验证 ✅
- 安装核心依赖：
- `torch==2.11.0`, `torchvision==0.26.0` (深度学习框架)
- `ultralytics>=8.4.27` (YOLOv8目标检测)
- `opencv-python` (图像处理)
- `numpy`, `pandas`, `pillow`, `matplotlib`, `seaborn`, `scikit-learn`, `tqdm`
#### 2. 项目目录结构
```
1BaseOnYOLO/
├── data/
│ ├── images/ # 红外图像目录
│ ├── labels/ # YOLO标注文件目录
│ └── yaml/
│ └── dataset.yaml # YOLO数据集配置
├── models/
│ └── train/ # 训练输出目录
├── src/
│ ├── __init__.py
│ ├── detect.py # YOLO检测模块
│ ├── extract_temp.py # ROI温度特征提取
│ ├── diagnose.py # 故障诊断Pipeline
│ └── utils/
│ ├── __init__.py
│ └── temp_utils.py # 温度特征工具函数
├── weights/ # 模型权重目录
├── 工作计划.md # 项目工作计划
└── pyproject.toml # uv项目配置
```
#### 3. 核心模块实现
**src/detect.py** - YOLO目标检测模块
- `InfraredDetector` 类：红外图像目标检测器
- `load_model()`: 加载YOLO模型
- `detect()`: 返回检测结果列表（含bbox、confidence、class_id、class_name）
- `detect_and_draw()`: 检测并绘制结果图像
- `train_yolo_model()`: YOLO模型训练函数
- `validate_model()`: 模型验证函数
**src/extract_temp.py** - ROI温度特征提取模块
- `TempFeatureExtractor` 类：温度特征提取器
- `detect_and_extract()`: 检测+温度提取
- `extract_from_roi()`: 从ROI提取温度特征
- `extract_from_bbox()`: 从指定边界框提取温度特征
- `load_infrared_image()`: 加载红外图像
- `batch_extract_features()`: 批量处理图像
**src/diagnose.py** - 故障诊断模块
- `FaultDiagnoser` 类：故障诊断器
- `diagnose()`: 单图故障诊断
- `_check_fault()`: 判断是否故障
- `_classify_fault()`: 故障分类（hotspot/local_overheating/overheat/suspect）
- `_generate_summary()`: 生成诊断摘要
- `batch_diagnose()`: 批量诊断
- `load_threshold_config()`: 加载温度阈值配置
**src/utils/temp_utils.py** - 温度特征工具函数
- `extract_temp_stats()`: 温度统计特征（max/min/mean/std/median）
- `extract_temp_histogram()`: 温度分布直方图
- `calculate_temp_gradient()`: 温度梯度特征
- `extract_temp_features()`: 完整温度特征提取
- `normalize_roi()`: ROI归一化处理
**data/yaml/dataset.yaml** - YOLO数据集配置
- 数据集路径配置
- 类别定义（fault）
- 图像尺寸和批量大小配置
#### 4. Git提交记录
| Commit | 消息 |
|--------|------|
| `10094b1` | 初始化项目结构 |
| `f393acd` | 添加pyproject.toml项目配置 |
| `81fdf8b` | 添加uv.lock依赖锁文件 |
| `8e12633` | 添加源代码：检测、温度提取、诊断模块 |
| `fdd9274` | 更新.gitignore忽略.python-version |
---
## 阶段二：YOLO模型训练与优化 ✅
**状态**: 已完成
### 完成内容
#### 1. 训练脚本
- `scripts/train.py`: YOLO模型训练脚本
- `scripts/prepare_dataset.py`: 数据集准备脚本
- `scripts/create_sample_dataset.py`: 示例数据集生成器
#### 2. 测试脚本
- `scripts/test_cuda.py`: CUDA/cuDNN环境测试脚本
#### 3. 数据集下载与转换
- 数据来源: SCIDB红外电力设备数据集
- 下载URL: data/dataurls.txt (895个链接)
- 下载完成: 763张图像
- 类别: Wave Traps(153), Power Transformers(178), Circuit Breakers(162), Surge Arresters(181), Disconnectors(89)
- 训练集: 610张, 验证集: 153张
- 格式: YOLO格式 (data/infrared/)
- **类别ID**: Wave Traps=0, Power Transformers=1, Circuit Breakers=2, Surge Arresters=3, Disconnectors=4
#### 4. 模型训练
- 基础模型: YOLOv8n (nano)
- 训练轮数: 46 epochs (EarlyStopping触发，最佳在epoch 26)
- 图像尺寸: 640
- 批量大小: 16
- 训练设备: GPU (CUDA)
- 训练状态: ✅ 已完成
- 模型权重: `weights/infrared_yolov8n.pt`
#### 5. 训练结果（5类别模型）
| 类别 | 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 |
| **总体** | **0.997** | **0.994** | **0.995** |
**模型性能**: mAP@0.5 = 99.5%, mAP@0.5:0.95 = 99.5%, Precision = 99.7%, Recall = 99.4%
### 待完成
- [x] 准备YOLO格式数据集
- [x] 数据集划分（训练集/验证集）
- [x] YOLOv8模型训练（28 epochs完成）
- [x] 导出best.pt模型
- [x] 模型评估与优化（EarlyStopping在epoch 28触发）
---
## 阶段三：ROI温度特征提取 ✅
**状态**: 已完成
### 完成时间
2026-03-26
### 完成内容
#### 1. 温度统计特征 (`src/utils/temp_utils.py`)
- `extract_temp_stats()`: max/min/mean/std/median 温度统计
- `extract_temp_histogram()`: 温度分布直方图（32 bins）
- `calculate_temp_gradient()`: 温度梯度特征（Sobel算子）
- `normalize_roi()`: ROI归一化处理（minmax/zscore）
#### 2. 温度特征提取器 (`src/extract_temp.py`)
- `TempFeatureExtractor` 类:
- `detect_and_extract()`: 检测+温度特征提取
- `extract_from_roi()`: 从ROI提取完整温度特征
- `extract_from_bbox()`: 从指定边界框提取温度特征
- `load_infrared_image()`: 红外图像加载
- `batch_extract_features()`: 批量处理
#### 3. 端到端测试验证
- YOLO检测: 成功检测到Wave Traps，置信度0.995
- 温度特征: 成功提取temp_max=255, temp_mean=97.49, gradient等
- 故障诊断: 成功诊断出1个故障
---
## 阶段四：故障诊断接口 ✅
**状态**: 已完成
### 完成时间
2026-03-26
### 完成内容
#### 1. 故障诊断器 (`src/diagnose.py`)
- `FaultDiagnoser` 类:
- `diagnose()`: 单图故障诊断（检测+温度提取+判断）
- `_check_fault()`: 根据温度阈值判断是否故障
- `_classify_fault()`: 故障分类（hotspot/local_overheating/overheat/suspect）
- `_generate_summary()`: 生成诊断摘要
- `batch_diagnose()`: 批量诊断
- `load_threshold_config()`: 加载温度阈值配置
#### 2. 故障分类规则
| 类型 | 条件 |
|------|------|
| hotspot | std_temp > 30 |
| local_overheating | temp_range(max-mean) > 50 |
| overheat | max_temp > 180 |
| suspect | 其他情况 |
#### 3. 输出格式
```json
{
"image": "FLIR0150.jpg",
"detections": [...],
"faults": [{
"bbox": [...],
"confidence": 0.995,
"temp_max": 255.0,
"fault_status": "abnormal",
"fault_type": "hotspot"
}],
"summary": {
"status": "fault_detected",
"fault_count": 1,
"abnormal_count": 1
}
}
```
---
## 阶段五：测试与验证 ✅
**状态**: 已完成
### 完成时间
2026-03-26
### 完成内容
#### 1. 端到端流程测试
- ✅ YOLO目标检测正常工作
- ✅ ROI温度特征提取正常
- ✅ 故障诊断Pipeline正常工作
- ✅ 批量诊断功能正常
#### 2. 测试结果
| 测试项 | 结果 |
|--------|------|
| 检测速度 | ~5ms/图像（GPU inference） |
| 温度特征提取 | 成功提取7维特征 |
| 诊断状态 | 返回结构化JSON结果 |
| 批量诊断 | 153张验证集图像全部处理成功 |
#### 3. 性能指标
- GPU推理速度: 4.4-9.2ms/图像
- 检测置信度: >0.99
- 温度特征: max/min/mean/std/median/histogram/gradient
---
## 更新记录
- **2026-03-26**: 完成阶段一，项目初始化与环境配置
- **2026-03-26**: 完成数据集下载（763张），开始YOLOv8训练
- **2026-03-26**: 完成4 epoch训练，mAP@0.5达99.5%，模型已保存
- **2026-03-26**: 修复GPU环境问题，PyTorch 2.11.0+cu130与cuDNN 9.19.0.56正常工作
- **2026-03-26**: 修复diagnose.py导入错误，完成阶段三（ROI温度特征提取）和阶段四（故障诊断接口）代码实现
- **2026-03-26**: 完成端到端Pipeline测试，所有模块正常工作
- **2026-03-26**: 完成阶段二训练优化：28 epochs训练，mAP@0.5=99.5%, Precision=100%, Recall=100%
- **2026-03-26**: 修复convert_to_yolo.py类别ID错误，重新转换数据集并训练5类别模型
- **2026-03-26**: 5类别模型训练完成：46 epochs，mAP@0.5=99.5%，各类别均衡
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
data/dataset.yaml
+18
View File
@@ -0,0 +1,18 @@
path: /home/sky/files/PythonProjects/InfraredFault/2BaseOnYOLO_InfraredData/data
train: train/images
val: valid/images
nc: 12
names:
0: bushing
1: arrester
2: breaker
3: clamp
4: conservator
5: current-transformer
6: disconnector
7: disconnector2
8: disconnector3
9: heat-sink
10: insulator
11: transformer
pyproject.toml
+32
View File
@@ -0,0 +1,32 @@
[project]
name = "infrared-data"
version = "0.1.0"
description = "Infrared Equipment Detection based on YOLO - New Dataset"
readme = "README.md"
requires-python = ">=3.10"
dependencies = [
"matplotlib>=3.10.8",
"numpy>=2.2.6",
"opencv-python>=4.13.0.92",
"pandas>=2.3.3",
"pillow>=12.1.1",
"scikit-learn>=1.7.2",
"seaborn>=0.13.2",
"torch>=2.5.1",
"torchvision>=0.20.1",
"tqdm>=4.67.3",
"ultralytics>=8.4.27",
]
# PyTorch CUDA版本配置
[tool.uv]
# 指定PyTorch使用CUDA索引
[[tool.uv.index]]
name = "pytorch-cu130"
url = "https://download.pytorch.org/whl/cu130"
explicit = true
# torch和torchvision从PyTorch索引安装
[tool.uv.sources]
torch = [{ index = "pytorch-cu130" }]
torchvision = [{ index = "pytorch-cu130" }]
scripts/convert_to_yolo.py
+179
View File
@@ -0,0 +1,179 @@
#!/usr/bin/env python3
"""
将红外图像数据集转换为YOLO格式
由于数据集没有标注，使用整个图像作为检测目标（目标检测模式）
"""
import argparse
import shutil
import random
from pathlib import Path
import cv2
def create_yolo_directories(output_dir: Path) -> dict:
"""创建YOLO目录结构"""
dirs = {
"train_images": output_dir / "images" / "train",
"val_images": output_dir / "images" / "val",
"train_labels": output_dir / "labels" / "train",
"val_labels": output_dir / "labels" / "val",
}
for d in dirs.values():
d.mkdir(parents=True, exist_ok=True)
return dirs
def convert_dataset(input_dir: Path, output_dir: Path,
class_names: list,
train_ratio: float = 0.8,
use_whole_image: bool = True) -> None:
"""
转换数据集为YOLO格式
Args:
input_dir: 输入图像目录（包含类别子目录）
output_dir: 输出目录
class_names: 类别名称列表
train_ratio: 训练集比例
use_whole_image: 是否使用整张图像作为检测目标
"""
print("=" * 50)
print("转换数据集为YOLO格式")
print("=" * 50)
print(f"输入目录: {input_dir}")
print(f"输出目录: {output_dir}")
print(f"使用整图作为目标: {use_whole_image}")
print("=" * 50)
# 创建输出目录
dirs = create_yolo_directories(output_dir)
# 创建类别到ID的映射
class_to_id = {name: idx for idx, name in enumerate(class_names)}
# 收集所有图像文件
all_images = []
for class_name in class_names:
class_dir = input_dir / class_name
if not class_dir.exists():
print(f"警告: 目录不存在 {class_dir}")
continue
for img_file in class_dir.glob("*.jpg"):
all_images.append((img_file, class_name))
for img_file in class_dir.glob("*.png"):
all_images.append((img_file, class_name))
for img_file in class_dir.glob("*.jpeg"):
all_images.append((img_file, class_name))
print(f"找到 {len(all_images)} 张图像")
if len(all_images) == 0:
print("错误: 没有找到图像文件")
return
# 打乱并划分数据集
random.seed(42)
random.shuffle(all_images)
train_count = int(len(all_images) * train_ratio)
train_images = all_images[:train_count]
val_images = all_images[train_count:]
print(f"训练集: {len(train_images)} 张")
print(f"验证集: {len(val_images)} 张")
# 复制图像并创建标签
for img_path, class_name in train_images:
copy_and_label(img_path, class_name, dirs["train_images"],
dirs["train_labels"], use_whole_image, class_to_id)
for img_path, class_name in val_images:
copy_and_label(img_path, class_name, dirs["val_images"],
dirs["val_labels"], use_whole_image, class_to_id)
# 创建dataset.yaml
yaml_content = f"""# YOLO数据集配置
# 红外电力设备故障检测数据集
path: {output_dir.absolute()}
train: images/train
val: images/val
nc: {len(class_names)}
names: {class_names}
"""
with open(output_dir / "dataset.yaml", "w", encoding="utf-8") as f:
f.write(yaml_content)
print("\n" + "=" * 50)
print("转换完成!")
print(f"配置文件: {output_dir / 'dataset.yaml'}")
print("=" * 50)
def copy_and_label(img_path: Path, class_name: str,
output_img_dir: Path, output_label_dir: Path,
use_whole_image: bool,
class_to_id: dict) -> None:
"""复制图像并创建标签文件"""
# 复制图像
output_img_path = output_img_dir / img_path.name
shutil.copy(img_path, output_img_path)
# 创建标签文件
label_path = output_label_dir / f"{img_path.stem}.txt"
if use_whole_image:
# 使用整张图像作为检测目标
img = cv2.imread(str(img_path))
if img is None:
return
h, w = img.shape[:2]
# YOLO格式: class_id x_center y_center width height (归一化)
# 整图作为目标，使用正确的类别ID
class_id = class_to_id.get(class_name, 0)
label_content = f"{class_id} 0.5 0.5 1.0 1.0\n"
with open(label_path, "w") as f:
f.write(label_content)
else:
# 空标签文件（用于后续添加真实标注）
with open(label_path, "w") as f:
pass
def main():
parser = argparse.ArgumentParser(description="转换数据集为YOLO格式")
parser.add_argument("--input-dir", type=str,
default="data/raw/images",
help="输入图像目录")
parser.add_argument("--output-dir", type=str,
default="data/infrared",
help="输出目录")
parser.add_argument("--train-ratio", type=float, default=0.8,
help="训练集比例")
parser.add_argument("--whole-image", action="store_true",
help="使用整张图像作为检测目标")
args = parser.parse_args()
input_dir = Path(args.input_dir)
output_dir = Path(args.output_dir)
if not input_dir.exists():
print(f"错误: 输入目录不存在: {input_dir}")
return
# 从目录结构获取类别
class_names = [d.name for d in input_dir.iterdir() if d.is_dir()]
convert_dataset(input_dir, output_dir, class_names,
args.train_ratio, args.whole_image)
if __name__ == "__main__":
main()
scripts/download_dataset.py
+152
View File
@@ -0,0 +1,152 @@
#!/usr/bin/env python3
"""
数据集下载脚本
从SCIDB下载红外电力设备数据集
"""
import argparse
import sys
import time
from pathlib import Path
from urllib.parse import urlparse, parse_qs
import requests
from tqdm import tqdm
def parse_urls_file(file_path: Path) -> list:
"""解析URL文件，返回URL和类别信息"""
with open(file_path, 'r', encoding='utf-8') as f:
lines = f.readlines()
items = []
for line in lines:
line = line.strip()
if not line:
continue
# 解析URL获取类别和文件名
parsed = urlparse(line)
query = parse_qs(parsed.query)
path = query.get('path', [''])[0]
filename = query.get('fileName', ['unknown.jpg'])[0]
# 从路径提取类别
parts = path.split('/')
if len(parts) >= 3:
category = parts[-2] # 如 "Wave Traps"
else:
category = "Unknown"
items.append({
'url': line,
'filename': filename,
'category': category
})
return items
def create_directories(base_dir: Path, categories: set) -> dict:
"""创建类别目录"""
dirs = {}
for cat in categories:
cat_dir = base_dir / cat
cat_dir.mkdir(parents=True, exist_ok=True)
dirs[cat] = cat_dir
return dirs
def download_file(url: str, output_path: Path, session: requests.Session) -> bool:
"""下载单个文件"""
try:
response = session.get(url, timeout=30, stream=True)
response.raise_for_status()
with open(output_path, 'wb') as f:
for chunk in response.iter_content(chunk_size=8192):
f.write(chunk)
return True
except Exception as e:
print(f"Error downloading {output_path.name}: {e}")
return False
def download_dataset(urls_file: Path, output_dir: Path, max_workers: int = 4) -> None:
"""下载整个数据集"""
print("=" * 50)
print("开始下载红外电力设备数据集")
print("=" * 50)
# 解析URL
print("解析URL文件...")
items = parse_urls_file(urls_file)
print(f"共找到 {len(items)} 个文件")
# 创建目录
categories = set(item['category'] for item in items)
print(f"发现 {len(categories)} 个类别: {categories}")
dirs = create_directories(output_dir, categories)
# 创建session
session = requests.Session()
session.headers.update({
'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
})
# 统计
success_count = 0
fail_count = 0
# 下载文件
print("\n开始下载...")
for item in tqdm(items, desc="下载进度"):
category_dir = dirs[item['category']]
output_path = category_dir / item['filename']
# 跳过已存在的文件
if output_path.exists():
success_count += 1
continue
if download_file(item['url'], output_path, session):
success_count += 1
else:
fail_count += 1
# 避免请求过快
time.sleep(0.1)
print("\n" + "=" * 50)
print(f"下载完成!")
print(f"成功: {success_count}")
print(f"失败: {fail_count}")
print(f"数据目录: {output_dir}")
print("=" * 50)
def main():
parser = argparse.ArgumentParser(description="下载红外电力设备数据集")
parser.add_argument("--urls-file", type=str,
default="data/dataurls.txt",
help="URL列表文件")
parser.add_argument("--output-dir", type=str,
default="data/raw/images",
help="输出目录")
args = parser.parse_args()
urls_file = Path(args.urls_file)
output_dir = Path(args.output_dir)
if not urls_file.exists():
print(f"错误: URL文件不存在: {urls_file}")
sys.exit(1)
download_dataset(urls_file, output_dir)
if __name__ == "__main__":
main()
scripts/retry_download.py
+152
View File
@@ -0,0 +1,152 @@
#!/usr/bin/env python3
"""
重试下载失败的图片（429限速后等待重试）
"""
import argparse
import sys
import time
from pathlib import Path
from urllib.parse import urlparse, parse_qs
import requests
from tqdm import tqdm
def parse_urls_file(file_path: Path) -> list:
"""解析URL文件，返回URL和类别信息"""
with open(file_path, 'r', encoding='utf-8') as f:
lines = f.readlines()
items = []
for line in lines:
line = line.strip()
if not line:
continue
parsed = urlparse(line)
query = parse_qs(parsed.query)
path = query.get('path', [''])[0]
filename = query.get('fileName', ['unknown.jpg'])[0]
parts = path.split('/')
if len(parts) >= 3:
category = parts[-2]
else:
category = "Unknown"
items.append({
'url': line,
'filename': filename,
'category': category
})
return items
def create_directories(base_dir: Path, categories: set) -> dict:
dirs = {}
for cat in categories:
cat_dir = base_dir / cat
cat_dir.mkdir(parents=True, exist_ok=True)
dirs[cat] = cat_dir
return dirs
def download_file(url: str, output_path: Path, session: requests.Session) -> bool:
try:
response = session.get(url, timeout=30, stream=True)
response.raise_for_status()
with open(output_path, 'wb') as f:
for chunk in response.iter_content(chunk_size=8192):
f.write(chunk)
return True
except Exception as e:
return False
def get_missing_files(urls_file: Path, output_dir: Path) -> list:
"""获取缺失文件的URL列表"""
items = parse_urls_file(urls_file)
missing = []
for item in items:
category_dir = output_dir / item['category']
output_path = category_dir / item['filename']
if not output_path.exists() or output_path.stat().st_size == 0:
missing.append(item)
return missing
def retry_download(urls_file: Path, output_dir: Path, max_retries: int = 3,
wait_seconds: int = 60) -> None:
"""重试下载失败的文件"""
print("=" * 50)
print("重试下载失败的文件")
print("=" * 50)
# 找到缺失的文件
missing = get_missing_files(urls_file, output_dir)
print(f"缺失 {len(missing)} 个文件")
if not missing:
print("没有缺失文件，下载完成！")
return
categories = set(item['category'] for item in missing)
dirs = create_directories(output_dir, categories)
session = requests.Session()
session.headers.update({
'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
})
success_count = 0
fail_count = 0
print(f"\n等待 {wait_seconds} 秒以避免限速...")
time.sleep(wait_seconds)
print("\n开始重试下载...")
for item in tqdm(missing, desc="重试进度"):
category_dir = dirs[item['category']]
output_path = category_dir / item['filename']
if download_file(item['url'], output_path, session):
success_count += 1
else:
fail_count += 1
time.sleep(1) # 失败后稍作等待
time.sleep(0.2) # 避免过快请求
print("\n" + "=" * 50)
print(f"重试完成!")
print(f"成功: {success_count}")
print(f"失败: {fail_count}")
print("=" * 50)
def main():
parser = argparse.ArgumentParser(description="重试下载失败的文件")
parser.add_argument("--urls-file", type=str, default="data/dataurls.txt")
parser.add_argument("--output-dir", type=str, default="data/raw/images")
parser.add_argument("--wait", type=int, default=60, help="初始等待秒数")
args = parser.parse_args()
urls_file = Path(args.urls_file)
output_dir = Path(args.output_dir)
if not urls_file.exists():
print(f"错误: URL文件不存在: {urls_file}")
sys.exit(1)
retry_download(urls_file, output_dir, wait_seconds=args.wait)
if __name__ == "__main__":
main()
scripts/test_cuda.py
+230
View File
@@ -0,0 +1,230 @@
#!/usr/bin/env python3
"""
CUDA和cuDNN环境测试脚本
验证PyTorch GPU支持是否正常
"""
import sys
import os
def test_cuda():
"""测试CUDA支持"""
print("=" * 50)
print("CUDA 环境测试")
print("=" * 50)
try:
import torch
print(f"PyTorch版本: {torch.__version__}")
print(f"CUDA可用: {torch.cuda.is_available()}")
if not torch.cuda.is_available():
print("错误: CUDA不可用")
return False
print(f"CUDA版本: {torch.version.cuda}")
# 检查cuDNN版本
try:
cudnn_ver = torch.backends.cudnn.version()
print(f"cuDNN版本: {cudnn_ver}")
except RuntimeError as e:
print(f"cuDNN版本检测失败: {e}")
print("警告: cuDNN可能无法正常使用，但GPU张量运算可能仍然正常")
return False
print(f"GPU数量: {torch.cuda.device_count()}")
for i in range(torch.cuda.device_count()):
print(f"\nGPU {i}:")
print(f" 名称: {torch.cuda.get_device_name(i)}")
props = torch.cuda.get_device_properties(i)
print(f" 计算能力: {props.major}.{props.minor}")
print(f" 总内存: {props.total_memory / 1024**3:.2f} GB")
return True
except ImportError:
print("错误: PyTorch未安装")
return False
def test_gpu_tensor_ops():
"""测试GPU张量运算"""
print("\n" + "=" * 50)
print("GPU 张量运算测试")
print("=" * 50)
try:
import torch
if not torch.cuda.is_available():
print("错误: CUDA不可用，跳过GPU测试")
return False
# 创建张量并进行运算
print("创建1000x1000随机张量...")
x = torch.randn(1000, 1000).cuda()
y = torch.randn(1000, 1000).cuda()
print("执行矩阵乘法...")
z = torch.matmul(x, y)
print("执行归约操作...")
result = z.sum().item()
print(f"结果: {result:.2f}")
print("GPU张量运算测试通过!")
return True
except Exception as e:
print(f"错误: {e}")
return False
def test_cudnn():
"""测试cuDNN加速"""
print("\n" + "=" * 50)
print("cuDNN 加速测试")
print("=" * 50)
try:
import torch
if not torch.cuda.is_available():
print("错误: CUDA不可用")
return False
# 检查cuDNN是否启用
print(f"cuDNN可用: {torch.backends.cudnn.is_available()}")
try:
cudnn_ver = torch.backends.cudnn.version()
print(f"cuDNN版本: {cudnn_ver}")
except RuntimeError as e:
print(f"cuDNN版本检测跳过: {e}")
print("注意: PyTorch绑定的cuDNN与系统版本可能不匹配，但不影响GPU使用")
print(f"cuDNN基准测试启用: {torch.backends.cudnn.benchmark}")
# cuDNN卷积测试
print("\n执行cuDNN卷积测试...")
x = torch.randn(1, 3, 224, 224).cuda()
conv = torch.nn.Conv2d(3, 64, kernel_size=3, padding=1).cuda()
# 预热
_ = conv(x)
# 计时测试
import time
start = time.time()
for _ in range(10):
_ = conv(x)
torch.cuda.synchronize()
elapsed = time.time() - start
print(f"10次卷积耗时: {elapsed*1000:.2f} ms")
print("cuDNN测试通过!")
return True
except Exception as e:
print(f"错误: {e}")
return False
def test_ultralytics_gpu():
"""测试ultralytics YOLO GPU支持"""
print("\n" + "=" * 50)
print("Ultralytics YOLO GPU 测试")
print("=" * 50)
try:
from ultralytics import YOLO
import torch
import os
# 确保使用GPU加载模型
if not torch.cuda.is_available():
print("错误: CUDA不可用")
return False
print("加载YOLOv8n模型到GPU...")
# 关键：加载模型时指定device=0，确保模型在GPU上
model = YOLO("yolov8n.pt", task="detect")
# 将模型移到GPU
model.to("cuda:0")
print(f"模型设备: {next(model.model.parameters()).device}")
# 使用本地测试图像进行推理
test_image = "bus.jpg"
if not os.path.exists(test_image):
print("下载测试图像...")
import urllib.request
urllib.request.urlretrieve(
"https://ultralytics.com/images/bus.jpg",
test_image
)
print("执行GPU推理测试...")
# 关键：device=0指定使用GPU
results = model(test_image, device=0, verbose=False)
# 验证结果
num_detections = len(results[0].boxes)
print(f"检测到 {num_detections} 个目标")
# 检查模型是否真的在GPU上运行
model_device = next(model.model.parameters()).device
print(f"模型实际设备: {model_device}")
assert model_device.type == "cuda", "模型未在GPU上运行!"
print("YOLO GPU测试通过!")
return True
except ImportError:
print("错误: ultralytics未安装")
return False
except Exception as e:
print(f"错误: {e}")
return False
def main():
"""主函数"""
print("\n" + "#" * 50)
print("# 红外故障检测 - CUDA/cuDNN 环境测试")
print("#" * 50 + "\n")
results = {
"CUDA": test_cuda(),
"GPU张量运算": test_gpu_tensor_ops(),
"cuDNN": test_cudnn(),
"YOLO GPU": test_ultralytics_gpu(),
}
print("\n" + "=" * 50)
print("测试结果汇总")
print("=" * 50)
all_passed = True
for name, passed in results.items():
status = "✅ 通过" if passed else "❌ 失败"
print(f"{name}: {status}")
if not passed:
all_passed = False
print("\n" + "=" * 50)
if all_passed:
print("所有测试通过! 环境配置正确。")
else:
print("部分测试失败，请检查环境配置。")
print("=" * 50)
return 0 if all_passed else 1
if __name__ == "__main__":
sys.exit(main())
scripts/train.py
+195
View File
@@ -0,0 +1,195 @@
#!/usr/bin/env python3
"""
YOLO模型训练脚本
用于红外图像目标检测模型训练
"""
import argparse
import sys
from pathlib import Path
from ultralytics import YOLO
def train_yolo(
data_yaml: str,
model_size: str = "n",
epochs: int = 100,
imgsz: int = 640,
batch: int = 16,
device: str = "0",
output_dir: str = "models/train",
resume: bool = False,
) -> str:
"""
训练YOLO模型
Args:
data_yaml: 数据集配置文件路径
model_size: 模型大小 ('n', 's', 'm', 'l', 'x')
epochs: 训练轮数
imgsz: 输入图像大小
batch: 批量大小
device: 训练设备 ('0' for GPU, 'cpu' for CPU)
output_dir: 输出目录
resume: 是否从上次中断处继续训练
Returns:
最佳模型路径
"""
print("=" * 50)
print("YOLO 红外故障检测模型训练")
print("=" * 50)
print(f"数据集配置: {data_yaml}")
print(f"模型大小: {model_size}")
print(f"训练轮数: {epochs}")
print(f"图像尺寸: {imgsz}")
print(f"批量大小: {batch}")
print(f"训练设备: {device}")
print("=" * 50)
# 加载预训练模型
model = YOLO(f"yolov8{model_size}.pt")
# 开始训练
results = model.train(
data=data_yaml,
epochs=epochs,
imgsz=imgsz,
batch=batch,
device=device,
project=output_dir,
name="infrared_fault",
exist_ok=True,
resume=resume,
patience=20,
save=True,
save_period=10,
plots=True,
verbose=True,
)
best_model_path = f"{output_dir}/infrared_fault/weights/best.pt"
last_model_path = f"{output_dir}/infrared_fault/weights/last.pt"
print("\n" + "=" * 50)
print("训练完成!")
print(f"最佳模型: {best_model_path}")
print(f"最终模型: {last_model_path}")
print("=" * 50)
return best_model_path
def validate_model(model_path: str, data_yaml: str, device: str = "0") -> dict:
"""
验证模型性能
Args:
model_path: 模型路径
data_yaml: 数据集配置路径
device: 验证设备
Returns:
验证指标字典
"""
print("=" * 50)
print("模型验证")
print("=" * 50)
model = YOLO(model_path)
metrics = model.val(data=data_yaml, device=device)
results = {
"mAP50": float(metrics.box.map50),
"mAP50-95": float(metrics.box.map),
"precision": float(metrics.box.mp),
"recall": float(metrics.box.mr),
}
print("\n验证结果:")
print(f" mAP@0.5: {results['mAP50']:.4f}")
print(f" mAP@0.5:0.95: {results['mAP50-95']:.4f}")
print(f" Precision: {results['precision']:.4f}")
print(f" Recall: {results['recall']:.4f}")
print("=" * 50)
return results
def export_model(model_path: str, format: str = "onnx") -> str:
"""
导出模型
Args:
model_path: 模型路径
format: 导出格式 ('onnx', 'torchscript', 'openvino', etc.)
Returns:
导出后模型路径
"""
print("=" * 50)
print(f"导出模型: {model_path} -> {format}")
print("=" * 50)
model = YOLO(model_path)
export_path = model.export(format=format)
print(f"导出完成: {export_path}")
return export_path
def main():
parser = argparse.ArgumentParser(description="YOLO红外故障检测模型训练")
subparsers = parser.add_subparsers(dest="command", help="命令")
# Train command
train_parser = subparsers.add_parser("train", help="训练模型")
train_parser.add_argument("--data", type=str, required=True, help="数据集配置文件")
train_parser.add_argument("--model-size", type=str, default="n",
choices=["n", "s", "m", "l", "x"],
help="模型大小 (default: n)")
train_parser.add_argument("--epochs", type=int, default=100, help="训练轮数")
train_parser.add_argument("--imgsz", type=int, default=640, help="图像尺寸")
train_parser.add_argument("--batch", type=int, default=16, help="批量大小")
train_parser.add_argument("--device", type=str, default="0", help="训练设备")
train_parser.add_argument("--output", type=str, default="models/train", help="输出目录")
train_parser.add_argument("--resume", action="store_true", help="继续训练")
# Validate command
val_parser = subparsers.add_parser("validate", help="验证模型")
val_parser.add_argument("--model", type=str, required=True, help="模型路径")
val_parser.add_argument("--data", type=str, required=True, help="数据集配置")
val_parser.add_argument("--device", type=str, default="0", help="验证设备")
# Export command
export_parser = subparsers.add_parser("export", help="导出模型")
export_parser.add_argument("--model", type=str, required=True, help="模型路径")
export_parser.add_argument("--format", type=str, default="onnx",
choices=["onnx", "torchscript", "openvino", "tensorrt"],
help="导出格式")
args = parser.parse_args()
if args.command == "train":
train_yolo(
data_yaml=args.data,
model_size=args.model_size,
epochs=args.epochs,
imgsz=args.imgsz,
batch=args.batch,
device=args.device,
output_dir=args.output,
resume=args.resume,
)
elif args.command == "validate":
validate_model(args.model, args.data, args.device)
elif args.command == "export":
export_model(args.model, args.format)
else:
parser.print_help()
sys.exit(1)
if __name__ == "__main__":
main()
scripts/usage_example.py
+134
View File
@@ -0,0 +1,134 @@
#!/usr/bin/env python3
"""
红外故障检测模块使用示例
"""
import sys
from pathlib import Path
# 添加项目根目录到路径
sys.path.insert(0, str(Path(__file__).parent.parent))
from src.detect import InfraredDetector
from src.extract_temp import TempFeatureExtractor
from src.diagnose import FaultDiagnoser
import cv2
import json
def example1_basic_detection():
"""示例1：基础目标检测"""
print("\n【示例1】基础目标检测")
detector = InfraredDetector("weights/infrared_yolov8n.pt", conf=0.25)
detections = detector.detect("data/infrared/images/val/FLIR0166.jpg")
print(f"检测到 {len(detections)} 个目标:")
for i, det in enumerate(detections):
print(f" [{i+1}] {det['class_name']}: conf={det['confidence']:.3f}")
def example2_detect_and_draw():
"""示例2：检测并绘制结果"""
print("\n【示例2】检测并绘制结果")
output_dir = Path(__file__).parent.parent / "output"
output_dir.mkdir(exist_ok=True)
detector = InfraredDetector("weights/infrared_yolov8n.pt")
detector.detect_and_draw(
"data/infrared/images/val/FLIR0166.jpg",
output_path=str(output_dir / "detection_result.jpg")
)
print(f"结果已保存至: {output_dir / 'detection_result.jpg'}")
def example3_temperature_extraction():
"""示例3：温度特征提取"""
print("\n【示例3】温度特征提取")
extractor = TempFeatureExtractor("weights/infrared_yolov8n.pt")
detector = InfraredDetector("weights/infrared_yolov8n.pt")
# 读取图像
gray_img = cv2.imread("data/infrared/images/val/FLIR0166.jpg", cv2.IMREAD_GRAYSCALE)
# 检测目标
detections = detector.detect("data/infrared/images/val/FLIR0166.jpg")
# 提取第一个检测框的温度特征
if detections:
bbox = detections[0]['bbox']
x1, y1, x2, y2 = map(int, bbox)
roi = gray_img[y1:y2, x1:x2]
features = extractor.extract_from_roi(roi)
print(f" - 最高温度: {features['temp_max']:.2f}")
print(f" - 平均温度: {features['temp_mean']:.2f}")
print(f" - 温度标准差: {features['temp_std']:.2f}")
def example4_fault_diagnosis():
"""示例4：故障诊断"""
print("\n【示例4】故障诊断")
diagnoser = FaultDiagnoser("weights/infrared_yolov8n.pt")
result = diagnoser.diagnose("data/infrared/images/val/FLIR0166.jpg")
print(f"诊断状态: {result['summary']['status']}")
print(f"故障类型: {result['summary']['fault_types']}")
for fault in result['faults']:
print(f" - {fault['fault_type']}: temp_max={fault['temp_max']:.2f}")
def example5_batch_diagnosis():
"""示例5：批量诊断"""
print("\n【示例5】批量诊断")
output_dir = Path(__file__).parent.parent / "output"
output_dir.mkdir(exist_ok=True)
diagnoser = FaultDiagnoser("weights/infrared_yolov8n.pt")
diagnoser.batch_diagnose(
"data/infrared/images/val",
str(output_dir / "batch_results.json")
)
print(f"批量诊断完成，结果已保存至: {output_dir / 'batch_results.json'}")
def example6_custom_threshold():
"""示例6：自定义温度阈值"""
print("\n【示例6】自定义温度阈值")
# 自定义温度阈值
threshold = {
"max_temp": 200, # 最高温度阈值
"std_temp": 40 # 标准差阈值
}
diagnoser = FaultDiagnoser("weights/infrared_yolov8n.pt", temp_threshold=threshold)
result = diagnoser.diagnose("data/infrared/images/val/FLIR0166.jpg")
print(f"使用自定义阈值: max_temp={threshold['max_temp']}, std_temp={threshold['std_temp']}")
print(f"诊断结果: {result['summary']['status']}")
def main():
print("=" * 60)
print("红外故障检测模块使用示例")
print("=" * 60)
example1_basic_detection()
example2_detect_and_draw()
example3_temperature_extraction()
example4_fault_diagnosis()
example5_batch_diagnosis()
example6_custom_threshold()
print("\n" + "=" * 60)
print("所有示例执行完成!")
print("=" * 60)
if __name__ == "__main__":
main()
src/__init__.py
+1
View File
@@ -0,0 +1 @@
# Infrared Fault Detection Package
src/detect.py
+208
View File
@@ -0,0 +1,208 @@
"""
YOLO目标检测脚本
用于红外图像中的目标检测
"""
import os
from pathlib import Path
from typing import Optional, List, Dict, Union
import numpy as np
from ultralytics import YOLO
import cv2
class InfraredDetector:
"""红外图像目标检测器"""
def __init__(self, model_path: Optional[str] = None, conf: float = 0.25):
"""
初始化检测器
Args:
model_path: YOLO模型路径
conf: 置信度阈值
"""
self.model = None
self.conf = conf
if model_path and Path(model_path).exists():
self.load_model(model_path)
def load_model(self, model_path: str) -> None:
"""加载YOLO模型"""
self.model = YOLO(model_path)
print(f"Model loaded from: {model_path}")
def detect(self, image: Union[str, np.ndarray],
conf: Optional[float] = None) -> List[Dict]:
"""
对图像进行目标检测
Args:
image: 图像路径或图像数组
conf: 置信度阈值（覆盖默认值）
Returns:
检测结果列表
"""
if self.model is None:
raise ValueError("Model not loaded. Please load a model first.")
threshold = conf if conf is not None else self.conf
results = self.model(image, conf=threshold)
detections = []
for result in results:
boxes = result.boxes
if boxes is None:
continue
for box in boxes:
xyxy = box.xyxy[0].cpu().numpy()
xywh = box.xywh[0].cpu().numpy()
conf_score = float(box.conf[0])
cls_id = int(box.cls[0])
detections.append({
'bbox': xyxy.tolist(), # [x1, y1, x2, y2]
'bbox_xywh': xywh.tolist(), # [x_center, y_center, w, h]
'confidence': conf_score,
'class_id': cls_id,
'class_name': result.names[cls_id]
})
return detections
def detect_and_draw(self, image: Union[str, np.ndarray],
output_path: Optional[str] = None,
conf: Optional[float] = None) -> np.ndarray:
"""
检测并绘制结果
Args:
image: 图像路径或图像数组
output_path: 输出图像路径
conf: 置信度阈值
Returns:
绘制了检测结果的图像
"""
if isinstance(image, str):
img = cv2.imread(image)
else:
img = image.copy()
detections = self.detect(image, conf)
for det in detections:
x1, y1, x2, y2 = map(int, det['bbox'])
conf_score = det['confidence']
class_name = det['class_name']
color = (0, 255, 0)
cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
label = f"{class_name} {conf_score:.2f}"
cv2.putText(img, label, (x1, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
if output_path:
cv2.imwrite(output_path, img)
print(f"Result saved to: {output_path}")
return img
def train_yolo_model(data_yaml: str,
model_size: str = 'n',
epochs: int = 100,
imgsz: int = 640,
batch: int = 16,
output_dir: str = 'models/train') -> str:
"""
训练YOLO模型
Args:
data_yaml: 数据集配置文件路径
model_size: 模型大小 ('n', 's', 'm', 'l', 'x')
epochs: 训练轮数
imgsz: 输入图像大小
batch: 批量大小
output_dir: 输出目录
Returns:
最佳模型路径
"""
model = YOLO(f'yolov8{model_size}.pt')
results = model.train(
data=data_yaml,
epochs=epochs,
imgsz=imgsz,
batch=batch,
project=output_dir,
name='infrared_fault',
exist_ok=True,
patience=20,
save=True,
plots=True
)
best_model_path = f"{output_dir}/infrared_fault/weights/best.pt"
return best_model_path
def validate_model(model_path: str, data_yaml: str) -> Dict:
"""
验证模型性能
Args:
model_path: 模型路径
data_yaml: 数据集配置路径
Returns:
验证指标
"""
model = YOLO(model_path)
metrics = model.val(data=data_yaml)
return {
'mAP50': float(metrics.box.map50),
'mAP50-95': float(metrics.box.map),
'precision': float(metrics.box.mp),
'recall': float(metrics.box.mr)
}
if __name__ == '__main__':
import argparse
parser = argparse.ArgumentParser(description='YOLO Infrared Detection')
subparsers = parser.add_subparsers(dest='command', help='Commands')
# Detect command
detect_parser = subparsers.add_parser('detect', help='Detect objects in image')
detect_parser.add_argument('--model', type=str, required=True, help='Model path')
detect_parser.add_argument('--image', type=str, required=True, help='Image path')
detect_parser.add_argument('--output', type=str, help='Output image path')
detect_parser.add_argument('--conf', type=float, default=0.25, help='Confidence threshold')
# Train command
train_parser = subparsers.add_parser('train', help='Train YOLO model')
train_parser.add_argument('--data', type=str, required=True, help='Data yaml path')
train_parser.add_argument('--model-size', type=str, default='n', choices=['n', 's', 'm', 'l', 'x'])
train_parser.add_argument('--epochs', type=int, default=100)
train_parser.add_argument('--imgsz', type=int, default=640)
train_parser.add_argument('--batch', type=int, default=16)
args = parser.parse_args()
if args.command == 'detect':
detector = InfraredDetector(args.model, args.conf)
img = detector.detect_and_draw(args.image, args.output)
print(f"Detected {len(detector.detect(args.image))} objects")
elif args.command == 'train':
best_path = train_yolo_model(
args.data, args.model_size, args.epochs, args.imgsz, args.batch
)
print(f"Training complete. Best model: {best_path}")
src/diagnose.py
+203
View File
@@ -0,0 +1,203 @@
"""
故障诊断模块
基于检测结果和温度特征进行故障诊断
"""
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)
src/extract_temp.py
+184
View File
@@ -0,0 +1,184 @@
"""
ROI温度特征提取模块
基于YOLO检测结果，提取ROI区域的温度特征
"""
import numpy as np
import cv2
import json
from pathlib import Path
from typing import Dict, List, Optional, Union
from ultralytics import YOLO
from .utils.temp_utils import (
extract_temp_stats,
extract_temp_histogram,
calculate_temp_gradient,
normalize_roi
)
class TempFeatureExtractor:
"""红外图像ROI温度特征提取器"""
def __init__(self, model_path: Optional[str] = None):
"""
初始化温度特征提取器
Args:
model_path: YOLO模型路径，如果为None则只提取温度特征
"""
self.model = None
if model_path and Path(model_path).exists():
self.model = YOLO(model_path)
def detect_and_extract(self, image: np.ndarray,
conf: float = 0.25) -> List[Dict]:
"""
对图像进行检测并提取ROI温度特征
Args:
image: 输入红外图像
conf: 检测置信度阈值
Returns:
检测结果列表，每个元素包含检测框和温度特征
"""
if self.model is None:
raise ValueError("Model not loaded. Please provide a valid model path.")
results = self.model(image, conf=conf)
detections = []
for result in results:
boxes = result.boxes
if boxes is None:
continue
for box in boxes:
xyxy = box.xyxy[0].cpu().numpy()
conf_score = float(box.conf[0])
cls_id = int(box.cls[0])
x1, y1, x2, y2 = map(int, xyxy)
roi = image[y1:y2, x1:x2]
temp_features = extract_temp_stats(roi)
temp_features.update({
'bbox': [int(x1), int(y1), int(x2), int(y2)],
'confidence': conf_score,
'class_id': cls_id,
'class_name': result.names[cls_id]
})
detections.append(temp_features)
return detections
def extract_from_roi(self, roi: np.ndarray,
include_hist: bool = True,
include_gradient: bool = True) -> Dict:
"""
从指定ROI提取温度特征
Args:
roi: ROI区域图像
include_hist: 是否包含直方图特征
include_gradient: 是否包含梯度特征
Returns:
温度特征字典
"""
features = extract_temp_stats(roi)
if include_hist:
features['histogram'] = extract_temp_histogram(roi).tolist()
if include_gradient:
features.update(calculate_temp_gradient(roi))
return features
def extract_from_bbox(self, image: np.ndarray,
bbox: List[int]) -> Dict:
"""
从给定边界框提取温度特征
Args:
image: 输入图像
bbox: [x1, y1, x2, y2] 边界框坐标
Returns:
温度特征字典
"""
x1, y1, x2, y2 = bbox
roi = image[y1:y2, x1:x2]
return self.extract_from_roi(roi)
def load_infrared_image(image_path: str) -> np.ndarray:
"""
加载红外图像
Args:
image_path: 图像路径
Returns:
图像数组
"""
image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
if image is None:
raise ValueError(f"Failed to load image: {image_path}")
return image
def batch_extract_features(image_dir: Union[str, Path],
output_path: Union[str, Path],
model_path: Optional[str] = None,
conf: float = 0.25) -> None:
"""
批量处理图像并提取温度特征
Args:
image_dir: 图像目录路径
output_path: 输出JSON文件路径
model_path: YOLO模型路径
conf: 检测置信度阈值
"""
image_dir = Path(image_dir)
extractor = TempFeatureExtractor(model_path)
all_results = []
for img_path in sorted(image_dir.glob('*')):
if img_path.suffix.lower() not in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']:
continue
try:
image = load_infrared_image(str(img_path))
detections = extractor.detect_and_extract(image, conf=conf)
all_results.append({
'image': img_path.name,
'detections': detections
})
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(all_results, f, indent=2, ensure_ascii=False)
print(f"Processed {len(all_results)} images, results saved to {output_path}")
if __name__ == '__main__':
import argparse
parser = argparse.ArgumentParser(description='Extract temperature features from infrared images')
parser.add_argument('--image-dir', type=str, required=True, help='Input image directory')
parser.add_argument('--output', type=str, required=True, help='Output JSON path')
parser.add_argument('--model', type=str, default=None, help='YOLO model path')
parser.add_argument('--conf', type=float, default=0.25, help='Detection confidence threshold')
args = parser.parse_args()
batch_extract_features(args.image_dir, args.output, args.model, args.conf)
src/utils/__init__.py
+1
View File
@@ -0,0 +1 @@
# Utils Package
src/utils/temp_utils.py
+127
View File
@@ -0,0 +1,127 @@
"""
温度特征提取工具函数
从红外图像ROI区域提取温度统计特征
"""
import numpy as np
import cv2
from typing import Dict, List, Tuple, Optional
def extract_temp_stats(roi: np.ndarray) -> Dict[str, float]:
"""
从ROI区域提取温度统计特征
Args:
roi: ROI区域图像数组
Returns:
温度统计特征字典
"""
if roi.size == 0:
return {
'temp_max': 0.0,
'temp_min': 0.0,
'temp_mean': 0.0,
'temp_std': 0.0,
'temp_median': 0.0,
}
return {
'temp_max': float(np.max(roi)),
'temp_min': float(np.min(roi)),
'temp_mean': float(np.mean(roi)),
'temp_std': float(np.std(roi)),
'temp_median': float(np.median(roi)),
}
def extract_temp_histogram(roi: np.ndarray, bins: int = 32) -> np.ndarray:
"""
提取温度分布直方图特征
Args:
roi: ROI区域图像数组
bins: 直方图bin数量
Returns:
归一化直方图向量
"""
if roi.size == 0:
return np.zeros(bins)
hist, _ = np.histogram(roi.flatten(), bins=bins, range=(0, 255))
hist = hist.astype(np.float32) / (hist.sum() + 1e-7)
return hist
def calculate_temp_gradient(roi: np.ndarray) -> Dict[str, float]:
"""
计算温度梯度特征（反映温度变化剧烈程度）
Args:
roi: ROI区域图像数组
Returns:
温度梯度统计
"""
if roi.size == 0:
return {'gradient_mean': 0.0, 'gradient_std': 0.0}
grad_x = cv2.Sobel(roi, cv2.CV_32F, 1, 0, ksize=3)
grad_y = cv2.Sobel(roi, cv2.CV_32F, 0, 1, ksize=3)
gradient_magnitude = np.sqrt(grad_x**2 + grad_y**2)
return {
'gradient_mean': float(np.mean(gradient_magnitude)),
'gradient_std': float(np.std(gradient_magnitude)),
}
def extract_temp_features(roi: np.ndarray, include_hist: bool = True,
include_gradient: bool = True) -> Dict:
"""
从ROI区域提取完整温度特征
Args:
roi: ROI区域图像数组
include_hist: 是否包含直方图特征
include_gradient: 是否包含梯度特征
Returns:
完整温度特征字典
"""
features = extract_temp_stats(roi)
if include_hist:
features['temp_histogram'] = extract_temp_histogram(roi)
if include_gradient:
features.update(calculate_temp_gradient(roi))
return features
def normalize_roi(roi: np.ndarray, method: str = 'minmax') -> np.ndarray:
"""
对ROI进行归一化处理
Args:
roi: ROI区域图像数组
method: 归一化方法 ('minmax' 或 'zscore')
Returns:
归一化后的图像
"""
if method == 'minmax':
min_val, max_val = roi.min(), roi.max()
if max_val - min_val < 1e-7:
return np.zeros_like(roi)
return (roi - min_val) / (max_val - min_val)
elif method == 'zscore':
mean_val, std_val = np.mean(roi), np.std(roi)
if std_val < 1e-7:
return np.zeros_like(roi)
return (roi - mean_val) / std_val
else:
raise ValueError(f"Unknown normalization method: {method}")
工作计划.md
+82
View File
@@ -0,0 +1,82 @@
# 红外故障检测工作计划
## 项目目标
基于YOLO的红外图像目标识别与ROI区域温度特征提取系统
## 项目背景
- 现有数据集：红外图像数据（电气设备）
- 目标：识别红外图像中的设备部件，提取ROI区域的温度特征，用于后续故障诊断
## 工作计划
### 阶段一：项目初始化与环境配置
- [ ] 创建项目目录结构
- [ ] 配置Python环境（Python 3.8+, PyTorch, OpenCV, NumPy）
- [ ] 安装YOLO相关依赖（ultralytics）
- [ ] 整理红外图像数据集，转换为YOLO格式
- [ ] 编写数据配置文件（data.yaml）
### 阶段二：YOLO模型训练与优化
- [ ] 选择合适的YOLO模型（YOLOv8）
- [ ] 编写模型训练脚本
- [ ] 进行模型训练与验证
- [ ] 评估模型性能（mAP, Precision, Recall）
- [ ] 导出训练好的模型（best.pt）
### 阶段三：ROI温度特征提取
- [ ] 设计ROI区域温度提取算法
- [ ] 实现基于检测框的温度统计（最大值、最小值、平均值、标准差）
- [ ] 实现温度分布直方图特征
- [ ] 实现纹理特征（可选：基于GLCM）
- [ ] 封装温度特征提取模块
### 阶段四：故障诊断接口
- [ ] 设计故障诊断流程
- [ ] 封装推理Pipeline（检测+温度提取+诊断）
- [ ] 输出结构化结果（JSON格式）
### 阶段五：测试与验证
- [ ] 准备测试数据集
- [ ] 端到端流程测试
- [ ] 性能优化（推理速度）
## 项目结构
```
1BaseOnYOLO/
├── data/ # 数据目录
│ ├── images/ # 红外图像
│ ├── labels/ # YOLO格式标注
│ └── yaml/ # 配置文件
├── models/ # 模型目录
│ └── train/ # 训练输出
├── src/ # 源代码
│ ├── detect.py # 检测脚本
│ ├── extract_temp.py # 温度特征提取
│ ├── diagnose.py # 故障诊断
│ └── utils/ # 工具函数
├── weights/ # 模型权重
├── 工作计划.md
└── README.md
```
## 技术栈
- **深度学习框架**: PyTorch
- **目标检测**: YOLOv8 (ultralytics)
- **图像处理**: OpenCV, NumPy
- **数据格式**: YOLO TXT格式
## 预期输出
1. 训练好的YOLO模型权重
2. ROI温度特征提取模块
3. 故障诊断Pipeline
4. 完整的项目代码和文档