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

chore: update project docs and scripts for new dataset

Browse Source 
- 更新README.md数据集描述（12类，6735训练图+765验证图）
- 更新模型性能指标（新12类模型结果）
- 删除已过时的download_dataset.py和retry_download.py脚本
- 更新usage_example.py使用正确的路径和模型名
This commit is contained in:
SkyLi
2026-03-27 00:56:10 +08:00
parent 11c5e062c2
commit 787d17a870
4 changed files with 56 additions and 358 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+39 -40
View File
@@ -94,23 +94,30 @@ cuDNN version: 9.1.0
### 数据来源
SCIDB红外电力设备数据集: https://www.scidb.cn/en/detail?dataSetId=e416c488169f484485ad7575dcfc43ce
电力设备红外图像数据集（12类）
### 数据规模
| 类别 | 图像数量 |
|------|----------|
| Wave Traps | 153 |
| Power Transformers | 178 |
| Circuit Breakers | 162 |
| Surge Arresters | 181 |
| Disconnectors | 89 |
| **总计** | **763** |
| 类别 | 类别ID |
|------|--------|
| bushing | 0 |
| arrester | 1 |
| breaker | 2 |
| clamp | 3 |
| conservator | 4 |
| current-transformer | 5 |
| disconnector | 6 |
| disconnector2 | 7 |
| disconnector3 | 8 |
| heat-sink | 9 |
| insulator | 10 |
| transformer | 11 |
| **总计** | **12类** |
### 数据划分
- 训练集: 610张 (80%)
- 验证集: 153张 (20%)
- 训练集: 6735张
- 验证集: 765张
---
@@ -118,22 +125,7 @@ SCIDB红外电力设备数据集: https://www.scidb.cn/en/detail?dataSetId=e416c
### 完整训练流程
**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 \
@@ -160,20 +152,27 @@ uv run python scripts/train.py train \
| 指标 | 值 |
|------|-----|
| mAP@0.5 | **99.5%** |
| mAP@0.5:0.95 | **99.5%** |
| Precision | **99.7%** |
| Recall | **99.4%** |
| mAP@0.5 | **76.42%** |
| mAP@0.5:0.95 | **61.64%** |
| Precision | **66.58%** |
| Recall | **87.93%** |
**各类别性能**:
| 类别 | 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 |
| bushing | 0.925 | 0.990 | 0.994 |
| arrester | 0.930 | 0.980 | 0.991 |
| breaker | 0.797 | 0.910 | 0.898 |
| clamp | 0.781 | 0.973 | 0.918 |
| conservator | 0.546 | 0.823 | 0.718 |
| current-transformer | 0.856 | 0.977 | 0.951 |
| disconnector | 0.651 | 0.975 | 0.813 |
| disconnector2 | 0.306 | 0.767 | 0.286 |
| disconnector3 | 0.566 | 0.940 | 0.886 |
| heat-sink | 0.612 | 0.429 | 0.546 |
| insulator | 0.731 | 0.966 | 0.880 |
| transformer | 0.289 | 0.822 | 0.290 |
---
@@ -209,7 +208,7 @@ for det in detections:
'bbox_xywh': [x_center, y_center, w, h], # 中心格式
'confidence': 0.995,
'class_id': 0,
'class_name': 'Wave Traps'
'class_name': 'bushing'
}
```
@@ -270,13 +269,13 @@ print(f"故障类型: {result['summary']['fault_types']}")
**诊断输出格式**:
```json
{
"image": "FLIR0150.jpg",
"image": "20230907095711003_jpg.rf.0e14e8b4344dce21114683dc3f3e7e2e.jpg",
"detections": [
{
"bbox": [x1, y1, x2, y2],
"confidence": 0.995,
"class_id": 0,
"class_name": "Wave Traps"
"class_name": "bushing"
}
],
"faults": [
@@ -552,7 +551,7 @@ A:
A:
```python
threshold = {"max_temp": 200, "std_temp": 40}
diagnoser = FaultDiagnoser("weights/infrared_yolov8n.pt", temp_threshold=threshold)
diagnoser = FaultDiagnoser("weights/best.pt", temp_threshold=threshold)
```
---
scripts/download_dataset.py
-152
View File
@@ -1,152 +0,0 @@
#!/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
@@ -1,152 +0,0 @@
#!/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/usage_example.py
+17 -14
View File
@@ -19,8 +19,8 @@ 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")
detector = InfraredDetector("weights/best.pt", conf=0.25)
detections = detector.detect("data/valid/images/20230907095711003_jpg.rf.0e14e8b4344dce21114683dc3f3e7e2e.jpg")
print(f"检测到 {len(detections)} 个目标:")
for i, det in enumerate(detections):
@@ -34,9 +34,9 @@ def example2_detect_and_draw():
output_dir = Path(__file__).parent.parent / "output"
output_dir.mkdir(exist_ok=True)
detector = InfraredDetector("weights/infrared_yolov8n.pt")
detector = InfraredDetector("weights/best.pt")
detector.detect_and_draw(
"data/infrared/images/val/FLIR0166.jpg",
"data/valid/images/20230907095711003_jpg.rf.0e14e8b4344dce21114683dc3f3e7e2e.jpg",
output_path=str(output_dir / "detection_result.jpg")
)
print(f"结果已保存至: {output_dir / 'detection_result.jpg'}")
@@ -46,14 +46,15 @@ def example3_temperature_extraction():
"""示例3：温度特征提取"""
print("\n【示例3】温度特征提取")
extractor = TempFeatureExtractor("weights/infrared_yolov8n.pt")
detector = InfraredDetector("weights/infrared_yolov8n.pt")
sample_image = "data/valid/images/20230907095711003_jpg.rf.0e14e8b4344dce21114683dc3f3e7e2e.jpg"
extractor = TempFeatureExtractor("weights/best.pt")
detector = InfraredDetector("weights/best.pt")
# 读取图像
gray_img = cv2.imread("data/infrared/images/val/FLIR0166.jpg", cv2.IMREAD_GRAYSCALE)
gray_img = cv2.imread(sample_image, cv2.IMREAD_GRAYSCALE)
# 检测目标
detections = detector.detect("data/infrared/images/val/FLIR0166.jpg")
detections = detector.detect(sample_image)
# 提取第一个检测框的温度特征
if detections:
@@ -71,8 +72,9 @@ def example4_fault_diagnosis():
"""示例4：故障诊断"""
print("\n【示例4】故障诊断")
diagnoser = FaultDiagnoser("weights/infrared_yolov8n.pt")
result = diagnoser.diagnose("data/infrared/images/val/FLIR0166.jpg")
sample_image = "data/valid/images/20230907095711003_jpg.rf.0e14e8b4344dce21114683dc3f3e7e2e.jpg"
diagnoser = FaultDiagnoser("weights/best.pt")
result = diagnoser.diagnose(sample_image)
print(f"诊断状态: {result['summary']['status']}")
print(f"故障类型: {result['summary']['fault_types']}")
@@ -88,9 +90,9 @@ def example5_batch_diagnosis():
output_dir = Path(__file__).parent.parent / "output"
output_dir.mkdir(exist_ok=True)
diagnoser = FaultDiagnoser("weights/infrared_yolov8n.pt")
diagnoser = FaultDiagnoser("weights/best.pt")
diagnoser.batch_diagnose(
"data/infrared/images/val",
"data/valid/images",
str(output_dir / "batch_results.json")
)
print(f"批量诊断完成，结果已保存至: {output_dir / 'batch_results.json'}")
@@ -100,14 +102,15 @@ def example6_custom_threshold():
"""示例6：自定义温度阈值"""
print("\n【示例6】自定义温度阈值")
sample_image = "data/valid/images/20230907095711003_jpg.rf.0e14e8b4344dce21114683dc3f3e7e2e.jpg"
# 自定义温度阈值
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")
diagnoser = FaultDiagnoser("weights/best.pt", temp_threshold=threshold)
result = diagnoser.diagnose(sample_image)
print(f"使用自定义阈值: max_temp={threshold['max_temp']}, std_temp={threshold['std_temp']}")
print(f"诊断结果: {result['summary']['status']}")