commit fd1a5f40502a421d70c93945f93b3999a395e2a5 Author: SkyLi Date: Thu Mar 26 21:36:08 2026 +0800 chore: initial commit for Infrared Equipment Detection System - 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 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..7126975 --- /dev/null +++ b/.gitignore @@ -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 diff --git a/NEW_PROJECT_PROMPT.md b/NEW_PROJECT_PROMPT.md new file mode 100644 index 0000000..3ae1ec0 --- /dev/null +++ b/NEW_PROJECT_PROMPT.md @@ -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 --model-size n --epochs 50 --device 0 + +# 验证 +uv run python scripts/train.py validate --model <模型路径> --data + +# 导出模型 +uv run python scripts/train.py export --model <模型路径> --format onnx +``` diff --git a/PROGRESS.md b/PROGRESS.md new file mode 100644 index 0000000..f61440d --- /dev/null +++ b/PROGRESS.md @@ -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%，各类别均衡 diff --git a/README.md b/README.md new file mode 100644 index 0000000..bf9c07b --- /dev/null +++ b/README.md @@ -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 diff --git a/data/dataset.yaml b/data/dataset.yaml new file mode 100644 index 0000000..cf93d17 --- /dev/null +++ b/data/dataset.yaml @@ -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 diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..3dd3bfc --- /dev/null +++ b/pyproject.toml @@ -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" }] diff --git a/scripts/convert_to_yolo.py b/scripts/convert_to_yolo.py new file mode 100644 index 0000000..53ef36a --- /dev/null +++ b/scripts/convert_to_yolo.py @@ -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() diff --git a/scripts/download_dataset.py b/scripts/download_dataset.py new file mode 100644 index 0000000..411735d --- /dev/null +++ b/scripts/download_dataset.py @@ -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() diff --git a/scripts/retry_download.py b/scripts/retry_download.py new file mode 100644 index 0000000..52637a3 --- /dev/null +++ b/scripts/retry_download.py @@ -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() diff --git a/scripts/test_cuda.py b/scripts/test_cuda.py new file mode 100644 index 0000000..68e6387 --- /dev/null +++ b/scripts/test_cuda.py @@ -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()) diff --git a/scripts/train.py b/scripts/train.py new file mode 100644 index 0000000..3c441bd --- /dev/null +++ b/scripts/train.py @@ -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() diff --git a/scripts/usage_example.py b/scripts/usage_example.py new file mode 100644 index 0000000..4ee37d7 --- /dev/null +++ b/scripts/usage_example.py @@ -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() diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..e00194a --- /dev/null +++ b/src/__init__.py @@ -0,0 +1 @@ +# Infrared Fault Detection Package diff --git a/src/detect.py b/src/detect.py new file mode 100644 index 0000000..0b6e403 --- /dev/null +++ b/src/detect.py @@ -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}") diff --git a/src/diagnose.py b/src/diagnose.py new file mode 100644 index 0000000..1b0f266 --- /dev/null +++ b/src/diagnose.py @@ -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) diff --git a/src/extract_temp.py b/src/extract_temp.py new file mode 100644 index 0000000..833fabd --- /dev/null +++ b/src/extract_temp.py @@ -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) diff --git a/src/utils/__init__.py b/src/utils/__init__.py new file mode 100644 index 0000000..1ef7647 --- /dev/null +++ b/src/utils/__init__.py @@ -0,0 +1 @@ +# Utils Package diff --git a/src/utils/temp_utils.py b/src/utils/temp_utils.py new file mode 100644 index 0000000..24eeecf --- /dev/null +++ b/src/utils/temp_utils.py @@ -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}") diff --git a/工作计划.md b/工作计划.md new file mode 100644 index 0000000..014ff7e --- /dev/null +++ b/工作计划.md @@ -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. 完整的项目代码和文档