From 787d17a870debfbe719309578c1e0c96fc146677 Mon Sep 17 00:00:00 2001 From: SkyLi Date: Fri, 27 Mar 2026 00:56:10 +0800 Subject: [PATCH] chore: update project docs and scripts for new dataset MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 更新README.md数据集描述（12类，6735训练图+765验证图） - 更新模型性能指标（新12类模型结果） - 删除已过时的download_dataset.py和retry_download.py脚本 - 更新usage_example.py使用正确的路径和模型名 --- README.md | 79 +++++++++---------- scripts/download_dataset.py | 152 ------------------------------------ scripts/retry_download.py | 152 ------------------------------------ scripts/usage_example.py | 31 ++++---- 4 files changed, 56 insertions(+), 358 deletions(-) delete mode 100644 scripts/download_dataset.py delete mode 100644 scripts/retry_download.py diff --git a/README.md b/README.md index bf9c07b..1a0a62e 100644 --- a/README.md +++ b/README.md @@ -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) ``` --- diff --git a/scripts/download_dataset.py b/scripts/download_dataset.py deleted file mode 100644 index 411735d..0000000 --- a/scripts/download_dataset.py +++ /dev/null @@ -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() diff --git a/scripts/retry_download.py b/scripts/retry_download.py deleted file mode 100644 index 52637a3..0000000 --- a/scripts/retry_download.py +++ /dev/null @@ -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() diff --git a/scripts/usage_example.py b/scripts/usage_example.py index 4ee37d7..66b038c 100644 --- a/scripts/usage_example.py +++ b/scripts/usage_example.py @@ -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']}")