From f8654f447b4641a4c2cb40674fdeed100c841a6b Mon Sep 17 00:00:00 2001 From: Shaun Zhou Date: Thu, 10 Sep 2026 11:26:02 +0800 Subject: [PATCH] =?UTF-8?q?feat(export):=20=E6=94=B9=E4=B8=BA=E6=8C=89=20c?= =?UTF-8?q?lusters/=20=E7=9B=AE=E5=BD=95=E5=AF=BC=E5=87=BA=EF=BC=8C?= =?UTF-8?q?=E5=B9=B6=E9=87=8D=E5=AE=9A=E4=BD=8D=E5=A4=B1=E6=95=88=E7=9A=84?= =?UTF-8?q?=E8=BD=AF=E9=93=BE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 原来按 assignments.csv 导出，反映的是聚类当时的划分；而 clusters/ 目录会被人工 整理——挪动分错的图、把 cluster_03 改名成 grease。现在以整理后的目录为准。 - 扫描 clusters/ 下的子目录，目录名即类别名并原样保留（不再是数字编号） - 文件名用原图的名字，去掉 clusters/ 里的 NNNN_ 排序前缀 - 失效软链的重定位：读出链接目标，按其中的 data/ 段落到本机 --source-root 上； 再不行用文件名在源目录里兜底索引（惰性构建，用不到就不扫盘）。 从服务器拷回来的 clusters/ 软链全部指向 /home/deep/... 在本机失效，这是常态 - 解析不到的条目列出并中止，--skip-missing 可忽略 - manifest 记录类别、每张图的原始路径与被跳过的条目 在真实目录上验证：11,212 张、15 个类别，软链 100% 失效但全部重定位成功， 导出后抽查 200 张与原图逐字节一致。跨数据集的数字文件名（blue/113.bmp 与 pearlescent/113.bmp）会被重名检查拦下，--prefix-parent 可化解。 Co-Authored-By: Claude Opus 5 --- README.md | 42 ++--- scripts/export_clusters.py | 274 ++++++++++++++++------------- tests/test_export_clusters.py | 321 ++++++++++++++++------------------ 3 files changed, 319 insertions(+), 318 deletions(-) diff --git a/README.md b/README.md index 06f415d..6f9efd9 100644 --- a/README.md +++ b/README.md @@ -262,38 +262,34 @@ scatter: ### 按聚类结果导出数据集 -人工核对完分组后，把结果落成一个可以直接使用的数据集目录： +聚类产出的 `clusters/` 里是**软链**，方便人工翻看和调整——可以把图在子目录间挪动、 +删掉分错的、把 `cluster_03` 改名成 `grease`。整理完之后，以这个目录为准把**真实文件** +复制出来，得到一个可以直接使用的数据集： ```bash -uv run python scripts/export_clusters.py --work-dir runs/simclr_r18 # 先看计划加 --dry-run +uv run python scripts/export_clusters.py --work-dir runs/simclr_r101 --prefix-parent +# 先看计划加 --dry-run ``` ``` -data/film_r18_cluster_k4/ # <源目录名>_<骨架标签>_cluster_k<簇数>，自动推导 -├── 0/ # 簇 0 的图，文件名保持原样 -├── 1/ -├── 2/ -├── 3/ -└── manifest.json # 来源 run、特征配置、每张图的原始路径 +data/simclr_r101/ # 默认取 work_dir 的名字，--out 可完全指定 +├── grease/ # 子目录名沿用 clusters/ 里的（人工整理过的类别名） +├── mosquito/ +├── normal/ +└── manifest.json # 每张图的原始路径与所属类别，便于追溯 ``` -**复制真实文件**，不是软链，所以这个目录可以直接拷走或喂给别的工具。 +导出的是**真实文件**（不是软链），文件名用原图的名字、去掉 `clusters/` 里的 `0125_` +排序前缀。加 `--prefix-parent` 会再拼上原父目录名（`反面纵黑点/abc.png` -> +`反面纵黑点_abc.png`），既保住来源信息，也化解不同目录同名文件撞车。 -默认保持原文件名。加 `--prefix-parent` 会把原父目录名拼到文件名前： +> **从服务器拷回来的 `clusters/` 软链全都是坏的**（指向 `/home/deep/...`），不影响使用： +> 脚本读出链接目标，按其中的 `data/` 段重新落到本机 `--source-root`（默认当前目录）上； +> 万一还找不到，再用文件名在源目录里兜底索引一次。解析不到的条目会列出来并中止， +> 确认要忽略就加 `--skip-missing`。 -```bash -uv run python scripts/export_clusters.py --work-dir runs/simclr_r18 --prefix-parent -# back-mosquito-G0/5.14_677mmX316.10m_....png -> back-mosquito-G0_5.14_677mmX316.10m_....png -``` - -导出后目录只剩簇号，原来的来源信息就没了——加前缀能留住它，人工核对时一眼看出 -某张图原本在哪个目录；顺带也化解不同子目录同名文件在同一个簇里撞车的情况 -（撞车时的报错也会提示可以试试这个开关）。 -目录名可用 `--tag r18ss` 换标签、`--out` 完全指定。几个保护：源文件缺失、 -同簇内出现重名（保持原名会互相覆盖）都会直接报错；目标目录非空时必须显式 `--overwrite`。 - -> 标签只反映骨架，不区分 pretrained 与 simclr——同一骨架的两种跑法会推出同名目录， -> 需要并存时用 `--tag` 手动区分。无论目录叫什么，`manifest.json` 都记着这批图来自哪次 run。 +几个保护：同类别内重名（保持原名会互相覆盖）直接报错并提示可用 `--prefix-parent`； +目标目录非空时必须显式 `--overwrite`。 ### 自监督预训练（SimCLR） diff --git a/scripts/export_clusters.py b/scripts/export_clusters.py index 59e2ad0..87bb54a 100755 --- a/scripts/export_clusters.py +++ b/scripts/export_clusters.py @@ -1,38 +1,35 @@ #!/usr/bin/env python -"""按聚类结果把原图复制成一个新数据集。 +"""按 clusters/ 目录里的分组把原图复制成一个新数据集。 - uv run python scripts/export_clusters.py --work-dir runs/simclr_r18 + uv run python scripts/export_clusters.py --work-dir runs/orgnized_from_r101_k40 -产出（目录名默认自动推导，见下）： +聚类脚本产出的 `clusters/` 里是**软链**，方便人工翻看、调整——可以把图在子目录之间 +挪动、删掉分错的、把 `cluster_03` 改名成 `grease`。这个脚本以整理后的目录为准， +把每张图的**真实文件**复制出来： ``` -data/film_r18_cluster_k4/ -├── 0/ ← 簇 0 的原图，文件名保持不变 -├── 1/ -├── 2/ -├── 3/ -└── manifest.json # 来源 run、每张图的原始路径，便于追溯 +data/orgnized_from_r101_k40/ +├── grease/ ← 子目录名沿用 clusters/ 里的名字（人工整理过的类别名） +├── mosquito/ +├── normal/ +└── manifest.json # 每张图的原始路径与所属类别，便于追溯 ``` -**复制真实文件**，不是软链——所以这个目录可以直接拷走或喂给别的工具。 -默认保持原文件名；加 `--prefix-parent` 会在文件名前拼上原父目录名 -（`grease/abc.png` -> `grease_abc.png`），便于导出后仍能一眼看出图片来自哪个原始目录， -也能避免不同子目录的同名文件在同一个簇里撞车。 +**软链失效也没关系。** 从服务器拷回来的 clusters/ 目录，软链指向的是服务器路径 +（`/home/deep/ModelTrainEval/data/raw/...`），在本机全部失效。脚本会读出链接目标， +按其中的 `data/` 段重新落到本机 `--source-root` 上；万一还找不到，再用文件名在 +源目录里兜底索引一次。 -目录名由三段拼成 `<源目录名>_<骨架标签>_cluster_k<簇数>`，取自聚类产出的 -summary.json：源目录名来自 assignments.csv 里各路径的公共父目录，骨架标签来自 -feature_source（resnet18 -> r18、tiny_conv -> tiny 等）。用 --tag / --out 可以覆盖。 - -注意标签只反映骨架，不区分 pretrained 与 simclr——同一个骨架的两种跑法会推出同名目录。 -目录非空时脚本会拒绝写入（除非 --overwrite），需要并存时用 `--tag r18ss` 这类手动区分； -不管叫什么名字，manifest.json 里都记着这批图具体来自哪次 run。 +导出的文件名用**原图的名字**，不带 clusters/ 里的 `0125_` 排序前缀。 +加 `--prefix-parent` 则再拼上原父目录名（`反面纵黑点/abc.png` -> `反面纵黑点_abc.png`）， +便于导出后仍能看出图片来自哪个原始目录。 """ from __future__ import annotations import argparse -import csv import json +import os import shutil import sys from collections import Counter, defaultdict @@ -44,78 +41,90 @@ from utils import setup_logging # noqa: E402 logger = None -#: 骨架 -> 目录名里的短标签 -BACKBONE_TAGS = { - "resnet18": "r18", - "resnet34": "r34", - "resnet50": "r50", - "mobilenet_v3_small": "mbv3s", - "mobilenet_v3_large": "mbv3l", - "shufflenet_v2_x1_0": "sfv2", - "tiny_conv": "tiny", -} +IMAGE_SUFFIXES = (".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp") -def backbone_tag(summary: dict) -> str: - """从 summary.json 的 feature_source 推出骨架标签。""" - source = summary.get("feature_source") or {} - backbone = source.get("backbone") or {} - name = backbone.get("arch") or backbone.get("name") or "unknown" - return BACKBONE_TAGS.get(name, str(name).replace("_", "")) +class SourceResolver: + """把 clusters/ 里的条目解析成本机上的真实图片路径。 + + 三种情况依次尝试： + 1. 条目本身就是真实文件（有人直接把图拷进去了）； + 2. 软链有效，直接 resolve； + 3. 软链失效——读出链接目标，按 `data/` 段重新落到本机 `source_root` 上。 + 再不行就用链接目标的文件名在源目录里查索引（索引惰性构建，用不到就不扫盘）。 + """ + + def __init__(self, source_root: Path, marker: str = "data") -> None: + self.source_root = source_root.resolve() + self.marker = marker + self._index: dict[str, Path] | None = None + + def _rebase(self, target: Path) -> Path | None: + """把 /home/deep/ModelTrainEval/data/raw/film/x.png 落到本机同名的 data/ 下。""" + parts = target.parts + if self.marker not in parts: + return None + candidate = self.source_root.joinpath(*parts[parts.index(self.marker) :]) + return candidate if candidate.exists() else None + + def _lookup(self, name: str) -> Path | None: + if self._index is None: + logger.info("建立文件名索引以兜底解析（扫描 %s）…", self.source_root / self.marker) + self._index = { + path.name: path + for path in (self.source_root / self.marker).rglob("*") + if path.suffix.lower() in IMAGE_SUFFIXES + } + logger.info("索引 %d 个文件", len(self._index)) + return self._index.get(name) + + def resolve(self, entry: Path) -> Path | None: + if entry.is_symlink(): + target = Path(os.readlink(entry)) + if entry.exists(): # 软链有效 + return entry.resolve() + return self._rebase(target) or self._lookup(target.name) + if entry.is_file(): + return entry.resolve() + return None -def read_assignments(path: Path) -> list[tuple[Path, int]]: - if not path.exists(): - raise FileNotFoundError(f"找不到 {path}，先跑一次 scripts/cluster_unlabeled.py") - rows: list[tuple[Path, int]] = [] - with path.open(encoding="utf-8") as handle: - for line_no, row in enumerate(csv.DictReader(handle), start=2): - try: - rows.append((Path(row["path"]), int(row["cluster"]))) - except (KeyError, ValueError) as exc: - raise ValueError(f"{path}:{line_no} 解析失败：{row}") from exc - if not rows: - raise ValueError(f"{path} 里没有任何记录") - return rows +def read_clusters(clusters_dir: Path) -> dict[str, list[Path]]: + """扫描 clusters/ 下的每个子目录，子目录名即类别名。""" + if not clusters_dir.is_dir(): + raise FileNotFoundError(f"找不到目录 {clusters_dir}，先跑一次 scripts/cluster_unlabeled.py") - -def common_root(paths: list[Path]) -> Path: - """所有图片的公共父目录，用来推源数据集名。""" - parts = [p.resolve().parent.parts for p in paths] - shortest = min(len(p) for p in parts) - common = [] - for index in range(shortest): - segment = {p[index] for p in parts} - if len(segment) != 1: - break - common.append(segment.pop()) - return Path(*common) if common else Path(".") + groups: dict[str, list[Path]] = {} + for category in sorted(p for p in clusters_dir.iterdir() if p.is_dir()): + entries = sorted( + p + for p in category.rglob("*") + if (p.is_symlink() or p.is_file()) and p.suffix.lower() in IMAGE_SUFFIXES + ) + if entries: + groups[category.name] = entries + if not groups: + raise ValueError(f"{clusters_dir} 下没有找到任何图片条目") + return groups def destination_name(source: Path, prefix_parent: bool = False) -> str: - """目标文件名。 + """目标文件名，用**原图**的名字（不带 clusters/ 里的 NNNN_ 排序前缀）。 - `prefix_parent` 时拼上原父目录名（`grease/abc.png` -> `grease_abc.png`）： - 导出后目录只剩簇号，原来的来源信息就没了，加前缀能留住它； - 顺带也化解不同子目录同名文件在同一个簇里撞车的情况。 + `prefix_parent` 时拼上原父目录名：导出后只剩类别目录，来源信息本来会丢失。 """ return f"{source.parent.name}_{source.name}" if prefix_parent else source.name -def check_conflicts(groups: dict[int, list[Path]], prefix_parent: bool = False) -> None: - """同一个簇里出现同名文件时直接报错——不改名就不能悄悄覆盖。 - - 按**目标文件名**判重：加了父目录前缀之后原本会撞车的文件可能就不撞了。 - """ +def check_conflicts(resolved: dict[str, list[Path]], prefix_parent: bool = False) -> None: + """同一个类别里出现同名文件时直接报错——不改名就不能悄悄覆盖。""" problems = [] - for cluster, files in sorted(groups.items()): + for category, files in sorted(resolved.items()): names = Counter(destination_name(f, prefix_parent) for f in files) for name, count in names.items(): if count > 1: - sources = [ - str(f) for f in files if destination_name(f, prefix_parent) == name - ][:3] - problems.append(f" 簇 {cluster} 里 {name!r} 出现 {count} 次: {', '.join(sources)} …") + sources = [str(f) for f in files if destination_name(f, prefix_parent) == name][:3] + problems.append(f" {category} 里 {name!r} 出现 {count} 次: {', '.join(sources)} …") if problems: hint = "" if prefix_parent else "可以试试 --prefix-parent 用父目录名区分，或" raise ValueError( @@ -128,72 +137,86 @@ def main() -> int: parser = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter ) - parser.add_argument("--work-dir", default=None, help="聚类产出目录（含 assignments.csv）") - parser.add_argument("--assignments", default=None, help="直接指定 assignments.csv") + parser.add_argument("--work-dir", default=None, help="聚类产出目录（取其中的 clusters/）") + parser.add_argument("--clusters", default=None, help="直接指定 clusters 目录") parser.add_argument("--out", default=None, help="输出目录；缺省自动推导到 data/ 下") parser.add_argument("--data-root", default="data", help="自动推导时的输出根目录，默认 data") - parser.add_argument("--tag", default=None, help="覆盖目录名里的骨架标签（如 r18）") + parser.add_argument( + "--source-root", + default=".", + help="本机存放原图的仓库根目录，用来重定位失效的软链，默认当前目录", + ) parser.add_argument("--overwrite", action="store_true", help="输出目录已存在时先清空") parser.add_argument("--dry-run", action="store_true", help="只打印计划，不真的复制") parser.add_argument("--no-manifest", action="store_true", help="不写 manifest.json") parser.add_argument( "--prefix-parent", action="store_true", - help="复制时在文件名前拼上原父目录名（grease/abc.png -> grease_abc.png）", + help="复制时在文件名前拼上原父目录名（反面纵黑点/abc.png -> 反面纵黑点_abc.png）", + ) + parser.add_argument( + "--skip-missing", action="store_true", help="解析不到的条目跳过而不是中止" ) args = parser.parse_args() logger = setup_logging() - if not args.work_dir and not args.assignments: - logger.error("必须给 --work-dir 或 --assignments 之一") + if not args.work_dir and not args.clusters: + logger.error("必须给 --work-dir 或 --clusters 之一") return 1 work_dir = Path(args.work_dir) if args.work_dir else None - assignments_path = Path(args.assignments) if args.assignments else work_dir / "assignments.csv" + clusters_dir = Path(args.clusters) if args.clusters else work_dir / "clusters" try: - rows = read_assignments(assignments_path) + groups = read_clusters(clusters_dir) except (FileNotFoundError, ValueError) as exc: logger.error("%s", exc) return 1 - summary_path = (work_dir or assignments_path.parent) / "summary.json" - summary = json.loads(summary_path.read_text(encoding="utf-8")) if summary_path.exists() else {} - - groups: dict[int, list[Path]] = defaultdict(list) - for path, cluster in rows: - groups[cluster].append(path) - - source_root = common_root([p for p, _ in rows]) - num_clusters = int(summary.get("num_clusters") or (max(groups) + 1)) - tag = args.tag or backbone_tag(summary) out_dir = ( - Path(args.out) - if args.out - else Path(args.data_root) / f"{source_root.name}_{tag}_cluster_k{num_clusters}" + Path(args.out) if args.out else Path(args.data_root) / clusters_dir.resolve().parent.name ) - - logger.info("来源: %s（%d 张，%d 个簇）", assignments_path, len(rows), len(groups)) - logger.info("源数据集: %s | 骨架标签: %s", source_root, tag) + total = sum(len(v) for v in groups.values()) + logger.info("来源: %s（%d 张，%d 个类别）", clusters_dir, total, len(groups)) logger.info("输出目录: %s", out_dir) - if args.prefix_parent: - example = rows[0][0] - logger.info( - "文件名加父目录前缀：%s -> %s", example.name, destination_name(example, True) - ) - missing = [p for p, _ in rows if not p.exists()] - if missing: - logger.error("有 %d 个源文件不存在，例如 %s", len(missing), missing[0]) - return 1 + resolver = SourceResolver(Path(args.source_root)) + resolved: dict[str, list[Path]] = defaultdict(list) + unresolved: list[Path] = [] + for category, entries in groups.items(): + for entry in entries: + source = resolver.resolve(entry) + if source is None: + unresolved.append(entry) + else: + resolved[category].append(source) + + if unresolved: + logger.error("有 %d 个条目在本机找不到对应的原图，例如：", len(unresolved)) + for entry in unresolved[:5]: + target = os.readlink(entry) if entry.is_symlink() else "(不是软链)" + logger.error(" %s -> %s", entry, target) + if not args.skip_missing: + logger.error( + "确认 --source-root 指向存放 data/ 的目录（当前 %s）；" + "确实要忽略这些条目请加 --skip-missing", + Path(args.source_root).resolve(), + ) + return 1 + logger.warning("--skip-missing：跳过上述 %d 个条目", len(unresolved)) + try: - check_conflicts(groups, args.prefix_parent) + check_conflicts(resolved, args.prefix_parent) except ValueError as exc: logger.error("%s", exc) return 1 - for cluster in sorted(groups): - logger.info(" 簇 %d -> %s/%d/ 共 %d 张", cluster, out_dir.name, cluster, len(groups[cluster])) + for category in sorted(resolved): + logger.info(" %-14s -> %s/%s/ 共 %d 张", category, out_dir.name, category, + len(resolved[category])) + if args.prefix_parent: + example = next(iter(next(iter(resolved.values())))) + logger.info("文件名加父目录前缀：%s -> %s", example.name, destination_name(example, True)) if args.dry_run: logger.info("--dry-run：未写入任何文件") return 0 @@ -208,30 +231,33 @@ def main() -> int: copied = 0 total_bytes = 0 manifest: dict[str, list[dict]] = {} - for cluster in sorted(groups): - target = out_dir / str(cluster) - target.mkdir(parents=True, exist_ok=True) + for category in sorted(resolved): + target_dir = out_dir / category + target_dir.mkdir(parents=True, exist_ok=True) entries = [] - for source in sorted(groups[cluster]): + for source in sorted(resolved[category]): name = destination_name(source, args.prefix_parent) - destination = target / name - shutil.copy2(source, destination) # copy2 连同修改时间一起复制 + destination = target_dir / name + shutil.copy2(source, destination) # copy2 连同修改时间一起复制 total_bytes += destination.stat().st_size copied += 1 entries.append({"file": name, "source": str(source)}) - manifest[str(cluster)] = entries - logger.info("簇 %d 完成：%d 张", cluster, len(entries)) + manifest[category] = entries + logger.info("%-14s 完成：%d 张", category, len(entries)) if not args.no_manifest: + summary_path = (work_dir or clusters_dir.parent) / "summary.json" + summary = json.loads(summary_path.read_text(encoding="utf-8")) if summary_path.exists() else {} (out_dir / "manifest.json").write_text( json.dumps( { - "source_root": str(source_root), - "assignments": str(assignments_path.resolve()), - "num_clusters": num_clusters, + "clusters_dir": str(clusters_dir.resolve()), + "source_root": str(Path(args.source_root).resolve()), + "categories": sorted(manifest), "prefix_parent": args.prefix_parent, + "skipped": [str(p) for p in unresolved], "feature_source": summary.get("feature_source"), - "cluster_sizes": {k: len(v) for k, v in sorted(manifest.items())}, + "category_sizes": {k: len(v) for k, v in sorted(manifest.items())}, "files": manifest, }, ensure_ascii=False, diff --git a/tests/test_export_clusters.py b/tests/test_export_clusters.py index e35cb9e..4192d2c 100644 --- a/tests/test_export_clusters.py +++ b/tests/test_export_clusters.py @@ -1,6 +1,5 @@ -"""按聚类结果导出数据集。""" +"""按 clusters/ 目录导出数据集。""" -import csv import filecmp import json import sys @@ -18,38 +17,37 @@ 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) +def curated(tmp_path): + """造出「源图在本机 + clusters/ 里全是失效软链」的现场。 + 模拟从服务器拷回来的 clusters/：软链指向 /server/proj/data/raw/...，本机没有该路径， + 但同样的相对路径在 tmp_path/data/raw/ 下存在。 + """ rng = np.random.default_rng(0) - files = [] - for folder, count in (("grease", 3), ("wrinkle", 2)): + 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): - path = source / folder / f"{folder}_{i}.png" - Image.fromarray(rng.integers(0, 255, (20, 30, 3), dtype=np.uint8)).save(path) - files.append(path) + 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) - 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", + 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 source, work_dir, files + return tmp_path, clusters, categories def run(argv, monkeypatch): @@ -57,167 +55,148 @@ def run(argv, monkeypatch): 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" +def test_broken_symlinks_are_relocated_to_local_source(curated, tmp_path, monkeypatch): + """从服务器拷回来的软链全部失效，仍应按 data/ 段落到本机原图上。""" + root, clusters, categories = curated + out = tmp_path / "out" - 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"], + assert run(["--clusters", str(clusters), "--out", str(out), "--source-root", str(root)], 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 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) - # 仍然是真实文件、内容不变 - for path in out.rglob("*.png"): - assert not path.is_symlink() and path.is_file() + +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: - source = Path(entry["source"]) - assert entry["file"] == f"{source.parent.name}_{source.name}" + assert Path(entry["source"]).is_file() -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) +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"] - exported = {p.name for p in (data_root / "film_r18_cluster_k2").rglob("*.png")} - assert exported == {f.name for f in files} + 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_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) +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() - 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 # 加前缀：共存 +def test_resolver_falls_back_to_basename_index(curated, tmp_path): + """链接目标里没有 data/ 段时，用文件名在源目录里兜底查找。""" + root, clusters, _ = curated + resolver = ec.SourceResolver(root) + ec.logger = ec.setup_logging() - out = data_root / "film_r18_cluster_k2" - assert {"grease_grease_0.png", "wrinkle_grease_0.png"} <= {p.name for p in out.rglob("*.png")} + 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" - - -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"