feat(scripts): 按聚类结果导出数据集
新增 scripts/export_clusters.py：读聚类产出的 assignments.csv，把原图复制成 data/<源目录名>_<骨架标签>_cluster_k<簇数>/，子目录用从 0 开始的数字编号。 - 复制真实文件（copy2 保留修改时间），不建软链、不改名 - 目录名自动推导：源目录名取各路径的公共父目录，骨架标签取自 summary.json 的 feature_source；--tag / --out 可覆盖 - manifest.json 记录来源 run、特征配置与每张图的原始路径，便于追溯 - 保护：源文件缺失、同簇内重名（保持原名会互相覆盖）直接报错； 目标目录非空需显式 --overwrite；--dry-run 只打印计划 - 聚类产出的 summary.json 新增 feature_source 字段（method / backbone / feature_stages），供上面的自动命名与日后追溯 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
7f13ec5d18
commit
f8a52c9382
@@ -247,6 +247,30 @@ scatter:
|
||||
> 完整的数学推导、算法流程图与全部消融数据见
|
||||
> **[docs/unsupervised-clustering.md](docs/unsupervised-clustering.md)**。
|
||||
|
||||
### 按聚类结果导出数据集
|
||||
|
||||
人工核对完分组后，把结果落成一个可以直接使用的数据集目录：
|
||||
|
||||
```bash
|
||||
uv run python scripts/export_clusters.py --work-dir runs/simclr_r18 # 先看计划加 --dry-run
|
||||
```
|
||||
|
||||
```
|
||||
data/film_r18_cluster_k4/ # <源目录名>_<骨架标签>_cluster_k<簇数>，自动推导
|
||||
├── 0/ # 簇 0 的图，文件名保持原样
|
||||
├── 1/
|
||||
├── 2/
|
||||
├── 3/
|
||||
└── manifest.json # 来源 run、特征配置、每张图的原始路径
|
||||
```
|
||||
|
||||
**复制真实文件**，不是软链也不改名，所以这个目录可以直接拷走或喂给别的工具。
|
||||
目录名可用 `--tag r18ss` 换标签、`--out` 完全指定。几个保护：源文件缺失、
|
||||
同簇内出现重名（保持原名会互相覆盖）都会直接报错；目标目录非空时必须显式 `--overwrite`。
|
||||
|
||||
> 标签只反映骨架，不区分 pretrained 与 simclr——同一骨架的两种跑法会推出同名目录，
|
||||
> 需要并存时用 `--tag` 手动区分。无论目录叫什么，`manifest.json` 都记着这批图来自哪次 run。
|
||||
|
||||
### 自监督预训练（SimCLR）
|
||||
|
||||
`method: pretrained`（默认）用 ImageNet 权重直接抽特征，不训练；`method: simclr` 则在
|
||||
|
||||
@@ -122,10 +122,26 @@ def write_outputs(cfg, work_dir, metas, result, num_clusters, stability):
|
||||
width, height = meta["size"]
|
||||
writer.writerow([meta["path"], label, f"{distance:.6f}", meta["source_dir"], width, height])
|
||||
|
||||
method = str(cfg.get("method", "pretrained"))
|
||||
backbone_cfg = (
|
||||
dict(cfg.get("simclr", {}).get("backbone", {}))
|
||||
if method == "simclr"
|
||||
else dict(cfg.get("backbone", {}))
|
||||
)
|
||||
summary = {
|
||||
"num_clusters": num_clusters,
|
||||
"num_images": len(metas),
|
||||
"method": cfg.get("method", "pretrained"),
|
||||
"method": method,
|
||||
# 记录特征来源，供 export_clusters.py 自动命名以及日后追溯
|
||||
"feature_source": {
|
||||
"method": method,
|
||||
"backbone": backbone_cfg,
|
||||
"feature_stages": (
|
||||
cfg.get("simclr", {}).get("feature_stages")
|
||||
if method == "simclr"
|
||||
else cfg.get("feature_stages")
|
||||
),
|
||||
},
|
||||
"image_size": cfg.get("image_size"),
|
||||
"cluster_sizes": result.sizes,
|
||||
"silhouette": result.silhouette,
|
||||
|
||||
Executable
+218
@@ -0,0 +1,218 @@
|
||||
#!/usr/bin/env python
|
||||
"""按聚类结果把原图复制成一个新数据集。
|
||||
|
||||
uv run python scripts/export_clusters.py --work-dir runs/simclr_r18
|
||||
|
||||
产出（目录名默认自动推导，见下）：
|
||||
|
||||
```
|
||||
data/film_r18_cluster_k4/
|
||||
├── 0/ ← 簇 0 的原图，文件名保持不变
|
||||
├── 1/
|
||||
├── 2/
|
||||
├── 3/
|
||||
└── manifest.json # 来源 run、每张图的原始路径，便于追溯
|
||||
```
|
||||
|
||||
**复制真实文件**，不是软链、也不改名——所以这个目录可以直接拷走或喂给别的工具。
|
||||
|
||||
目录名由三段拼成 `<源目录名>_<骨架标签>_cluster_k<簇数>`，取自聚类产出的
|
||||
summary.json：源目录名来自 assignments.csv 里各路径的公共父目录，骨架标签来自
|
||||
feature_source（resnet18 -> r18、tiny_conv -> tiny 等）。用 --tag / --out 可以覆盖。
|
||||
|
||||
注意标签只反映骨架，不区分 pretrained 与 simclr——同一个骨架的两种跑法会推出同名目录。
|
||||
目录非空时脚本会拒绝写入（除非 --overwrite），需要并存时用 `--tag r18ss` 这类手动区分；
|
||||
不管叫什么名字，manifest.json 里都记着这批图具体来自哪次 run。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import json
|
||||
import shutil
|
||||
import sys
|
||||
from collections import Counter, defaultdict
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
|
||||
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",
|
||||
}
|
||||
|
||||
|
||||
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("_", ""))
|
||||
|
||||
|
||||
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 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(".")
|
||||
|
||||
|
||||
def check_conflicts(groups: dict[int, list[Path]]) -> None:
|
||||
"""同一个簇里出现同名文件时直接报错——要求保留原文件名，就不能悄悄覆盖。"""
|
||||
problems = []
|
||||
for cluster, files in sorted(groups.items()):
|
||||
names = Counter(f.name for f in files)
|
||||
for name, count in names.items():
|
||||
if count > 1:
|
||||
sources = [str(f) for f in files if f.name == name][:3]
|
||||
problems.append(f" 簇 {cluster} 里 {name!r} 出现 {count} 次: {', '.join(sources)} …")
|
||||
if problems:
|
||||
raise ValueError(
|
||||
"以下文件重名，保持原名会互相覆盖；请先重命名源文件或改用别的导出方式：\n"
|
||||
+ "\n".join(problems[:20])
|
||||
)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
global logger
|
||||
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("--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("--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")
|
||||
args = parser.parse_args()
|
||||
|
||||
logger = setup_logging()
|
||||
|
||||
if not args.work_dir and not args.assignments:
|
||||
logger.error("必须给 --work-dir 或 --assignments 之一")
|
||||
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"
|
||||
|
||||
try:
|
||||
rows = read_assignments(assignments_path)
|
||||
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}"
|
||||
)
|
||||
|
||||
logger.info("来源: %s（%d 张，%d 个簇）", assignments_path, len(rows), len(groups))
|
||||
logger.info("源数据集: %s | 骨架标签: %s", source_root, tag)
|
||||
logger.info("输出目录: %s", out_dir)
|
||||
|
||||
missing = [p for p, _ in rows if not p.exists()]
|
||||
if missing:
|
||||
logger.error("有 %d 个源文件不存在，例如 %s", len(missing), missing[0])
|
||||
return 1
|
||||
try:
|
||||
check_conflicts(groups)
|
||||
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]))
|
||||
if args.dry_run:
|
||||
logger.info("--dry-run：未写入任何文件")
|
||||
return 0
|
||||
|
||||
if out_dir.exists() and any(out_dir.iterdir()):
|
||||
if not args.overwrite:
|
||||
logger.error("%s 已存在且非空；确认要覆盖请加 --overwrite", out_dir)
|
||||
return 1
|
||||
shutil.rmtree(out_dir)
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
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)
|
||||
entries = []
|
||||
for source in sorted(groups[cluster]):
|
||||
destination = target / source.name # 保持原文件名
|
||||
shutil.copy2(source, destination) # copy2 连同修改时间一起复制
|
||||
total_bytes += destination.stat().st_size
|
||||
copied += 1
|
||||
entries.append({"file": source.name, "source": str(source)})
|
||||
manifest[str(cluster)] = entries
|
||||
logger.info("簇 %d 完成：%d 张", cluster, len(entries))
|
||||
|
||||
if not args.no_manifest:
|
||||
(out_dir / "manifest.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"source_root": str(source_root),
|
||||
"assignments": str(assignments_path.resolve()),
|
||||
"num_clusters": num_clusters,
|
||||
"feature_source": summary.get("feature_source"),
|
||||
"cluster_sizes": {k: len(v) for k, v in sorted(manifest.items())},
|
||||
"files": manifest,
|
||||
},
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
logger.info("完成：复制 %d 张，共 %.1f MB -> %s", copied, total_bytes / 1024**2, out_dir)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,161 @@
|
||||
"""按聚类结果导出数据集。"""
|
||||
|
||||
import csv
|
||||
import filecmp
|
||||
import json
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("PIL")
|
||||
|
||||
import numpy as np # noqa: E402
|
||||
from PIL import Image # noqa: E402
|
||||
|
||||
sys.path.insert(0, "scripts")
|
||||
import export_clusters as ec # noqa: E402
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def cluster_run(tmp_path):
|
||||
"""造一份源图片 + assignments.csv + summary.json。"""
|
||||
source = tmp_path / "film"
|
||||
(source / "grease").mkdir(parents=True)
|
||||
(source / "wrinkle").mkdir(parents=True)
|
||||
|
||||
rng = np.random.default_rng(0)
|
||||
files = []
|
||||
for folder, count in (("grease", 3), ("wrinkle", 2)):
|
||||
for i in range(count):
|
||||
path = source / folder / f"{folder}_{i}.png"
|
||||
Image.fromarray(rng.integers(0, 255, (20, 30, 3), dtype=np.uint8)).save(path)
|
||||
files.append(path)
|
||||
|
||||
work_dir = tmp_path / "run"
|
||||
work_dir.mkdir()
|
||||
with (work_dir / "assignments.csv").open("w", newline="", encoding="utf-8") as handle:
|
||||
writer = csv.writer(handle)
|
||||
writer.writerow(["path", "cluster", "distance_to_centroid", "source_dir", "width", "height"])
|
||||
for index, path in enumerate(files):
|
||||
writer.writerow([str(path), index % 2, "0.1", path.parent.name, 30, 20])
|
||||
(work_dir / "summary.json").write_text(
|
||||
json.dumps({
|
||||
"num_clusters": 2,
|
||||
"feature_source": {
|
||||
"method": "simclr",
|
||||
"backbone": {"name": "torchvision", "arch": "resnet18"},
|
||||
},
|
||||
}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
return source, work_dir, files
|
||||
|
||||
|
||||
def run(argv, monkeypatch):
|
||||
monkeypatch.setattr(sys, "argv", ["export_clusters.py", *argv])
|
||||
return ec.main()
|
||||
|
||||
|
||||
def test_exports_real_files_with_original_names(cluster_run, tmp_path, monkeypatch):
|
||||
source, work_dir, files = cluster_run
|
||||
data_root = tmp_path / "data"
|
||||
|
||||
assert run(["--work-dir", str(work_dir), "--data-root", str(data_root)], monkeypatch) == 0
|
||||
|
||||
out = data_root / "film_r18_cluster_k2" # <源目录名>_<骨架标签>_cluster_k<簇数>
|
||||
assert out.is_dir()
|
||||
assert sorted(p.name for p in out.iterdir() if p.is_dir()) == ["0", "1"]
|
||||
|
||||
exported = sorted(p for p in out.rglob("*.png"))
|
||||
assert len(exported) == len(files)
|
||||
for path in exported:
|
||||
assert not path.is_symlink(), "必须是真实文件，不是软链"
|
||||
assert path.is_file()
|
||||
# 文件名保持不变、内容逐字节一致
|
||||
by_name = {p.name: p for p in exported}
|
||||
for original in files:
|
||||
assert original.name in by_name
|
||||
assert filecmp.cmp(original, by_name[original.name], shallow=False)
|
||||
|
||||
|
||||
def test_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"
|
||||
Reference in New Issue
Block a user