import json import requests from pathlib import Path from tqdm import tqdm import concurrent.futures # =====================配置===================== NDJSON_FILE = "dataset.ndjson" # 你的ndjson文本保存文件名 SAVE_ROOT = Path("shuttlecock_dataset") THREAD_COUNT = 10 # 并发下载线程 # ============================================== # 创建目录 for part in ["train", "val", "test"]: (SAVE_ROOT / "images" / part).mkdir(parents=True, exist_ok=True) (SAVE_ROOT / "labels" / part).mkdir(parents=True, exist_ok=True) def task_download(line): try: data = json.loads(line.strip()) if data["type"] != "image": return img_url = data["url"] img_name = data["file"] split = data["split"] boxes = data["annotations"]["boxes"] img_save_path = SAVE_ROOT / "images" / split / img_name label_save_path = SAVE_ROOT / "labels" / split / Path(img_name).with_suffix(".txt") # 下载图片 resp = requests.get(img_url, timeout=20) if resp.status_code == 200: with open(img_save_path, "wb") as f: f.write(resp.content) else: print(f"下载失败 {img_name} code:{resp.status_code}") return # 写入YOLO标签 label_lines = [] for box in boxes: cls, cx, cy, w, h = box label_lines.append(f"{cls} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}") with open(label_save_path, "w", encoding="utf-8") as f: f.write("\n".join(label_lines)) except Exception as err: print(f"异常: {err}") if __name__ == "__main__": with open(NDJSON_FILE, "r", encoding="utf-8") as f: lines = [ln for ln in f.readlines() if ln.strip()] print(f"总共待处理图片：{len(lines)}") with concurrent.futures.ThreadPoolExecutor(max_workers=THREAD_COUNT) as pool: list(tqdm(pool.map(task_download, lines), total=len(lines))) print("✅ 所有任务执行完毕！")
