SkyLi
fd1a5f4050
- 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
153 lines
3.9 KiB
Python
153 lines
3.9 KiB
Python
#!/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()
|