ModelTrainEval/tests/test_export_clusters.py

"""按聚类结果导出数据集。"""

import csv
import filecmp
import json
import sys

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 cluster_run(tmp_path):
    """造一份源图片 + assignments.csv + summary.json。"""
    source = tmp_path / "film"
    (source / "grease").mkdir(parents=True)
    (source / "wrinkle").mkdir(parents=True)

    rng = np.random.default_rng(0)
    files = []
    for folder, count in (("grease", 3), ("wrinkle", 2)):
        for i in range(count):
            path = source / folder / f"{folder}_{i}.png"
            Image.fromarray(rng.integers(0, 255, (20, 30, 3), dtype=np.uint8)).save(path)
            files.append(path)

    work_dir = tmp_path / "run"
    work_dir.mkdir()
    with (work_dir / "assignments.csv").open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(["path", "cluster", "distance_to_centroid", "source_dir", "width", "height"])
        for index, path in enumerate(files):
            writer.writerow([str(path), index % 2, "0.1", path.parent.name, 30, 20])
    (work_dir / "summary.json").write_text(
        json.dumps({
            "num_clusters": 2,
            "feature_source": {
                "method": "simclr",
                "backbone": {"name": "torchvision", "arch": "resnet18"},
            },
        }),
        encoding="utf-8",
    )
    return source, work_dir, files


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


def test_exports_real_files_with_original_names(cluster_run, tmp_path, monkeypatch):
    source, work_dir, files = cluster_run
    data_root = tmp_path / "data"

    assert run(["--work-dir", str(work_dir), "--data-root", str(data_root)], monkeypatch) == 0

    out = data_root / "film_r18_cluster_k2"  # <源目录名>_<骨架标签>_cluster_k<簇数>
    assert out.is_dir()
    assert sorted(p.name for p in out.iterdir() if p.is_dir()) == ["0", "1"]

    exported = sorted(p for p in out.rglob("*.png"))
    assert len(exported) == len(files)
    for path in exported:
        assert not path.is_symlink(), "必须是真实文件，不是软链"
        assert path.is_file()
    # 文件名保持不变、内容逐字节一致
    by_name = {p.name: p for p in exported}
    for original in files:
        assert original.name in by_name
        assert filecmp.cmp(original, by_name[original.name], shallow=False)


def test_cluster_dirs_are_numbered_from_zero(cluster_run, tmp_path, monkeypatch):
    _, work_dir, files = cluster_run
    data_root = tmp_path / "data"
    run(["--work-dir", str(work_dir), "--data-root", str(data_root)], monkeypatch)

    out = data_root / "film_r18_cluster_k2"
    assert (out / "0").is_dir() and (out / "1").is_dir()
    assert len(list((out / "0").glob("*.png"))) + len(list((out / "1").glob("*.png"))) == len(files)


def test_manifest_records_provenance(cluster_run, tmp_path, monkeypatch):
    _, work_dir, files = cluster_run
    data_root = tmp_path / "data"
    run(["--work-dir", str(work_dir), "--data-root", str(data_root)], monkeypatch)

    manifest = json.loads((data_root / "film_r18_cluster_k2" / "manifest.json").read_text("utf-8"))
    assert manifest["num_clusters"] == 2
    assert manifest["feature_source"]["backbone"]["arch"] == "resnet18"
    assert sum(manifest["cluster_sizes"].values()) == len(files)
    # 原图路径可追溯
    recorded = {e["source"] for entries in manifest["files"].values() for e in entries}
    assert recorded == {str(f) for f in files}


def test_tag_and_out_overrides(cluster_run, tmp_path, monkeypatch):
    _, work_dir, _ = cluster_run
    data_root = tmp_path / "data"

    run(["--work-dir", str(work_dir), "--data-root", str(data_root), "--tag", "r18ss"], monkeypatch)
    assert (data_root / "film_r18ss_cluster_k2").is_dir()

    explicit = tmp_path / "somewhere" / "my_dataset"
    run(["--work-dir", str(work_dir), "--out", str(explicit)], monkeypatch)
    assert explicit.is_dir()


def test_dry_run_writes_nothing(cluster_run, tmp_path, monkeypatch):
    _, work_dir, _ = cluster_run
    data_root = tmp_path / "data"
    assert run(["--work-dir", str(work_dir), "--data-root", str(data_root), "--dry-run"],
               monkeypatch) == 0
    assert not data_root.exists()


def test_refuses_to_overwrite_unless_asked(cluster_run, tmp_path, monkeypatch):
    _, work_dir, _ = cluster_run
    data_root = tmp_path / "data"
    argv = ["--work-dir", str(work_dir), "--data-root", str(data_root)]

    assert run(argv, monkeypatch) == 0
    assert run(argv, monkeypatch) == 1                      # 已存在且非空 -> 拒绝
    assert run([*argv, "--overwrite"], monkeypatch) == 0     # 明确要求才覆盖


def test_duplicate_basenames_in_one_cluster_are_rejected(cluster_run, tmp_path, monkeypatch):
    """要求保留原文件名，同簇重名就会互相覆盖，必须报错而不是悄悄丢文件。"""
    source, work_dir, _ = cluster_run
    twin = source / "wrinkle" / "grease_0.png"   # 与 grease/grease_0.png 同名
    Image.fromarray(np.zeros((20, 30, 3), dtype=np.uint8)).save(twin)

    rows = list(csv.DictReader((work_dir / "assignments.csv").open(encoding="utf-8")))
    with (work_dir / "assignments.csv").open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)
        writer.writerow({**rows[0], "path": str(twin), "cluster": rows[0]["cluster"]})

    assert run(["--work-dir", str(work_dir), "--data-root", str(tmp_path / "data")], monkeypatch) == 1


def test_missing_source_file_is_reported(cluster_run, tmp_path, monkeypatch):
    _, work_dir, files = cluster_run
    files[0].unlink()
    assert run(["--work-dir", str(work_dir), "--data-root", str(tmp_path / "data")], monkeypatch) == 1


def test_backbone_tag_mapping():
    assert ec.backbone_tag({"feature_source": {"backbone": {"arch": "resnet18"}}}) == "r18"
    assert ec.backbone_tag({"feature_source": {"backbone": {"name": "tiny_conv"}}}) == "tiny"
    assert ec.backbone_tag({"feature_source": {"backbone": {"arch": "shufflenet_v2_x1_0"}}}) == "sfv2"
    assert ec.backbone_tag({}) == "unknown"