Shaun ZhouandClaude Opus 5
9b31473b27
新增 --prefix-parent：复制时把图片原来的父目录名拼到文件名前 （grease/abc.png -> grease_abc.png）。导出后目录只剩簇号，来源信息本来会丢失， 加前缀能留住它，人工核对时一眼看出某张图原本在哪个目录。 - 重名检查改为按**目标文件名**判重：加前缀后原本会撞车的文件可能就不撞了； 不加前缀而撞车时，报错会提示可以试试这个开关 - manifest.json 记下 prefix_parent，导出的 file 字段也用实际写出的文件名 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
224 lines
9.1 KiB
Python
224 lines
9.1 KiB
Python
"""按聚类结果导出数据集。"""
|
|
|
|
import csv
|
|
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 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_prefix_parent_prepends_source_folder(cluster_run, tmp_path, monkeypatch):
|
|
"""--prefix-parent 把原父目录名拼到文件名前，导出后仍能看出图来自哪个目录。"""
|
|
_, work_dir, files = cluster_run
|
|
data_root = tmp_path / "data"
|
|
|
|
assert run(["--work-dir", str(work_dir), "--data-root", str(data_root), "--prefix-parent"],
|
|
monkeypatch) == 0
|
|
|
|
out = data_root / "film_r18_cluster_k2"
|
|
exported = {p.name for p in out.rglob("*.png")}
|
|
assert exported == {f"{f.parent.name}_{f.name}" for f in files}
|
|
assert all(name.startswith(("grease_", "wrinkle_")) for name in exported)
|
|
|
|
# 仍然是真实文件、内容不变
|
|
for path in out.rglob("*.png"):
|
|
assert not path.is_symlink() and path.is_file()
|
|
|
|
manifest = json.loads((out / "manifest.json").read_text("utf-8"))
|
|
assert manifest["prefix_parent"] is True
|
|
for entries in manifest["files"].values():
|
|
for entry in entries:
|
|
source = Path(entry["source"])
|
|
assert entry["file"] == f"{source.parent.name}_{source.name}"
|
|
|
|
|
|
def test_without_prefix_names_are_unchanged(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)
|
|
|
|
exported = {p.name for p in (data_root / "film_r18_cluster_k2").rglob("*.png")}
|
|
assert exported == {f.name for f in files}
|
|
|
|
|
|
def test_prefix_parent_resolves_same_name_across_folders(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"]})
|
|
|
|
data_root = tmp_path / "data"
|
|
argv = ["--work-dir", str(work_dir), "--data-root", str(data_root)]
|
|
assert run(argv, monkeypatch) == 1 # 不加前缀：撞车
|
|
assert run([*argv, "--prefix-parent"], monkeypatch) == 0 # 加前缀：共存
|
|
|
|
out = data_root / "film_r18_cluster_k2"
|
|
assert {"grease_grease_0.png", "wrinkle_grease_0.png"} <= {p.name for p in out.rglob("*.png")}
|
|
|
|
|
|
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"
|
|
|
|
|
|
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"
|