ModelTrainEval/tests/test_prepare_yolo_dataset.py

"""把 generate_boxes.py 的产出整理成 YOLO 训练集。"""

import json
import sys
from pathlib import Path

import pytest

pytest.importorskip("PIL")

import numpy as np  # noqa: E402
import yaml  # noqa: E402
from PIL import Image  # noqa: E402

sys.path.insert(0, "scripts")
import prepare_yolo_dataset as pyd  # noqa: E402

CLASSES = ["grease", "normal", "whitedot"]


@pytest.fixture
def source(tmp_path):
    """模拟 generate_boxes.py 的产出：images/<类别>/ + labels/<类别>/ + data.yaml。"""
    root = tmp_path / "runs" / "boxes"
    rng = np.random.default_rng(0)
    # whitedot 只有 2 张（小类），normal 里有 1 张无框（负样本）
    for class_id, (name, count) in enumerate(zip(CLASSES, (6, 5, 2), strict=True)):
        image_dir = root / "images" / name
        label_dir = root / "labels" / name
        image_dir.mkdir(parents=True)
        label_dir.mkdir(parents=True)
        for i in range(count):
            stem = f"{name}_{i}"
            Image.fromarray(rng.integers(0, 255, (32, 32, 3), dtype=np.uint8)).save(
                image_dir / f"{stem}.png"
            )
            empty = name == "normal" and i == 0
            (label_dir / f"{stem}.txt").write_text(
                "" if empty else f"{class_id} 0.5 0.5 0.2 0.2\n", encoding="utf-8"
            )
    (root / "data.yaml").write_text(
        yaml.safe_dump({"nc": len(CLASSES), "names": dict(enumerate(CLASSES))}), encoding="utf-8"
    )
    return root


def run(source, out, *extra):
    argv = ["prepare_yolo_dataset.py", "--source", str(source), "--out", str(out), *extra]
    original, sys.argv = sys.argv, argv
    try:
        return pyd.main()
    finally:
        sys.argv = original


def test_flattens_classes_and_splits(source, tmp_path):
    """类别目录被拍平，train/val 分开，图片与标签一一对应。"""
    out = tmp_path / "data" / "boxes"
    assert run(source, out) == 0

    for split in ("train", "val"):
        images = sorted(p.stem for p in (out / "images" / split).iterdir())
        labels = sorted(p.stem for p in (out / "labels" / split).iterdir())
        assert images == labels
    total = sum(len(list((out / "images" / s).iterdir())) for s in ("train", "val"))
    assert total == 13
    # 拍平后不应再有 <类别>/ 这一层
    assert not (out / "images" / "train" / "grease").exists()


def test_hard_links_share_inode(source, tmp_path):
    """用硬链接而不是复制，不占额外磁盘。"""
    out = tmp_path / "data" / "boxes"
    assert run(source, out) == 0
    linked = next((out / "images" / "train").iterdir())
    original = next(p for p in (source / "images").rglob("*.png") if p.name == linked.name)
    assert linked.stat().st_ino == original.stat().st_ino


def test_small_class_present_in_both_splits(source, tmp_path):
    """只有 2 张的小类也要在验证集里留下一张，否则它的指标是空的。"""
    out = tmp_path / "data" / "boxes"
    assert run(source, out) == 0
    per_class = json.loads((out / "summary.json").read_text(encoding="utf-8"))["per_class"]
    assert per_class["train"]["whitedot"] == 1
    assert per_class["val"]["whitedot"] == 1


def test_empty_labels_kept_as_negatives_by_default(source, tmp_path):
    """无框的图默认保留——YOLO 里它们是压低误检的负样本。"""
    out = tmp_path / "data" / "boxes"
    assert run(source, out) == 0
    summary = json.loads((out / "summary.json").read_text(encoding="utf-8"))
    assert summary["empty_labels"] == 1
    assert summary["num_images"] == 13

    out_dropped = tmp_path / "data" / "dropped"
    assert run(source, out_dropped, "--drop-empty") == 0
    dropped = json.loads((out_dropped / "summary.json").read_text(encoding="utf-8"))
    assert dropped["num_images"] == 12


def test_data_yaml_uses_absolute_path_and_source_names(source, tmp_path):
    """data.yaml 要能被 ultralytics 从任意工作目录读到，类别顺序沿用来源。"""
    out = tmp_path / "data" / "boxes"
    assert run(source, out) == 0
    meta = yaml.safe_load((out / "data.yaml").read_text(encoding="utf-8"))
    assert Path(meta["path"]).is_absolute()
    assert meta["nc"] == 3
    assert [meta["names"][i] for i in sorted(meta["names"])] == CLASSES
    assert (Path(meta["path"]) / meta["train"]).is_dir()
    assert (Path(meta["path"]) / meta["val"]).is_dir()


def test_reproducible_split(source, tmp_path):
    """同一个种子两次运行，划分必须一致。"""
    first, second = tmp_path / "a", tmp_path / "b"
    assert run(source, first) == 0
    assert run(source, second) == 0
    for split in ("train", "val"):
        assert sorted(p.name for p in (first / "images" / split).iterdir()) == sorted(
            p.name for p in (second / "images" / split).iterdir()
        )


def test_name_collision_after_flattening_is_reported(source, tmp_path, capsys):
    """不同类别目录下的同名文件，拍平后会互相覆盖——必须报错而不是静默丢数据。"""
    duplicate = source / "images" / "normal" / "grease_0.png"
    Image.fromarray(np.zeros((32, 32, 3), dtype=np.uint8)).save(duplicate)
    (source / "labels" / "normal" / "grease_0.txt").write_text("1 0.5 0.5 0.2 0.2\n", encoding="utf-8")

    assert run(source, tmp_path / "data" / "boxes") != 0
    assert "冲突" in capsys.readouterr().out


def test_refuses_to_clobber_existing_output(source, tmp_path):
    """输出目录非空时要求显式 --overwrite。"""
    out = tmp_path / "data" / "boxes"
    out.mkdir(parents=True)
    (out / "keep.txt").write_text("x", encoding="utf-8")
    assert run(source, out) != 0
    assert (out / "keep.txt").exists()
    assert run(source, out, "--overwrite") == 0
    assert not (out / "keep.txt").exists()


def test_missing_label_is_an_error(source, tmp_path):
    """图片多于标签说明来源目录不完整，不能悄悄跳过。"""
    (source / "labels" / "grease" / "grease_0.txt").unlink()
    assert run(source, tmp_path / "out") != 0


def test_dry_run_writes_nothing(source, tmp_path):
    out = tmp_path / "data" / "boxes"
    assert run(source, out, "--dry-run") == 0
    assert not out.exists()