Files
Infrared-ai-research-lab/scripts/retry_download.py
T
SkyLi fd1a5f4050 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
2026-03-26 21:36:08 +08:00

153 lines
4.1 KiB
Python

#!/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()