feat(detect): 把弱监督框数据集链接到 data/ 并接上 YOLO 训练
第五步产出的框标签按类别分目录，便于人工核对；检测训练要的是 train/val 划分——类别写在标签第一列，不靠目录名。补上两个入口： - scripts/prepare_yolo_dataset.py：拍平类别目录、按类别分层切分、硬链接到 data/。分层是必需的，black 只有 53 张，随机切容易让小类在验证集里挂零。 无框的图默认保留（检测训练里是压低误检的负样本）；拍平后的同名冲突直接 报错而不是静默覆盖。 - detect.py + configs/detect/film.yaml：包一层 ultralytics，沿用本仓库的 配置继承、设备解析与 SwanLab 可视化。 选 YOLOv8 而不是 v11/v12：CLAUDE.md 限定只用 convolution 模型，YOLO11 骨架 里有 C2PSA、YOLO12 以注意力为主，只有 v8 的 C2f + SPPF 是纯卷积。 mosaic/mixup 关掉——每张图只有一个弱监督框，拼接后小目标被缩得更小，对 blackdot/whitedot 这类只占千分之几面积的类别是净伤害。 ultralytics 的 project 必须传绝对路径，相对路径会被它再拼到全局 settings 的 runs_dir 上，产出会跑到 runs/detect/runs/detect/<name> 去。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
86c86b105c
commit
8dc2fe6a2e
@@ -0,0 +1,220 @@
|
||||
#!/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/<name>/
|
||||
├── 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/<source 目录名>")
|
||||
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())
|
||||
Reference in New Issue
Block a user