#!/usr/bin/env python """按 clusters/ 目录里的分组把原图复制成一个新数据集。 uv run python scripts/export_clusters.py --work-dir runs/orgnized_from_r101_k40 聚类脚本产出的 `clusters/` 默认是**硬链接**（`link_mode: hard`），方便人工翻看、调整—— 可以把图在子目录之间挪动、删掉分错的、把 `cluster_03` 改名成 `grease`。这个脚本以整理后的 目录为准，把每张图**复制**成独立的副本： ``` data/orgnized_from_r101_k40/ ├── grease/ ← 子目录名沿用 clusters/ 里的名字（人工整理过的类别名） ├── mosquito/ ├── normal/ └── manifest.json # 每张图的原始路径与所属类别，便于追溯 ``` **硬链接条目本身就是原图**，只是路径在 runs/ 下、文件名带着排序前缀。脚本按 inode 反查 assignments.csv 认出它对应哪张原图，好让文件名和 `--prefix-parent` 拿到真正的来源。 **软链失效也没关系**（`link_mode: sym` 时）。从服务器拷回来的 clusters/ 目录，软链指向的 是服务器路径（`/home/deep/ModelTrainEval/data/raw/...`），在本机全部失效。脚本会读出链接 目标，按其中的 `data/` 段重新落到本机 `--source-root` 上；万一还找不到，再用文件名在 源目录里兜底索引一次。 导出的文件名用**原图的名字**，不带 clusters/ 里的 `0125_` 排序前缀。 加 `--prefix-parent` 则再拼上原父目录名（`反面纵黑点/abc.png` -> `反面纵黑点_abc.png`）， 便于导出后仍能看出图片来自哪个原始目录。 """ from __future__ import annotations import argparse import csv import json import os import re 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 IMAGE_SUFFIXES = (".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp") RANK_PREFIX = re.compile(r"^\d{4}_") class SourceResolver: """把 clusters/ 里的条目解析成本机上的**原图**路径。 条目可能是三种东西，依次处理： 1. **硬链接**（link_mode: hard 的默认产物）——它就是原图本身，但路径在 runs/ 下、 文件名还带着 `0125_` 排序前缀。按 inode 反查 assignments.csv 里的原始路径， 这样导出的文件名和 `--prefix-parent` 才能拿到真正的来源目录；查不到就退而用 条目本身、并剥掉排序前缀。 2. **有效软链**——直接 resolve。 3. **失效软链**（从服务器拷回来的常态）——读出链接目标，按其中的 `data/` 段 重新落到本机 `source_root` 上；再不行用文件名查索引。 """ def __init__( self, source_root: Path, marker: str = "data", assignments: Path | None = None ) -> None: self.source_root = source_root.resolve() self.marker = marker self._index: dict[str, Path] | None = None self._by_inode = self._build_inode_index(assignments) @staticmethod def _build_inode_index(assignments: Path | None) -> dict[tuple[int, int], Path]: """assignments.csv 里列出的原图 -> inode，用来认出硬链接指向谁。""" if not assignments or not assignments.exists(): return {} index: dict[tuple[int, int], Path] = {} with assignments.open(encoding="utf-8") as handle: for row in csv.DictReader(handle): path = Path(row.get("path", "")) try: stat = path.stat() except OSError: continue index[(stat.st_dev, stat.st_ino)] = path.resolve() return index 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(): stat = entry.stat() original = self._by_inode.get((stat.st_dev, stat.st_ino)) if original is not None: # 硬链接：认出它指向的原图 return original return self._lookup(RANK_PREFIX.sub("", entry.name)) or entry.resolve() return None 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") 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` 时拼上原父目录名：导出后只剩类别目录，来源信息本来会丢失。 """ return f"{source.parent.name}_{source.name}" if prefix_parent else source.name def check_conflicts(resolved: dict[str, list[Path]], prefix_parent: bool = False) -> None: """同一个类别里出现同名文件时直接报错——不改名就不能悄悄覆盖。""" problems = [] 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" {category} 里 {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="聚类产出目录（取其中的 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( "--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="复制时在文件名前拼上原父目录名（反面纵黑点/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.clusters: logger.error("必须给 --work-dir 或 --clusters 之一") return 1 work_dir = Path(args.work_dir) if args.work_dir else None clusters_dir = Path(args.clusters) if args.clusters else work_dir / "clusters" try: groups = read_clusters(clusters_dir) except (FileNotFoundError, ValueError) as exc: logger.error("%s", exc) return 1 out_dir = ( Path(args.out) if args.out else Path(args.data_root) / clusters_dir.resolve().parent.name ) total = sum(len(v) for v in groups.values()) logger.info("来源: %s（%d 张，%d 个类别）", clusters_dir, total, len(groups)) logger.info("输出目录: %s", out_dir) resolver = SourceResolver( Path(args.source_root), assignments=(work_dir or clusters_dir.parent) / "assignments.csv" ) 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(resolved, args.prefix_parent) except ValueError as exc: logger.error("%s", exc) return 1 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 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 category in sorted(resolved): target_dir = out_dir / category target_dir.mkdir(parents=True, exist_ok=True) entries = [] for source in sorted(resolved[category]): name = destination_name(source, args.prefix_parent) 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[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( { "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"), "category_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())