#!/usr/bin/env python """把 generate_boxes.py 的产出整理成 YOLO 训练所需的数据集，放到 data/ 下。 uv run python scripts/prepare_yolo_dataset.py --source runs/boxes_gradcam generate_boxes.py 的产出是**按类别**分目录的（便于人工核对）： ``` runs/boxes_gradcam/ ├── images/<类别>/<图名>.png └── labels/<类别>/<图名>.txt ``` 而 YOLO 检测训练要的是 **train/val 划分**——类别信息在标签文件的第一列里， 不靠目录名。这一步做三件事： 1. 把按类别分的目录拍平，按 train/val 重新组织； 2. **按类别分层**切分，小类不至于在验证集里挂零； 3. 全部用**硬链接**，不占额外磁盘。 ``` data// ├── images/{train,val}/<图名>.png ├── labels/{train,val}/<图名>.txt └── data.yaml ``` 没有框的图（标签为空）默认保留：YOLO 训练里它们是负样本，有助于压低误检。 """ from __future__ import annotations import argparse import errno import json import os import random import shutil import sys from collections import Counter, defaultdict from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parents[1])) import yaml # noqa: E402 from datasets.yolo import IMAGE_SUFFIXES, write_data_yaml # noqa: E402 from utils import setup_logging # noqa: E402 logger = None def link_or_copy(source: Path, destination: Path) -> str: """硬链接；跨文件系统时退回复制。""" if destination.exists() or destination.is_symlink(): destination.unlink() destination.parent.mkdir(parents=True, exist_ok=True) try: os.link(source, destination) return "hard" except OSError as exc: if exc.errno != errno.EXDEV: raise shutil.copy2(source, destination) return "copy" def first_class_id(label_path: Path) -> int | None: """取标签里第一个框的类别下标；空标签（无框）返回 None。""" for line in label_path.read_text(encoding="utf-8").splitlines(): line = line.strip() if line: return int(float(line.split()[0])) return None def collect(source: Path) -> list[tuple[Path, Path, int | None]]: """配对 images/ 与 labels/，返回 (图片, 标签, 类别下标) 列表。""" images_root, labels_root = source / "images", source / "labels" if not images_root.is_dir() or not labels_root.is_dir(): raise FileNotFoundError(f"{source} 下缺少 images/ 或 labels/，先跑 scripts/generate_boxes.py") pairs = [] missing = [] for image in sorted(p for p in images_root.rglob("*") if p.suffix.lower() in IMAGE_SUFFIXES): label = labels_root / image.relative_to(images_root).with_suffix(".txt") if not label.exists(): missing.append(image) continue pairs.append((image, label, first_class_id(label))) if missing: raise FileNotFoundError( f"有 {len(missing)} 张图没有对应的标签，例如 {missing[0]}；" "产出目录可能不完整，重新跑一次 generate_boxes.py" ) if not pairs: raise ValueError(f"{source} 下没有找到任何图片") return pairs def main() -> int: global logger parser = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter ) parser.add_argument("--source", required=True, help="generate_boxes.py 的产出目录") parser.add_argument("--out", default=None, help="输出目录；缺省 data/") parser.add_argument("--data-root", default="data", help="自动推导时的根目录，默认 data") parser.add_argument("--val-ratio", type=float, default=0.2) parser.add_argument("--seed", type=int, default=42) parser.add_argument( "--drop-empty", action="store_true", help="丢弃没有框的图。默认保留——YOLO 训练里它们是负样本，有助于压低误检", ) parser.add_argument("--overwrite", action="store_true", help="输出目录已存在时先清空") parser.add_argument("--dry-run", action="store_true", help="只打印计划") args = parser.parse_args() source = Path(args.source) out_dir = Path(args.out) if args.out else Path(args.data_root) / source.name logger = setup_logging() try: pairs = collect(source) except (FileNotFoundError, ValueError) as exc: logger.error("%s", exc) return 1 meta_path = source / "data.yaml" if not meta_path.exists(): logger.error("找不到 %s", meta_path) return 1 meta = yaml.safe_load(meta_path.read_text(encoding="utf-8")) names = [meta["names"][i] for i in sorted(meta["names"])] empty = [p for p in pairs if p[2] is None] if args.drop_empty: pairs = [p for p in pairs if p[2] is not None] logger.info("来源: %s（%d 张，%d 类）", source, len(pairs), len(names)) logger.info( "无框图片 %d 张：%s", len(empty), "已丢弃" if args.drop_empty else "保留为负样本" ) # 按类别分层切分；无框的图单独作为一组分配，避免它们全落进同一侧 grouped: dict[int | None, list] = defaultdict(list) for pair in pairs: grouped[pair[2]].append(pair) rng = random.Random(args.seed) splits: dict[str, list] = {"train": [], "val": []} for _class_id, group in sorted(grouped.items(), key=lambda kv: (kv[0] is None, kv[0])): order = group[:] rng.shuffle(order) n_val = max(1, round(len(order) * args.val_ratio)) if len(order) > 1 else 0 splits["val"].extend(order[:n_val]) splits["train"].extend(order[n_val:]) counts = { split: Counter(names[c] if c is not None else "(无框)" for _, _, c in items) for split, items in splits.items() } logger.info("划分: train %d / val %d", len(splits["train"]), len(splits["val"])) logger.info("%-14s %8s %8s", "类别", "train", "val") for name in names + ["(无框)"]: if counts["train"][name] or counts["val"][name]: logger.info("%-14s %8d %8d", name, counts["train"][name], counts["val"][name]) if args.dry_run: logger.info("--dry-run：未写入任何文件") return 0 if out_dir.exists() and any(out_dir.iterdir()): if not args.overwrite: logger.error("%s 已存在且非空；确认要覆盖请加 --overwrite", out_dir) return 1 shutil.rmtree(out_dir) modes: Counter[str] = Counter() seen: dict[str, Path] = {} for split, items in splits.items(): for image, label, _ in items: # 拍平之后不同类别目录下的同名文件会撞车，显式挡住 if image.name in seen: logger.error( "拍平后文件名冲突：%s 与 %s 同名。源数据集的文件名需要全局唯一，" "用 scripts/export_clusters.py --prefix-parent 重新导出即可", seen[image.name], image, ) return 1 seen[image.name] = image modes[link_or_copy(image, out_dir / "images" / split / image.name)] += 1 link_or_copy(label, out_dir / "labels" / split / f"{image.stem}.txt") write_data_yaml(out_dir / "data.yaml", names, root=out_dir.resolve(), train="images/train", val="images/val") (out_dir / "summary.json").write_text( json.dumps( { "source": str(source.resolve()), "num_images": len(pairs), "empty_labels": len(empty), "dropped_empty": args.drop_empty, "val_ratio": args.val_ratio, "seed": args.seed, "split_sizes": {k: len(v) for k, v in splits.items()}, "per_class": {k: dict(v) for k, v in counts.items()}, "link_modes": dict(modes), }, ensure_ascii=False, indent=2, ), encoding="utf-8", ) logger.info("链接方式: %s", dict(modes)) logger.info("完成 -> %s", out_dir) return 0 if __name__ == "__main__": raise SystemExit(main())