#!/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、每张图的原始路径，便于追溯 ``` **复制真实文件**，不是软链——所以这个目录可以直接拷走或喂给别的工具。 默认保持原文件名；加 `--prefix-parent` 会在文件名前拼上原父目录名 （`grease/abc.png` -> `grease_abc.png`），便于导出后仍能一眼看出图片来自哪个原始目录， 也能避免不同子目录的同名文件在同一个簇里撞车。 目录名由三段拼成 `<源目录名>_<骨架标签>_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 destination_name(source: Path, prefix_parent: bool = False) -> str: """目标文件名。 `prefix_parent` 时拼上原父目录名（`grease/abc.png` -> `grease_abc.png`）： 导出后目录只剩簇号，原来的来源信息就没了，加前缀能留住它； 顺带也化解不同子目录同名文件在同一个簇里撞车的情况。 """ 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: """同一个簇里出现同名文件时直接报错——不改名就不能悄悄覆盖。 按**目标文件名**判重：加了父目录前缀之后原本会撞车的文件可能就不撞了。 """ problems = [] for cluster, files in sorted(groups.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)} …") if problems: hint = "" if prefix_parent else "可以试试 --prefix-parent 用父目录名区分，或" raise ValueError( f"以下文件重名，同名会互相覆盖；{hint}先重命名源文件：\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") parser.add_argument( "--prefix-parent", action="store_true", help="复制时在文件名前拼上原父目录名（grease/abc.png -> grease_abc.png）", ) 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) 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 try: check_conflicts(groups, 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])) 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]): name = destination_name(source, args.prefix_parent) destination = target / 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)) 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, "prefix_parent": args.prefix_parent, "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())