Shaun ZhouandClaude Opus 5
33d28dcc7c
软链接记的是绝对路径，从服务器拷到本机后全部失效，也不便于直接在 runs/ 下整理。 硬链接不占额外磁盘却表现得和真实文件一样：看图软件能开、可以随意挪动、建子目录 继续分子类，整个目录拷到别的机器也不会断链。 - 新增 link_mode: hard（默认）/ sym / copy；跨文件系统时 hard 自动退回软链并警告 - 落盘后打印各方式的条目数，便于确认实际用的是哪种 - export_clusters.py 适配硬链接：条目本身就是原图，但路径在 runs/ 下、名字还带 0125_ 排序前缀。按 inode 反查 assignments.csv 还原原图路径，这样导出的文件名与 --prefix-parent 才能拿到真正的来源目录；查不到则剥掉前缀后用文件名索引兜底 实测：14,579 张跑聚类后 clusters/ 全为硬链接（nlink=2、可直接打开）； 把簇改名、建子目录挪图后再导出，排序前缀已剥除、来源目录正确还原为「反面纵黑点」， 120 张内容与原图逐字节一致。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
237 lines
10 KiB
Python
237 lines
10 KiB
Python
"""按 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_hard_links_are_traced_back_to_the_original(curated, tmp_path, monkeypatch):
|
|
"""硬链接条目本身就是原图，但路径在 runs/ 下、名字还带排序前缀。
|
|
应当按 inode 反查 assignments.csv 还原出原图路径，
|
|
这样文件名和 --prefix-parent 才能拿到真正的来源目录。"""
|
|
import csv as csv_mod
|
|
import os as os_mod
|
|
|
|
root, clusters, categories = curated
|
|
# 新建一张原图（不能复用已在该类别里的，否则和已有条目重名）
|
|
original = root / "data" / "raw" / "film" / "grease" / "extra.png"
|
|
Image.fromarray(np.full((12, 16, 3), 3, dtype=np.uint8)).save(original)
|
|
|
|
# 造一个硬链接条目 + 对应的 assignments.csv
|
|
hard = clusters / "grease" / f"0500_{original.name}"
|
|
os_mod.link(original, hard)
|
|
with (clusters.parent / "assignments.csv").open("w", newline="", encoding="utf-8") as handle:
|
|
writer = csv_mod.writer(handle)
|
|
writer.writerow(["path", "cluster"])
|
|
for sources in list(categories.values()) + [[original]]:
|
|
for source in sources:
|
|
writer.writerow([str(source), 0])
|
|
|
|
out = tmp_path / "out"
|
|
assert run(["--work-dir", str(clusters.parent), "--out", str(out),
|
|
"--source-root", str(root), "--prefix-parent"], monkeypatch) == 0
|
|
|
|
manifest = json.loads((out / "manifest.json").read_text("utf-8"))
|
|
entries = {e["file"]: e["source"] for e in manifest["files"]["grease"]}
|
|
expected = f"{original.parent.name}_{original.name}"
|
|
assert expected in entries, f"应还原成 {expected}，实际有 {sorted(entries)}"
|
|
assert entries[expected] == str(original.resolve()) # 指向原图而不是 runs/ 下的条目
|
|
assert not any(name[:4].isdigit() and name[4] == "_" for name in entries)
|
|
|
|
|
|
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"
|