feat(export): 改为按 clusters/ 目录导出,并重定位失效的软链

原来按 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 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
2026-09-10 11:26:02 +08:00
co-authored by Claude Opus 5
parent 9b31473b27
commit f8654f447b
3 changed files with 319 additions and 318 deletions
+19 -23
View File
@@ -262,38 +262,34 @@ scatter:
### 按聚类结果导出数据集 ### 按聚类结果导出数据集
人工核对完分组后,把结果落成一个可以直接使用的数据集目录: 聚类产出的 `clusters/` 里是**软链**,方便人工翻看和调整——可以把图在子目录间挪动、
删掉分错的、把 `cluster_03` 改名成 `grease`。整理完之后,以这个目录为准把**真实文件**
复制出来,得到一个可以直接使用的数据集:
```bash ```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<簇数>,自动推导 data/simclr_r101/ # 默认取 work_dir 的名字,--out 可完全指定
├── 0/ # 簇 0 的图,文件名保持原样 ├── grease/ # 子目录名沿用 clusters/ 里的(人工整理过的类别名)
├── 1/ ├── mosquito/
├── 2/ ├── normal/
├── 3/ └── manifest.json # 每张图的原始路径与所属类别,便于追溯
└── manifest.json # 来源 run、特征配置、每张图的原始路径
``` ```
**复制真实文件**,不是软链,所以这个目录可以直接拷走或喂给别的工具。 导出的是**真实文件**(不是软链),文件名用原图的名字、去掉 `clusters/` 里的 `0125_`
排序前缀。加 `--prefix-parent` 会再拼上原父目录名(`反面纵黑点/abc.png` ->
`反面纵黑点_abc.png`),既保住来源信息,也化解不同目录同名文件撞车。
默认保持原文件名。加 `--prefix-parent` 会把原父目录名拼到文件名前: > **从服务器拷回来的 `clusters/` 软链全都是坏的**(指向 `/home/deep/...`),不影响使用:
> 脚本读出链接目标,按其中的 `data/` 段重新落到本机 `--source-root`(默认当前目录)上;
> 万一还找不到,再用文件名在源目录里兜底索引一次。解析不到的条目会列出来并中止,
> 确认要忽略就加 `--skip-missing`。
```bash 几个保护:同类别内重名(保持原名会互相覆盖)直接报错并提示可用 `--prefix-parent`;
uv run python scripts/export_clusters.py --work-dir runs/simclr_r18 --prefix-parent 目标目录非空时必须显式 `--overwrite`。
# 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。
### 自监督预训练(SimCLR) ### 自监督预训练(SimCLR)
+150 -124
View File
@@ -1,38 +1,35 @@
#!/usr/bin/env python #!/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/ data/orgnized_from_r101_k40/
├── 0/ ← 簇 0 的原图,文件名保持不变 ├── grease/ ← 子目录名沿用 clusters/ 里的名字(人工整理过的类别名)
├── 1/ ├── mosquito/
├── 2/ ├── normal/
├── 3/ └── manifest.json # 每张图的原始路径与所属类别,便于追溯
└── manifest.json # 来源 run、每张图的原始路径,便于追溯
``` ```
**复制真实文件**,不是软链——所以这个目录可以直接拷走或喂给别的工具。 **软链失效也没关系。** 从服务器拷回来的 clusters/ 目录,软链指向的是服务器路径
默认保持原文件名;加 `--prefix-parent` 会在文件名前拼上原父目录名 (`/home/deep/ModelTrainEval/data/raw/...`),在本机全部失效。脚本会读出链接目标,
(`grease/abc.png` -> `grease_abc.png`),便于导出后仍能一眼看出图片来自哪个原始目录, 按其中的 `data/` 段重新落到本机 `--source-root` 上;万一还找不到,再用文件名在
也能避免不同子目录的同名文件在同一个簇里撞车。 源目录里兜底索引一次。
目录名由三段拼成 `<源目录名>_<骨架标签>_cluster_k<簇数>`,取自聚类产出的 导出的文件名用**原图的名字**,不带 clusters/ 里的 `0125_` 排序前缀。
summary.json:源目录名来自 assignments.csv 里各路径的公共父目录,骨架标签来自 加 `--prefix-parent` 则再拼上原父目录名(`反面纵黑点/abc.png` -> `反面纵黑点_abc.png`),
feature_source(resnet18 -> r18、tiny_conv -> tiny 等)。用 --tag / --out 可以覆盖。 便于导出后仍能看出图片来自哪个原始目录。
注意标签只反映骨架,不区分 pretrained 与 simclr——同一个骨架的两种跑法会推出同名目录。
目录非空时脚本会拒绝写入(除非 --overwrite),需要并存时用 `--tag r18ss` 这类手动区分;
不管叫什么名字,manifest.json 里都记着这批图具体来自哪次 run。
""" """
from __future__ import annotations from __future__ import annotations
import argparse import argparse
import csv
import json import json
import os
import shutil import shutil
import sys import sys
from collections import Counter, defaultdict from collections import Counter, defaultdict
@@ -44,78 +41,90 @@ from utils import setup_logging # noqa: E402
logger = None logger = None
#: 骨架 -> 目录名里的短标签 IMAGE_SUFFIXES = (".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp")
BACKBONE_TAGS = {
"resnet18": "r18",
"resnet34": "r34",
"resnet50": "r50",
"mobilenet_v3_small": "mbv3s",
"mobilenet_v3_large": "mbv3l",
"shufflenet_v2_x1_0": "sfv2",
"tiny_conv": "tiny",
}
def backbone_tag(summary: dict) -> str: class SourceResolver:
"""从 summary.json 的 feature_source 推出骨架标签。""" """把 clusters/ 里的条目解析成本机上的真实图片路径。
source = summary.get("feature_source") or {}
backbone = source.get("backbone") or {} 三种情况依次尝试:
name = backbone.get("arch") or backbone.get("name") or "unknown" 1. 条目本身就是真实文件(有人直接把图拷进去了);
return BACKBONE_TAGS.get(name, str(name).replace("_", "")) 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]]: def read_clusters(clusters_dir: Path) -> dict[str, list[Path]]:
if not path.exists(): """扫描 clusters/ 下的每个子目录,子目录名即类别名。"""
raise FileNotFoundError(f"找不到 {path},先跑一次 scripts/cluster_unlabeled.py") if not clusters_dir.is_dir():
rows: list[tuple[Path, int]] = [] raise FileNotFoundError(f"找不到目录 {clusters_dir},先跑一次 scripts/cluster_unlabeled.py")
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
groups: dict[str, list[Path]] = {}
def common_root(paths: list[Path]) -> Path: for category in sorted(p for p in clusters_dir.iterdir() if p.is_dir()):
"""所有图片的公共父目录,用来推源数据集名。""" entries = sorted(
parts = [p.resolve().parent.parts for p in paths] p
shortest = min(len(p) for p in parts) for p in category.rglob("*")
common = [] if (p.is_symlink() or p.is_file()) and p.suffix.lower() in IMAGE_SUFFIXES
for index in range(shortest): )
segment = {p[index] for p in parts} if entries:
if len(segment) != 1: groups[category.name] = entries
break if not groups:
common.append(segment.pop()) raise ValueError(f"{clusters_dir} 下没有找到任何图片条目")
return Path(*common) if common else Path(".") return groups
def destination_name(source: Path, prefix_parent: bool = False) -> str: 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 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 = [] 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) names = Counter(destination_name(f, prefix_parent) for f in files)
for name, count in names.items(): for name, count in names.items():
if count > 1: if count > 1:
sources = [ sources = [str(f) for f in files if destination_name(f, prefix_parent) == name][:3]
str(f) for f in files if destination_name(f, prefix_parent) == name problems.append(f" {category} 里 {name!r} 出现 {count} 次: {', '.join(sources)} …")
][:3]
problems.append(f" 簇 {cluster} 里 {name!r} 出现 {count} 次: {', '.join(sources)} …")
if problems: if problems:
hint = "" if prefix_parent else "可以试试 --prefix-parent 用父目录名区分,或" hint = "" if prefix_parent else "可以试试 --prefix-parent 用父目录名区分,或"
raise ValueError( raise ValueError(
@@ -128,72 +137,86 @@ def main() -> int:
parser = argparse.ArgumentParser( parser = argparse.ArgumentParser(
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
) )
parser.add_argument("--work-dir", default=None, help="聚类产出目录(含 assignments.csv)") parser.add_argument("--work-dir", default=None, help="聚类产出目录(取其中的 clusters/)")
parser.add_argument("--assignments", default=None, help="直接指定 assignments.csv") parser.add_argument("--clusters", default=None, help="直接指定 clusters 目录")
parser.add_argument("--out", default=None, help="输出目录;缺省自动推导到 data/ 下") parser.add_argument("--out", default=None, help="输出目录;缺省自动推导到 data/ 下")
parser.add_argument("--data-root", default="data", 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("--overwrite", action="store_true", help="输出目录已存在时先清空")
parser.add_argument("--dry-run", 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("--no-manifest", action="store_true", help="不写 manifest.json")
parser.add_argument( parser.add_argument(
"--prefix-parent", "--prefix-parent",
action="store_true", 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() args = parser.parse_args()
logger = setup_logging() logger = setup_logging()
if not args.work_dir and not args.assignments: if not args.work_dir and not args.clusters:
logger.error("必须给 --work-dir 或 --assignments 之一") logger.error("必须给 --work-dir 或 --clusters 之一")
return 1 return 1
work_dir = Path(args.work_dir) if args.work_dir else None 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: try:
rows = read_assignments(assignments_path) groups = read_clusters(clusters_dir)
except (FileNotFoundError, ValueError) as exc: except (FileNotFoundError, ValueError) as exc:
logger.error("%s", exc) logger.error("%s", exc)
return 1 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 = ( out_dir = (
Path(args.out) Path(args.out) if args.out else Path(args.data_root) / clusters_dir.resolve().parent.name
if args.out
else Path(args.data_root) / f"{source_root.name}_{tag}_cluster_k{num_clusters}"
) )
total = sum(len(v) for v in groups.values())
logger.info("来源: %s(%d 张,%d 个簇)", assignments_path, len(rows), len(groups)) logger.info("来源: %s(%d 张,%d 个类别)", clusters_dir, total, len(groups))
logger.info("源数据集: %s | 骨架标签: %s", source_root, tag)
logger.info("输出目录: %s", out_dir) 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()] resolver = SourceResolver(Path(args.source_root))
if missing: resolved: dict[str, list[Path]] = defaultdict(list)
logger.error("有 %d 个源文件不存在,例如 %s", len(missing), missing[0]) unresolved: list[Path] = []
return 1 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: try:
check_conflicts(groups, args.prefix_parent) check_conflicts(resolved, args.prefix_parent)
except ValueError as exc: except ValueError as exc:
logger.error("%s", exc) logger.error("%s", exc)
return 1 return 1
for cluster in sorted(groups): for category in sorted(resolved):
logger.info(" 簇 %d -> %s/%d/ 共 %d 张", cluster, out_dir.name, cluster, len(groups[cluster])) 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: if args.dry_run:
logger.info("--dry-run:未写入任何文件") logger.info("--dry-run:未写入任何文件")
return 0 return 0
@@ -208,30 +231,33 @@ def main() -> int:
copied = 0 copied = 0
total_bytes = 0 total_bytes = 0
manifest: dict[str, list[dict]] = {} manifest: dict[str, list[dict]] = {}
for cluster in sorted(groups): for category in sorted(resolved):
target = out_dir / str(cluster) target_dir = out_dir / category
target.mkdir(parents=True, exist_ok=True) target_dir.mkdir(parents=True, exist_ok=True)
entries = [] entries = []
for source in sorted(groups[cluster]): for source in sorted(resolved[category]):
name = destination_name(source, args.prefix_parent) name = destination_name(source, args.prefix_parent)
destination = target / name destination = target_dir / name
shutil.copy2(source, destination) # copy2 连同修改时间一起复制 shutil.copy2(source, destination) # copy2 连同修改时间一起复制
total_bytes += destination.stat().st_size total_bytes += destination.stat().st_size
copied += 1 copied += 1
entries.append({"file": name, "source": str(source)}) entries.append({"file": name, "source": str(source)})
manifest[str(cluster)] = entries manifest[category] = entries
logger.info("簇 %d 完成:%d 张", cluster, len(entries)) logger.info("%-14s 完成:%d 张", category, len(entries))
if not args.no_manifest: 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( (out_dir / "manifest.json").write_text(
json.dumps( json.dumps(
{ {
"source_root": str(source_root), "clusters_dir": str(clusters_dir.resolve()),
"assignments": str(assignments_path.resolve()), "source_root": str(Path(args.source_root).resolve()),
"num_clusters": num_clusters, "categories": sorted(manifest),
"prefix_parent": args.prefix_parent, "prefix_parent": args.prefix_parent,
"skipped": [str(p) for p in unresolved],
"feature_source": summary.get("feature_source"), "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, "files": manifest,
}, },
ensure_ascii=False, ensure_ascii=False,
+150 -171
View File
@@ -1,6 +1,5 @@
"""按聚类结果导出数据集。""" """按 clusters/ 目录导出数据集。"""
import csv
import filecmp import filecmp
import json import json
import sys import sys
@@ -18,38 +17,37 @@ import export_clusters as ec # noqa: E402
@pytest.fixture @pytest.fixture
def cluster_run(tmp_path): def curated(tmp_path):
"""造一份源图片 + assignments.csv + summary.json。""" """造出「源图在本机 + clusters/ 里全是失效软链」的现场。
source = tmp_path / "film"
(source / "grease").mkdir(parents=True)
(source / "wrinkle").mkdir(parents=True)
模拟从服务器拷回来的 clusters/:软链指向 /server/proj/data/raw/...,本机没有该路径,
但同样的相对路径在 tmp_path/data/raw/ 下存在。
"""
rng = np.random.default_rng(0) rng = np.random.default_rng(0)
files = [] sources = {}
for folder, count in (("grease", 3), ("wrinkle", 2)): 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): for i in range(count):
path = source / folder / f"{folder}_{i}.png" name = f"{dataset}{folder}{i}.png"
Image.fromarray(rng.integers(0, 255, (20, 30, 3), dtype=np.uint8)).save(path) path = directory / name
files.append(path) 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" clusters = tmp_path / "run" / "clusters"
work_dir.mkdir() categories = {"grease": [], "normal": []}
with (work_dir / "assignments.csv").open("w", newline="", encoding="utf-8") as handle: for index, (local, server) in enumerate(sorted(sources.items())):
writer = csv.writer(handle) category = "grease" if index % 2 == 0 else "normal"
writer.writerow(["path", "cluster", "distance_to_centroid", "source_dir", "width", "height"]) (clusters / category).mkdir(parents=True, exist_ok=True)
for index, path in enumerate(files): link = clusters / category / f"{index:04d}_{local.name}"
writer.writerow([str(path), index % 2, "0.1", path.parent.name, 30, 20]) link.symlink_to(server) # 指向服务器路径,本机必然失效
(work_dir / "summary.json").write_text( assert not link.exists()
json.dumps({ categories[category].append(local)
"num_clusters": 2, (tmp_path / "run" / "summary.json").write_text(
"feature_source": { json.dumps({"feature_source": {"backbone": {"arch": "resnet101"}}}), encoding="utf-8"
"method": "simclr",
"backbone": {"name": "torchvision", "arch": "resnet18"},
},
}),
encoding="utf-8",
) )
return source, work_dir, files return tmp_path, clusters, categories
def run(argv, monkeypatch): def run(argv, monkeypatch):
@@ -57,167 +55,148 @@ def run(argv, monkeypatch):
return ec.main() return ec.main()
def test_exports_real_files_with_original_names(cluster_run, tmp_path, monkeypatch): def test_broken_symlinks_are_relocated_to_local_source(curated, tmp_path, monkeypatch):
source, work_dir, files = cluster_run """从服务器拷回来的软链全部失效,仍应按 data/ 段落到本机原图上。"""
data_root = tmp_path / "data" root, clusters, categories = curated
out = tmp_path / "out"
assert run(["--work-dir", str(work_dir), "--data-root", str(data_root)], monkeypatch) == 0 assert run(["--clusters", str(clusters), "--out", str(out), "--source-root", str(root)],
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 monkeypatch) == 0
out = data_root / "film_r18_cluster_k2" for category, sources in categories.items():
exported = {p.name for p in out.rglob("*.png")} exported = sorted((out / category).glob("*.png"))
assert exported == {f"{f.parent.name}_{f.name}" for f in files} assert len(exported) == len(sources)
assert all(name.startswith(("grease_", "wrinkle_")) for name in exported) 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"): def test_category_names_are_preserved(curated, tmp_path, monkeypatch):
assert not path.is_symlink() and path.is_file() """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")) 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["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 entries in manifest["files"].values():
for entry in entries: for entry in entries:
source = Path(entry["source"]) assert Path(entry["source"]).is_file()
assert entry["file"] == f"{source.parent.name}_{source.name}"
def test_without_prefix_names_are_unchanged(cluster_run, tmp_path, monkeypatch): def test_dry_run_and_overwrite_guard(curated, tmp_path, monkeypatch):
_, work_dir, files = cluster_run root, clusters, _ = curated
data_root = tmp_path / "data" out = tmp_path / "out"
run(["--work-dir", str(work_dir), "--data-root", str(data_root)], monkeypatch) 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 run([*argv, "--dry-run"], monkeypatch) == 0
assert exported == {f.name for f in files} 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): def test_default_out_dir_uses_work_dir_name(curated, tmp_path, monkeypatch):
"""不同子目录的同名文件,不加前缀会撞车报错,加了前缀就能共存。""" root, clusters, _ = curated
source, work_dir, _ = cluster_run data_root = tmp_path / "exported"
twin = source / "wrinkle" / "grease_0.png" # 与 grease/grease_0.png 同名 assert run(["--work-dir", str(clusters.parent), "--data-root", str(data_root),
Image.fromarray(np.zeros((20, 30, 3), dtype=np.uint8)).save(twin) "--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" def test_resolver_falls_back_to_basename_index(curated, tmp_path):
argv = ["--work-dir", str(work_dir), "--data-root", str(data_root)] """链接目标里没有 data/ 段时,用文件名在源目录里兜底查找。"""
assert run(argv, monkeypatch) == 1 # 不加前缀:撞车 root, clusters, _ = curated
assert run([*argv, "--prefix-parent"], monkeypatch) == 0 # 加前缀:共存 resolver = ec.SourceResolver(root)
ec.logger = ec.setup_logging()
out = data_root / "film_r18_cluster_k2" odd = clusters / "grease" / "0100_odd.png"
assert {"grease_grease_0.png", "wrinkle_grease_0.png"} <= {p.name for p in out.rglob("*.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(): 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")) == "abc.png"
assert ec.destination_name(Path("data/raw/film/grease/abc.png"), True) == "grease_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"