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
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)
+150 -124
View File
@@ -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,
+150 -171
View File
@@ -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"