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
co-authored by
Claude Opus 5
Claude Opus 5
parent
9b31473b27
commit
f8654f447b
+150
-124
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user