ModelTrainEval/tests/test_export_clusters.py

"""按 clusters/ 目录导出数据集。"""

import filecmp
import json
import sys
from pathlib import Path

import pytest

pytest.importorskip("PIL")

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

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


@pytest.fixture
def curated(tmp_path):
    """造出「源图在本机 + clusters/ 里全是失效软链」的现场。

    模拟从服务器拷回来的 clusters/：软链指向 /server/proj/data/raw/...，本机没有该路径，
    但同样的相对路径在 tmp_path/data/raw/ 下存在。
    """
    rng = np.random.default_rng(0)
    sources = {}
    for dataset, folder, count in (("film", "grease", 3), ("film", "wrinkle", 2), ("blue", "", 2)):
        relative = Path("data/raw") / dataset / folder if folder else Path("data/raw") / dataset
        directory = tmp_path / relative
        directory.mkdir(parents=True, exist_ok=True)
        for i in range(count):
            name = f"{dataset}{folder}{i}.png"
            path = directory / name
            Image.fromarray(rng.integers(0, 255, (12, 16, 3), dtype=np.uint8)).save(path)
            sources[path] = str(Path("/server/proj") / relative / name)

    clusters = tmp_path / "run" / "clusters"
    categories = {"grease": [], "normal": []}
    for index, (local, server) in enumerate(sorted(sources.items())):
        category = "grease" if index % 2 == 0 else "normal"
        (clusters / category).mkdir(parents=True, exist_ok=True)
        link = clusters / category / f"{index:04d}_{local.name}"
        link.symlink_to(server)              # 指向服务器路径，本机必然失效
        assert not link.exists()
        categories[category].append(local)
    (tmp_path / "run" / "summary.json").write_text(
        json.dumps({"feature_source": {"backbone": {"arch": "resnet101"}}}), encoding="utf-8"
    )
    return tmp_path, clusters, categories


def run(argv, monkeypatch):
    monkeypatch.setattr(sys, "argv", ["export_clusters.py", *argv])
    return ec.main()


def test_broken_symlinks_are_relocated_to_local_source(curated, tmp_path, monkeypatch):
    """从服务器拷回来的软链全部失效，仍应按 data/ 段落到本机原图上。"""
    root, clusters, categories = curated
    out = tmp_path / "out"

    assert run(["--clusters", str(clusters), "--out", str(out), "--source-root", str(root)],
               monkeypatch) == 0

    for category, sources in categories.items():
        exported = sorted((out / category).glob("*.png"))
        assert len(exported) == len(sources)
        for path in exported:
            assert not path.is_symlink(), "必须是真实文件"
        by_name = {p.name: p for p in exported}
        for source in sources:
            assert filecmp.cmp(source, by_name[source.name], shallow=False)


def test_category_names_are_preserved(curated, tmp_path, monkeypatch):
    """clusters/ 里的子目录名是人工整理出来的类别，导出后要原样保留。"""
    root, clusters, categories = curated
    out = tmp_path / "out"
    run(["--clusters", str(clusters), "--out", str(out), "--source-root", str(root)], monkeypatch)

    assert sorted(p.name for p in out.iterdir() if p.is_dir()) == sorted(categories)


def test_rank_prefix_is_stripped(curated, tmp_path, monkeypatch):
    """clusters/ 里的文件名带 0000_ 排序前缀，导出用原图名。"""
    root, clusters, _ = curated
    out = tmp_path / "out"
    run(["--clusters", str(clusters), "--out", str(out), "--source-root", str(root)], monkeypatch)

    assert not any(p.name[:4].isdigit() and p.name[4] == "_" for p in out.rglob("*.png"))


def test_prefix_parent_resolves_cross_dataset_collisions(curated, tmp_path, monkeypatch):
    """不同数据集用同样的数字文件名（blue/7.png 与 pearlescent/7.png）时会撞车，
    加父目录前缀就能共存——这正是真实数据里 blue/uv/pearlescent 的情况。"""
    root, clusters, _ = curated
    twins = []
    for dataset in ("blue", "pearlescent"):
        directory = root / "data" / "raw" / dataset
        directory.mkdir(parents=True, exist_ok=True)
        path = directory / "7.png"
        Image.fromarray(np.full((12, 16, 3), 7, dtype=np.uint8)).save(path)
        link = clusters / "normal" / f"05{len(twins)}_7.png"
        link.symlink_to(f"/server/proj/data/raw/{dataset}/7.png")
        twins.append(path)

    out = tmp_path / "out"
    argv = ["--clusters", str(clusters), "--out", str(out), "--source-root", str(root)]

    assert run(argv, monkeypatch) == 1                        # 同名撞车
    assert run([*argv, "--prefix-parent"], monkeypatch) == 0  # 用父目录名区分

    names = {p.name for p in out.rglob("*.png")}
    assert {"blue_7.png", "pearlescent_7.png"} <= names


def test_real_files_in_clusters_are_used_directly(curated, tmp_path, monkeypatch):
    """有人直接把真实图片拷进 clusters/ 时也要能用。"""
    root, clusters, _ = curated
    extra = clusters / "normal" / "9999_real.png"
    Image.fromarray(np.zeros((10, 10, 3), dtype=np.uint8)).save(extra)

    out = tmp_path / "out"
    assert run(["--clusters", str(clusters), "--out", str(out), "--source-root", str(root),
                "--prefix-parent"], monkeypatch) == 0
    assert (out / "normal" / "normal_9999_real.png").is_file()


def test_unresolvable_entries_abort_unless_skipped(curated, tmp_path, monkeypatch):
    root, clusters, _ = curated
    orphan = clusters / "normal" / "0099_gone.png"
    orphan.symlink_to("/server/proj/data/raw/film/grease/does_not_exist.png")

    out = tmp_path / "out"
    argv = ["--clusters", str(clusters), "--out", str(out), "--source-root", str(root),
            "--prefix-parent"]
    assert run(argv, monkeypatch) == 1
    assert run([*argv, "--skip-missing"], monkeypatch) == 0

    manifest = json.loads((out / "manifest.json").read_text("utf-8"))
    assert len(manifest["skipped"]) == 1


def test_manifest_records_sources_and_categories(curated, tmp_path, monkeypatch):
    root, clusters, categories = curated
    out = tmp_path / "out"
    run(["--clusters", str(clusters), "--out", str(out), "--source-root", str(root),
         "--prefix-parent"], monkeypatch)

    manifest = json.loads((out / "manifest.json").read_text("utf-8"))
    assert sorted(manifest["categories"]) == sorted(categories)
    assert manifest["prefix_parent"] is True
    assert manifest["feature_source"]["backbone"]["arch"] == "resnet101"
    total = sum(manifest["category_sizes"].values())
    assert total == sum(len(v) for v in categories.values())
    for entries in manifest["files"].values():
        for entry in entries:
            assert Path(entry["source"]).is_file()


def test_dry_run_and_overwrite_guard(curated, tmp_path, monkeypatch):
    root, clusters, _ = curated
    out = tmp_path / "out"
    argv = ["--clusters", str(clusters), "--out", str(out), "--source-root", str(root),
            "--prefix-parent"]

    assert run([*argv, "--dry-run"], monkeypatch) == 0
    assert not out.exists()

    assert run(argv, monkeypatch) == 0
    assert run(argv, monkeypatch) == 1                       # 已存在且非空
    assert run([*argv, "--overwrite"], monkeypatch) == 0


def test_default_out_dir_uses_work_dir_name(curated, tmp_path, monkeypatch):
    root, clusters, _ = curated
    data_root = tmp_path / "exported"
    assert run(["--work-dir", str(clusters.parent), "--data-root", str(data_root),
                "--source-root", str(root), "--prefix-parent"], monkeypatch) == 0
    assert (data_root / clusters.parent.name).is_dir()


def test_resolver_falls_back_to_basename_index(curated, tmp_path):
    """链接目标里没有 data/ 段时，用文件名在源目录里兜底查找。"""
    root, clusters, _ = curated
    resolver = ec.SourceResolver(root)
    ec.logger = ec.setup_logging()

    odd = clusters / "grease" / "0100_odd.png"
    target = sorted((root / "data" / "raw" / "film" / "grease").glob("*.png"))[0]
    odd.symlink_to(f"/elsewhere/{target.name}")   # 路径里没有 data/ 段
    assert resolver.resolve(odd) == target.resolve()


def test_missing_clusters_dir_is_reported(tmp_path, monkeypatch):
    assert run(["--clusters", str(tmp_path / "nope"), "--out", str(tmp_path / "o")], monkeypatch) == 1


def test_destination_name_helper():
    assert ec.destination_name(Path("data/raw/film/grease/abc.png")) == "abc.png"
    assert ec.destination_name(Path("data/raw/film/grease/abc.png"), True) == "grease_abc.png"