diff --git a/README.md b/README.md index f7dab86..a00aafa 100644 --- a/README.md +++ b/README.md @@ -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` 则在 diff --git a/scripts/cluster_unlabeled.py b/scripts/cluster_unlabeled.py index b5bfd5e..3951d94 100644 --- a/scripts/cluster_unlabeled.py +++ b/scripts/cluster_unlabeled.py @@ -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, diff --git a/scripts/export_clusters.py b/scripts/export_clusters.py new file mode 100755 index 0000000..b38f92a --- /dev/null +++ b/scripts/export_clusters.py @@ -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()) diff --git a/tests/test_export_clusters.py b/tests/test_export_clusters.py new file mode 100644 index 0000000..59511e8 --- /dev/null +++ b/tests/test_export_clusters.py @@ -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"