From 788bc36ab8eb3352f87e9b1803dc9e074661e116 Mon Sep 17 00:00:00 2001 From: Shaun Zhou Date: Thu, 17 Sep 2026 21:47:02 +0800 Subject: [PATCH] =?UTF-8?q?feat(detect):=20=E5=85=A8=E9=87=8F=E6=8E=A8?= =?UTF-8?q?=E7=90=86=E5=8F=AF=E8=A7=86=E5=8C=96=EF=BC=8C=E9=A2=84=E6=B5=8B?= =?UTF-8?q?=E6=A1=86=E4=B8=8E=E7=9C=9F=E5=80=BC=E6=A1=86=E7=94=BB=E5=9C=A8?= =?UTF-8?q?=E4=B8=80=E8=B5=B7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit scripts/visualize_detections.py：在整个数据集上推理，绿=真值、蓝=命中的预测、 红=没对上真值的预测，按结果类型分目录（ok / missed / extra / none），再配一份 index.csv（逐图的真值数、预测数、命中数、最高置信度）与 summary.json（逐类 TP/FP/FN）。 --only errors 只输出有问题的图。 匹配逻辑放在 engine/detection.py：按置信度贪心、IoU 阈值、类别需一致，与 mAP 的规则 一致，所以图上看到的漏检误检和指标算的是同一回事。same_class=False 可以单看定位准不准 ——定位对但分类错，和完全没找到是两种问题。 train 与 val 分开放：训练集的图模型见过，那部分结果不能当泛化能力看。 Co-Authored-By: Claude Opus 5 --- engine/detection.py | 131 +++++++++++++++ scripts/visualize_detections.py | 257 +++++++++++++++++++++++++++++ tests/test_detection.py | 131 +++++++++++++++ tests/test_visualize_detections.py | 89 ++++++++++ 4 files changed, 608 insertions(+) create mode 100644 engine/detection.py create mode 100644 scripts/visualize_detections.py create mode 100644 tests/test_detection.py create mode 100644 tests/test_visualize_detections.py diff --git a/engine/detection.py b/engine/detection.py new file mode 100644 index 0000000..3fd01f3 --- /dev/null +++ b/engine/detection.py @@ -0,0 +1,131 @@ +"""检测结果的匹配与绘制：把预测框和真值框放在一起看。 + +匹配用标准的贪心策略——预测按置信度从高到低，each 认领 IoU 最高且未被认领的真值框。 +这与 mAP 的匹配规则一致，所以图上看到的"漏检/误检"和指标算的是同一回事。 +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import NamedTuple + +#: 真值绿、命中的预测蓝、未命中的预测红。红绿对比最刺眼，留给"错"的那一类 +COLOR_GT = (0, 220, 0) +COLOR_MATCHED = (0, 160, 255) +COLOR_UNMATCHED = (255, 40, 40) + + +class Box(NamedTuple): + """归一化的 xyxy 框。`score` 仅预测框有。""" + + class_id: int + x1: float + y1: float + x2: float + y2: float + score: float | None = None + + +def from_yolo(class_id: int, cx: float, cy: float, width: float, height: float) -> Box: + """YOLO 的中心点记法 -> xyxy。""" + return Box(class_id, cx - width / 2, cy - height / 2, cx + width / 2, cy + height / 2) + + +def parse_label(text: str) -> list[Box]: + """读一个 YOLO 标签文件的内容；空文件返回空列表（负样本）。""" + boxes = [] + for line in text.splitlines(): + parts = line.split() + if len(parts) == 5: + boxes.append(from_yolo(int(parts[0]), *(float(v) for v in parts[1:]))) + return boxes + + +def iou(a: Box, b: Box) -> float: + """两个框的交并比。不相交时为 0。""" + left, top = max(a.x1, b.x1), max(a.y1, b.y1) + right, bottom = min(a.x2, b.x2), min(a.y2, b.y2) + if right <= left or bottom <= top: + return 0.0 + overlap = (right - left) * (bottom - top) + area_a = max(a.x2 - a.x1, 0) * max(a.y2 - a.y1, 0) + area_b = max(b.x2 - b.x1, 0) * max(b.y2 - b.y1, 0) + union = area_a + area_b - overlap + return overlap / union if union > 0 else 0.0 + + +class Match(NamedTuple): + """一张图的匹配结果。下标分别指向 predictions 与 ground_truth。""" + + pairs: list[tuple[int, int]] # (预测下标, 真值下标) + unmatched_predictions: list[int] # 误检 + unmatched_ground_truth: list[int] # 漏检 + + +def match_boxes( + predictions: Sequence[Box], + ground_truth: Sequence[Box], + iou_threshold: float = 0.5, + same_class: bool = True, +) -> Match: + """贪心匹配，与 mAP 的规则一致。 + + Args: + iou_threshold: 低于它不算命中。 + same_class: 是否要求类别一致。设 False 可以单看"位置对不对"—— + 定位准但分类错，和完全没找到是两种问题，值得分开看。 + """ + order = sorted(range(len(predictions)), key=lambda i: -(predictions[i].score or 0.0)) + taken: set[int] = set() + pairs = [] + unmatched_predictions = [] + for index in order: + prediction = predictions[index] + best, best_iou = None, iou_threshold + for candidate, truth in enumerate(ground_truth): + if candidate in taken: + continue + if same_class and truth.class_id != prediction.class_id: + continue + overlap = iou(prediction, truth) + if overlap >= best_iou: + best, best_iou = candidate, overlap + if best is None: + unmatched_predictions.append(index) + else: + taken.add(best) + pairs.append((index, best)) + return Match( + sorted(pairs), + sorted(unmatched_predictions), + [i for i in range(len(ground_truth)) if i not in taken], + ) + + +def draw_boxes(image, ground_truth: Sequence[Box], predictions: Sequence[Box], + classes: Sequence[str], match: Match, width: int = 2): + """把真值与预测画在图上（就地修改传入的 PIL 图像并返回）。 + + 真值绿色、命中的预测蓝色、未命中的预测红色；预测标注类别与置信度。 + """ + from PIL import ImageDraw + + draw = ImageDraw.Draw(image) + w, h = image.size + matched = {index for index, _ in match.pairs} + + def rectangle(box: Box, color, label: str | None) -> None: + xy = [box.x1 * w, box.y1 * h, box.x2 * w, box.y2 * h] + draw.rectangle(xy, outline=color, width=width) + if label: + # 文字画在框上方，贴边时改画到框内，免得被裁掉 + y = xy[1] - 11 if xy[1] > 12 else xy[1] + 1 + draw.text((xy[0] + 1, y), label, fill=color) + + for box in ground_truth: + rectangle(box, COLOR_GT, None) + for index, box in enumerate(predictions): + color = COLOR_MATCHED if index in matched else COLOR_UNMATCHED + name = classes[box.class_id] if box.class_id < len(classes) else str(box.class_id) + rectangle(box, color, f"{name} {box.score:.2f}" if box.score is not None else name) + return image diff --git a/scripts/visualize_detections.py b/scripts/visualize_detections.py new file mode 100644 index 0000000..02e19c5 --- /dev/null +++ b/scripts/visualize_detections.py @@ -0,0 +1,257 @@ +#!/usr/bin/env python +"""用训练好的检测模型在整个数据集上推理，把预测框与真值框画在一起供人工核对。 + + uv run python scripts/visualize_detections.py \\ + --weights runs/detect/film/weights/best.pt \\ + --data data/film_yolo/data.yaml \\ + --out runs/detect/film_vis + +图上三种颜色：**绿=真值**，**蓝=命中的预测**，**红=没对上真值的预测**（误检）。 +只有绿框没有蓝框，就是漏检。匹配规则与 mAP 一致（按置信度贪心、IoU 阈值、类别需一致）， +所以图上看到的和指标算的是同一回事。 + +产出按"结果类型"分目录，直接对着看失败模式： + +``` +runs// +├── train/ val/ # 分开放：训练集的图模型见过，指标不可信 +│ ├── ok/<图名>.png # 全部命中，没有多余预测 +│ ├── missed/<图名>.png # 有漏检 +│ ├── extra/<图名>.png # 有误检（含背景图上凭空框出的） +│ └── none/<图名>.png # 既无真值也无预测——负样本判对了 +├── index.csv # 每张图一行：真值数/预测数/命中数/最高置信度 +└── summary.json # 逐类 TP/FP/FN 与各结果类型的张数 +``` + +`--only errors` 只输出有漏检或误检的图，扫起来快得多。 +""" + +from __future__ import annotations + +import argparse +import csv +import json +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 # noqa: E402 +from engine.detection import Box, draw_boxes, match_boxes, parse_label # noqa: E402 +from utils import setup_logging # noqa: E402 +from utils.device import DeviceNotAvailableError, describe_device, resolve_device # noqa: E402 + +logger = None + +OUTCOMES = ("ok", "missed", "extra", "none") + + +def outcome(n_gt: int, n_pred: int, match) -> str: + """一张图的结果类型。missed 优先于 extra：漏检比误检更该先看。""" + if n_gt == 0 and n_pred == 0: + return "none" + if match.unmatched_ground_truth: + return "missed" + if match.unmatched_predictions: + return "extra" + return "ok" + + +def load_dataset(data_yaml: Path) -> tuple[list[str], dict[str, list[Path]]]: + """从 data.yaml 读类别与各 split 的图片列表。""" + meta = yaml.safe_load(data_yaml.read_text(encoding="utf-8")) + names = meta["names"] + classes = [names[i] for i in sorted(names)] if isinstance(names, dict) else list(names) + root = Path(meta.get("path", data_yaml.parent)) + if not root.is_absolute(): + root = (data_yaml.parent / root).resolve() + + splits = {} + for split in ("train", "val"): + entry = meta.get(split) + if not entry: + continue + directory = root / entry + if not directory.is_dir(): + raise FileNotFoundError(f"{data_yaml} 里的 {split} 指向 {directory}，但它不存在") + images = sorted(p for p in directory.iterdir() if p.suffix.lower() in IMAGE_SUFFIXES) + # train 与 val 指向同一目录时（未划分的数据集）只跑一次 + if any(directory == root / meta.get(other, "") for other in splits): + continue + splits[split] = images + if not splits: + raise FileNotFoundError(f"{data_yaml} 里没有可用的 train / val") + return classes, splits + + +def label_path_for(image: Path) -> Path: + """YOLO 约定：把路径里的 images/ 换成 labels/，后缀换成 .txt。""" + parts = list(image.parts) + for index in range(len(parts) - 1, -1, -1): + if parts[index] == "images": + parts[index] = "labels" + break + return Path(*parts).with_suffix(".txt") + + +def main() -> int: + global logger + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("--weights", required=True, help="训练好的 .pt") + parser.add_argument("--data", required=True, help="数据集的 data.yaml") + parser.add_argument("--out", required=True, help="产出目录，建议放在 runs/ 下") + parser.add_argument("--device", default="auto", help="auto / cuda / mps / cpu") + parser.add_argument("--imgsz", type=int, default=640) + parser.add_argument("--conf", type=float, default=0.25, help="置信度阈值") + parser.add_argument("--iou-threshold", type=float, default=0.5, help="判定命中的 IoU") + parser.add_argument("--batch", type=int, default=16) + parser.add_argument( + "--only", choices=("all", "errors"), default="all", + help="errors=只输出有漏检或误检的图，统计仍按全量计", + ) + parser.add_argument("--limit", type=int, default=0, help="每个 split 只跑前 N 张（试参数用）") + parser.add_argument("--overwrite", action="store_true", help="产出目录已存在时先清空") + args = parser.parse_args() + + out_dir = Path(args.out) + logger = setup_logging() + + try: + device = resolve_device(args.device) + classes, splits = load_dataset(Path(args.data)) + except (DeviceNotAvailableError, FileNotFoundError) as exc: + logger.error("%s", exc) + return 1 + + if args.limit: + splits = {name: images[: args.limit] for name, images in splits.items()} + total = sum(len(images) for images in splits.values()) + logger.info("设备: %s | 权重: %s", describe_device(device), args.weights) + logger.info( + "数据集: %s（%d 类，%s）", args.data, len(classes), + "，".join(f"{name} {len(images)} 张" for name, images in splits.items()), + ) + logger.info("conf=%.2f IoU=%.2f imgsz=%d", args.conf, args.iou_threshold, args.imgsz) + + 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) + + from PIL import Image + from ultralytics import YOLO + + model = YOLO(args.weights) + device_arg = device.index or 0 if device.type == "cuda" else device.type + + outcomes: dict[str, Counter[str]] = {name: Counter() for name in splits} + stats: dict[str, dict[str, Counter[str]]] = { + name: {"tp": Counter(), "fp": Counter(), "fn": Counter()} for name in splits + } + scores: dict[str, list[float]] = defaultdict(list) + rows = [] + processed = 0 + + for split, images in splits.items(): + results = model.predict( + source=[str(p) for p in images], stream=True, imgsz=args.imgsz, conf=args.conf, + device=device_arg, batch=args.batch, verbose=False, + ) + for image_path, result in zip(images, results, strict=True): + predictions = [ + Box(int(cls), *(float(v) for v in xyxy), float(score)) + for xyxy, cls, score in zip( + result.boxes.xyxyn.tolist(), result.boxes.cls.tolist(), + result.boxes.conf.tolist(), strict=True, + ) + ] + label_file = label_path_for(image_path) + truth = parse_label(label_file.read_text(encoding="utf-8")) if label_file.exists() else [] + match = match_boxes(predictions, truth, args.iou_threshold) + + kind = outcome(len(truth), len(predictions), match) + outcomes[split][kind] += 1 + for _, truth_index in match.pairs: + stats[split]["tp"][classes[truth[truth_index].class_id]] += 1 + for index in match.unmatched_predictions: + stats[split]["fp"][classes[predictions[index].class_id]] += 1 + for index in match.unmatched_ground_truth: + stats[split]["fn"][classes[truth[index].class_id]] += 1 + for box in predictions: + scores[classes[box.class_id]].append(box.score) + + rows.append({ + "split": split, "name": image_path.name, "outcome": kind, + "n_gt": len(truth), "n_pred": len(predictions), "n_matched": len(match.pairs), + "max_conf": round(max((b.score for b in predictions), default=0.0), 4), + "gt_classes": "|".join(sorted({classes[b.class_id] for b in truth})), + "pred_classes": "|".join(sorted({classes[b.class_id] for b in predictions})), + }) + + if args.only == "all" or kind in ("missed", "extra"): + destination = out_dir / split / kind / f"{image_path.stem}.png" + destination.parent.mkdir(parents=True, exist_ok=True) + with Image.open(image_path) as raw: + canvas = raw.convert("RGB") + draw_boxes(canvas, truth, predictions, classes, match) + canvas.save(destination) + + processed += 1 + if processed % 1000 == 0: + logger.info("处理 [%d/%d]", processed, total) + + with (out_dir / "index.csv").open("w", encoding="utf-8-sig", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + + summary = { + "weights": str(Path(args.weights).resolve()), + "data": str(Path(args.data).resolve()), + "conf": args.conf, + "iou_threshold": args.iou_threshold, + "imgsz": args.imgsz, + "outcomes": {split: dict(counter) for split, counter in outcomes.items()}, + "per_class": { + split: { + name: { + "tp": counters["tp"][name], "fp": counters["fp"][name], + "fn": counters["fn"][name], + } + for name in classes + } + for split, counters in stats.items() + }, + } + (out_dir / "summary.json").write_text( + json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8" + ) + + for split, counter in outcomes.items(): + shown = sum(counter[k] for k in ("missed", "extra")) + logger.info( + "%s：全部命中 %d，漏检 %d，误检 %d，负样本判对 %d（有问题的 %d 张，占 %.1f%%）", + split, counter["ok"], counter["missed"], counter["extra"], counter["none"], + shown, 100 * shown / max(sum(counter.values()), 1), + ) + logger.info("%-12s %8s %8s %8s %8s", "类别", "TP", "FP", "FN", "召回") + for name in classes: + tp, fp, fn = (stats[split][k][name] for k in ("tp", "fp", "fn")) + if tp or fp or fn: + logger.info( + "%-12s %8d %8d %8d %7.1f%%", + name, tp, fp, fn, 100 * tp / max(tp + fn, 1), + ) + logger.info("产出目录: %s（index.csv 可按 outcome / max_conf 排序筛选）", out_dir) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_detection.py b/tests/test_detection.py new file mode 100644 index 0000000..4e1e4f1 --- /dev/null +++ b/tests/test_detection.py @@ -0,0 +1,131 @@ +"""检测结果的匹配与绘制。""" + +import pytest + +pytest.importorskip("PIL") + +from engine.detection import ( # noqa: E402 + COLOR_GT, + COLOR_MATCHED, + COLOR_UNMATCHED, + Box, + draw_boxes, + from_yolo, + iou, + match_boxes, + parse_label, +) + + +def test_from_yolo_converts_center_to_corners(): + assert from_yolo(2, 0.5, 0.5, 0.2, 0.4) == Box(2, 0.4, 0.3, 0.6, 0.7) + + +def test_parse_label_reads_boxes_and_empty_files(): + boxes = parse_label("0 0.5 0.5 0.2 0.2\n3 0.1 0.1 0.1 0.1\n") + assert [b.class_id for b in boxes] == [0, 3] + assert parse_label("") == [] # 负样本 + assert parse_label("\n\n") == [] + + +@pytest.mark.parametrize( + "a,b,expected", + [ + ((0, 0, 1, 1), (0, 0, 1, 1), 1.0), # 完全重合 + ((0, 0, 1, 1), (2, 2, 3, 3), 0.0), # 不相交 + ((0, 0, 1, 1), (1, 1, 2, 2), 0.0), # 只碰到角 + ((0, 0, 2, 2), (1, 1, 3, 3), 1 / 7), # 交 1，并 7 + ], +) +def test_iou(a, b, expected): + assert iou(Box(0, *a), Box(0, *b)) == pytest.approx(expected) + + +def test_matches_by_confidence_first(): + """两个预测抢同一个真值时，置信度高的赢——与 mAP 的规则一致。""" + truth = [from_yolo(0, 0.5, 0.5, 0.2, 0.2)] + weak = Box(0, 0.4, 0.4, 0.6, 0.6, 0.3) + strong = Box(0, 0.41, 0.41, 0.59, 0.59, 0.9) + match = match_boxes([weak, strong], truth, iou_threshold=0.3) + + assert match.pairs == [(1, 0)] # 下标 1 = strong + assert match.unmatched_predictions == [0] + assert match.unmatched_ground_truth == [] + + +def test_low_iou_counts_as_both_miss_and_false_positive(): + truth = [from_yolo(0, 0.2, 0.2, 0.1, 0.1)] + prediction = Box(0, 0.7, 0.7, 0.9, 0.9, 0.9) + match = match_boxes([prediction], truth) + assert match.pairs == [] + assert match.unmatched_predictions == [0] # 误检 + assert match.unmatched_ground_truth == [0] # 同时也是漏检 + + +def test_class_mismatch_is_not_a_hit_unless_allowed(): + """定位准但分类错：默认不算命中，same_class=False 时可单看位置。""" + truth = [from_yolo(1, 0.5, 0.5, 0.2, 0.2)] + prediction = Box(0, 0.4, 0.4, 0.6, 0.6, 0.9) + + assert match_boxes([prediction], truth).pairs == [] + assert match_boxes([prediction], truth, same_class=False).pairs == [(0, 0)] + + +def test_each_ground_truth_is_claimed_once(): + """一个真值只能被一个预测认领，多出来的预测算误检。""" + truth = [from_yolo(0, 0.5, 0.5, 0.4, 0.4)] + predictions = [ + Box(0, 0.31, 0.31, 0.69, 0.69, 0.9), + Box(0, 0.32, 0.32, 0.68, 0.68, 0.8), + ] + match = match_boxes(predictions, truth) + assert len(match.pairs) == 1 + assert len(match.unmatched_predictions) == 1 + + +def test_no_prediction_on_background_is_clean(): + match = match_boxes([], []) + assert match == ([], [], []) + + +def test_multiple_boxes_are_matched_pairwise(): + truth = [from_yolo(0, 0.2, 0.2, 0.1, 0.1), from_yolo(1, 0.8, 0.8, 0.1, 0.1)] + predictions = [ + Box(1, 0.75, 0.75, 0.85, 0.85, 0.7), + Box(0, 0.15, 0.15, 0.25, 0.25, 0.9), + ] + match = match_boxes(predictions, truth) + assert sorted(match.pairs) == [(0, 1), (1, 0)] + assert not match.unmatched_predictions and not match.unmatched_ground_truth + + +def _colors(image): + return {pixel for pixel in image.getdata()} + + +def test_draw_uses_a_distinct_colour_per_role(): + """三种角色三种颜色：绿=真值，蓝=命中的预测，红=没对上的预测。""" + from PIL import Image + + truth = [from_yolo(0, 0.25, 0.25, 0.2, 0.2), from_yolo(0, 0.75, 0.25, 0.2, 0.2)] + predictions = [ + Box(0, 0.16, 0.16, 0.34, 0.34, 0.9), # 命中第一个 + Box(0, 0.5, 0.8, 0.6, 0.9, 0.4), # 没有对应真值 + ] + match = match_boxes(predictions, truth, iou_threshold=0.5) + assert match.pairs == [(0, 0)] + + image = Image.new("RGB", (200, 200), "black") + draw_boxes(image, truth, predictions, ["blackdot"], match) + present = _colors(image) + assert COLOR_GT in present and COLOR_MATCHED in present and COLOR_UNMATCHED in present + + +def test_draw_handles_boxes_at_the_border(): + """框贴着上边缘时，标签文字要画到框内，不能画到图外去。""" + from PIL import Image + + image = Image.new("RGB", (64, 64), "black") + prediction = Box(0, 0.0, 0.0, 0.5, 0.05, 0.9) + draw_boxes(image, [], [prediction], ["blackdot"], match_boxes([prediction], [])) + assert COLOR_UNMATCHED in _colors(image) diff --git a/tests/test_visualize_detections.py b/tests/test_visualize_detections.py new file mode 100644 index 0000000..76362d1 --- /dev/null +++ b/tests/test_visualize_detections.py @@ -0,0 +1,89 @@ +"""可视化脚本里不依赖模型的部分。""" + +import sys +from pathlib import Path + +import pytest + +pytest.importorskip("yaml") +import yaml # noqa: E402 + +sys.path.insert(0, "scripts") +import visualize_detections as vd # noqa: E402 + +from engine.detection import Box, from_yolo, match_boxes # noqa: E402 + + +def _match(predictions, truth): + return match_boxes(predictions, truth) + + +@pytest.mark.parametrize( + "truth,predictions,expected", + [ + ([], [], "none"), # 负样本判对 + ([from_yolo(0, .5, .5, .2, .2)], [], "missed"), # 完全没检出 + ([], [Box(0, .4, .4, .6, .6, .9)], "extra"), # 背景上凭空框出 + ([from_yolo(0, .5, .5, .2, .2)], [Box(0, .41, .41, .59, .59, .9)], "ok"), + ], +) +def test_outcome_categories(truth, predictions, expected): + assert vd.outcome(len(truth), len(predictions), _match(predictions, truth)) == expected + + +def test_missed_takes_priority_over_extra(): + """一张图同时有漏检和误检时归为 missed——漏检更该先看。""" + truth = [from_yolo(0, .2, .2, .1, .1)] + predictions = [Box(0, .7, .7, .9, .9, .9)] + assert vd.outcome(1, 1, _match(predictions, truth)) == "missed" + + +def test_label_path_swaps_images_for_labels(): + image = Path("/data/film/images/val/blackdot_0001.png") + assert vd.label_path_for(image) == Path("/data/film/labels/val/blackdot_0001.txt") + + +def test_label_path_uses_the_last_images_segment(): + """路径里多处出现 images 时，只换最靠后的那一段。""" + image = Path("/srv/images/film/images/train/a.png") + assert vd.label_path_for(image) == Path("/srv/images/film/labels/train/a.txt") + + +@pytest.fixture +def dataset(tmp_path): + for split in ("train", "val"): + (tmp_path / "images" / split).mkdir(parents=True) + (tmp_path / "images" / split / f"{split}_0001.png").write_bytes(b"") + data_yaml = tmp_path / "data.yaml" + data_yaml.write_text( + yaml.safe_dump({ + "path": str(tmp_path), "train": "images/train", "val": "images/val", + "nc": 2, "names": {0: "blackdot", 1: "grease"}, + }), + encoding="utf-8", + ) + return data_yaml + + +def test_load_dataset(dataset): + classes, splits = vd.load_dataset(dataset) + assert classes == ["blackdot", "grease"] + assert sorted(splits) == ["train", "val"] + assert splits["val"][0].name == "val_0001.png" + + +def test_load_dataset_skips_a_split_pointing_at_the_same_directory(dataset): + """未划分的数据集 train 与 val 指向同一目录，不该把同一批图跑两遍。""" + meta = yaml.safe_load(dataset.read_text(encoding="utf-8")) + meta["val"] = meta["train"] + dataset.write_text(yaml.safe_dump(meta), encoding="utf-8") + _, splits = vd.load_dataset(dataset) + assert list(splits) == ["train"] + + +def test_load_dataset_reports_a_missing_directory(dataset): + meta = yaml.safe_load(dataset.read_text(encoding="utf-8")) + meta["val"] = "images/nope" + dataset.write_text(yaml.safe_dump(meta), encoding="utf-8") + with pytest.raises(FileNotFoundError, match="不存在"): + vd.load_dataset(dataset)