#!/usr/bin/env python """把未分类的样本做无监督聚类，产出可人工核对的分组结果。 # 默认 4 类，ImageNet 预训练特征（不训练，CPU 可跑） uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml # 改簇数 / 换成在本地数据上自监督训练的特征（Mac 上自动走 MPS） # 断点续训：--set method=simclr simclr.resume=runs/simclr/simclr_last.pth uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \ --set num_clusters=6 method=simclr 产出（work_dir 下）： assignments.csv 每张图的簇号、到质心距离、原目录 summary.json 簇大小、轮廓系数、簇 × 原目录交叉表 clusters/cluster_00/... 按簇建的软链目录，直接用看图软件翻 cluster_00_preview.png 每簇代表样本（离质心最近）的拼图 """ from __future__ import annotations import argparse import csv import json import shutil import sys from collections import defaultdict from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parents[1])) import torch # noqa: E402 from datasets.build import build_dataset # noqa: E402 from datasets.unlabeled import unlabeled_collate # noqa: E402 from engine.cluster import ( # noqa: E402 StageFeatureExtractor, cluster_features, cluster_stability, extract_features, random_fit_index, ) from models.registry import BACKBONES # noqa: E402 from utils import load_config, seed_everything, setup_logging # noqa: E402 from utils.device import ( # noqa: E402 DeviceNotAvailableError, describe_device, resolve_device, ) logger = None def build_encoder(cfg, device, work_dir): """按 `method` 拿一个特征提取器。 pretrained: ImageNet 预训练卷积骨架，纯推理，CPU 也能跑。用顶层的 `backbone`。 simclr: 在本目录数据上自监督训练一个项目自己的卷积骨架。 用 `simclr.backbone`，与 pretrained 的骨架配置各自独立：两者的字段完全不同 （arch/pretrained vs widths/depths），共用一个块会在 `_base_` 合并时串味。 """ method = str(cfg.get("method", "pretrained")) if method == "pretrained": backbone = BACKBONES.build(dict(cfg.get("backbone", {"name": "torchvision"}))) return StageFeatureExtractor(backbone, cfg.get("feature_stages")).to(device).eval() if method != "simclr": raise ValueError(f"未知 method: {method}，可选 pretrained / simclr") from torch.utils.data import DataLoader from engine.simclr import SimCLRTrainer simclr_cfg = dict(cfg.get("simclr", {})) backbone_cfg = dict(simclr_cfg.pop("backbone", {"name": "tiny_conv"})) # 自研骨架不是 nn.Sequential，默认取最终输出而不是中间 stage stages = simclr_cfg.pop("feature_stages", None) backbone = BACKBONES.build(backbone_cfg).to(device) resume = simclr_cfg.pop("resume", None) checkpoint = simclr_cfg.pop("checkpoint", None) if resume and checkpoint: raise ValueError( "simclr.resume 与 simclr.checkpoint 只能二选一：" "resume=接着训练，checkpoint=跳过训练直接聚类。" ) if checkpoint and Path(checkpoint).exists(): state = torch.load(checkpoint, map_location=device, weights_only=False) backbone.load_state_dict(state["model"]) logger.info("复用已有的自监督骨架: %s", checkpoint) return StageFeatureExtractor(backbone, stages).to(device).eval() view_dataset = build_dataset( {**dict(cfg["dataset"]), "views": 2, "view_transforms": cfg.get("view_transforms", [])} ) loader = DataLoader( view_dataset, batch_size=int(simclr_cfg.pop("batch_size", 128)), shuffle=True, drop_last=True, num_workers=int(cfg.get("loader", {}).get("num_workers", 4)), collate_fn=unlabeled_collate, ) full_config = cfg.to_dict() if hasattr(cfg, "to_dict") else dict(cfg) SimCLRTrainer( backbone, loader, simclr_cfg, device, work_dir, full_config=full_config, resume=resume ).fit() return StageFeatureExtractor(backbone, stages).to(device).eval() def write_outputs(cfg, work_dir, metas, result, num_clusters, stability): """CSV + JSON + 软链目录 + 预览拼图。人工只需要翻软链目录就能确认分组是否合理。""" labels = result.labels.tolist() distances = result.distances.tolist() 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 meta, label, distance in zip(metas, labels, distances, strict=True): width, height = meta["size"] writer.writerow([meta["path"], label, f"{distance:.6f}", meta["source_dir"], width, height]) summary = { "num_clusters": num_clusters, "num_images": len(metas), "method": cfg.get("method", "pretrained"), "image_size": cfg.get("image_size"), "cluster_sizes": result.sizes, "silhouette": result.silhouette, "stability": stability, "inertia": result.inertia, "mean_distance_per_cluster": [ float(result.distances[result.labels == index].mean()) if size else None for index, size in enumerate(result.sizes) ], } (work_dir / "summary.json").write_text( json.dumps(summary, indent=2, ensure_ascii=False), encoding="utf-8" ) link_limit = int(cfg.get("link_limit", 300)) preview_count = int(cfg.get("preview_count", 25)) clusters_dir = work_dir / "clusters" if clusters_dir.exists(): shutil.rmtree(clusters_dir) order = sorted(range(len(metas)), key=lambda i: distances[i]) # 离质心近的排前面 per_cluster: dict[int, list[int]] = defaultdict(list) for index in order: per_cluster[labels[index]].append(index) for cluster, indices in sorted(per_cluster.items()): target = clusters_dir / f"cluster_{cluster:02d}" target.mkdir(parents=True, exist_ok=True) # 近的取大头，末尾补几张最远的——离群样本最能暴露这个簇是不是混了东西 chosen = indices[: max(link_limit - 20, 1)] + indices[-20:] if len(indices) > link_limit else indices for rank, index in enumerate(chosen): source = Path(metas[index]["path"]).resolve() link = target / f"{rank:04d}_{source.stem}{source.suffix}" if not link.exists(): link.symlink_to(source) save_preview(work_dir / f"cluster_{cluster:02d}_preview.png", metas, indices[:preview_count]) return summary def save_preview(path: Path, metas: list[dict], indices: list[int], cell: int = 96, columns: int = 5): """把代表样本拼成一张图，方便一眼看出这个簇是什么缺陷。""" if not indices: return from PIL import Image rows = (len(indices) + columns - 1) // columns canvas = Image.new("RGB", (columns * cell, rows * cell), (32, 32, 32)) for position, index in enumerate(indices): with Image.open(metas[index]["path"]) as raw: image = raw.convert("RGB") scale = min(cell / image.width, cell / image.height) image = image.resize((max(1, round(image.width * scale)), max(1, round(image.height * scale)))) x = (position % columns) * cell + (cell - image.width) // 2 y = (position // columns) * cell + (cell - image.height) // 2 canvas.paste(image, (x, y)) canvas.save(path) def sweep_clusters(embeddings, spec: str, cluster_kwargs: dict) -> None: """扫一遍簇数，同时看轮廓系数和换种子的稳定性。 无标签时没有"正确答案"，只能看：簇是否紧致（轮廓），以及换初始化后划分是否还一样 （稳定性）。稳定性低说明那个 k 是在硬切一团连续的数据。 """ start, end = (int(v) for v in spec.split("-")) logger.info("%3s %9s %10s %12s 簇大小", "k", "轮廓", "稳定性", "最大簇占比") for k in range(start, end + 1): result = cluster_features(embeddings, k, seed=42, **cluster_kwargs) stability = cluster_stability(embeddings, k, seeds=(0, 7), **cluster_kwargs) logger.info( "%3d %9.3f %10.3f %12.2f %s", k, result.silhouette, stability, max(result.sizes) / sum(result.sizes), result.sizes, ) def main() -> int: global logger parser = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter ) parser.add_argument("--config", default="configs/cluster/film.yaml") parser.add_argument("--root", default=None, help="覆盖配置里的图片根目录") parser.add_argument("--num-clusters", type=int, default=None, help="覆盖配置里的簇数") parser.add_argument("--work-dir", default=None) parser.add_argument("--device", default=None, help="覆盖配置中的 device：auto / cuda / mps / cpu") parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE") parser.add_argument("--save-features", action="store_true", help="额外保存 embedding 供复用") parser.add_argument( "--sweep-k", default=None, metavar="A-B", help="只扫簇数（如 2-8）打印指标表，不产出分组" ) args = parser.parse_args() cfg = load_config(args.config, args.set) if args.root: cfg["dataset"]["root"] = args.root if args.num_clusters: cfg["num_clusters"] = args.num_clusters work_dir = Path(args.work_dir or cfg.get("work_dir", "runs/cluster")) work_dir.mkdir(parents=True, exist_ok=True) logger = setup_logging(work_dir / "cluster.log") seed_everything(int(cfg.get("seed", 42))) num_clusters = int(cfg.get("num_clusters", 4)) try: device = resolve_device(args.device or cfg.get("device", "auto")) except DeviceNotAvailableError as exc: logger.error("%s", exc) return 1 dataset = build_dataset(dict(cfg["dataset"])) logger.info( "扫描到 %d 张图，来自 %d 个目录；簇数 %d，method=%s，设备 %s", len(dataset), len({dataset.source_dir(path) for path in dataset.images}), num_clusters, cfg.get("method", "pretrained"), describe_device(device), ) encoder = build_encoder(cfg, device, work_dir) from torch.utils.data import DataLoader loader = DataLoader( dataset, batch_size=int(cfg.get("loader", {}).get("batch_size", 64)), shuffle=False, num_workers=int(cfg.get("loader", {}).get("num_workers", 4)), collate_fn=unlabeled_collate, ) embeddings, metas = extract_features( encoder, loader, device, use_shape_features=bool(cfg.get("use_shape_features", False)) ) logger.info("embedding: %s", tuple(embeddings.shape)) cluster_kwargs = { "pca_dim": int(cfg.get("pca_dim", 32)), "whiten": bool(cfg.get("whiten", False)), "n_init": int(cfg.get("n_init", 10)), } fit_index = random_fit_index( len(metas), int(cfg.get("fit_max_samples", 0)), seed=int(cfg.get("seed", 42)) ) if fit_index is not None: logger.info("随机抽 %d 个样本拟合质心（拟合后仍分配全部样本）", len(fit_index)) if args.sweep_k: sweep_clusters(embeddings, args.sweep_k, cluster_kwargs) return 0 result = cluster_features( embeddings, num_clusters, seed=int(cfg.get("seed", 42)), fit_index=fit_index, **cluster_kwargs ) stability = cluster_stability(embeddings, num_clusters, seeds=(0, 7), **cluster_kwargs) logger.info( "簇大小: %s | 轮廓系数: %.4f | 换种子稳定性: %.4f", result.sizes, result.silhouette, stability, ) if args.save_features: torch.save({"embeddings": embeddings, "metas": metas}, work_dir / "features.pt") write_outputs(cfg, work_dir, metas, result, num_clusters, stability) logger.info("结果已写入 %s（clusters/ 下是按簇建的软链目录，直接翻图即可核对）", work_dir) for index, size in enumerate(result.sizes): mean_distance = float(result.distances[result.labels == index].mean()) if size else float("nan") logger.info(" cluster_%02d %6d 张 平均距质心 %.3f", index, size, mean_distance) return 0 if __name__ == "__main__": raise SystemExit(main())