feat: 未分类样本的无监督聚类（多 stage 特征 + KMeans，可选 SimCLR）
把 data/film 下所有图片拉平做聚类，不使用任何目录名，簇数由 num_clusters 决定。 预处理针对该批数据的实际情况（14.5k 张已裁好的缺陷小图，中位 152x113， 宽高比 0.24~39.9）：尺寸与亮度都是 ROI 裁剪和光照的产物、不是缺陷属性， 因此拉伸到 128x128（不用 letterbox，其填充比例等价于编码了宽高比）并逐图 z-score。 - cluster: torch 实现的 kmeans++/Lloyd（含空簇处理与 n_init）、PCA、轮廓系数、 NMI 稳定性、多 stage 特征提取；不引 scikit-learn - 白化默认关闭：CNN 特征谱极陡（λ1/λ32≈165x），白化会把噪声方向抬到与信号同权， 实测最大簇一度吃掉 96% 的样本 - simclr: NT-Xent 自监督预训练，用项目自己的卷积骨架，权重可直接初始化 DefectNet； 属于训练，走 require_cuda() 守卫 - 产出 assignments.csv、summary.json、按簇的软链目录与代表样本拼图，供人工核对 全量 14,579 张：轮廓系数 0.326、换种子稳定性 0.964，CPU 约 100 秒。 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
4b1147ed8f
commit
f433e5cf59
@@ -0,0 +1,296 @@
|
||||
#!/usr/bin/env python
|
||||
"""把未分类的样本做无监督聚类，产出可人工核对的分组结果。
|
||||
|
||||
# 默认 4 类，ImageNet 预训练特征（不训练，CPU 可跑）
|
||||
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml
|
||||
|
||||
# 改簇数 / 换成在本地数据上自监督训练的特征（需要 Nvidia GPU）
|
||||
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 NoNvidiaGpuError, require_cuda, resolve_device # noqa: E402
|
||||
|
||||
logger = None
|
||||
|
||||
|
||||
def build_encoder(cfg, device, work_dir):
|
||||
"""按 `method` 拿一个特征提取器。
|
||||
|
||||
pretrained: ImageNet 预训练卷积骨架，纯推理，CPU 也能跑。用顶层的 `backbone`。
|
||||
simclr: 在本目录数据上自监督训练一个项目自己的卷积骨架——属于训练，先过 GPU 守卫。
|
||||
用 `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")
|
||||
|
||||
require_cuda() # 自监督预训练是训练，禁止在无 Nvidia GPU 的设备上跑
|
||||
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)
|
||||
|
||||
checkpoint = simclr_cfg.pop("checkpoint", None)
|
||||
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,
|
||||
)
|
||||
SimCLRTrainer(backbone, loader, simclr_cfg, device, work_dir).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("--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(str(cfg.get("device", "cuda")))
|
||||
except NoNvidiaGpuError as exc:
|
||||
logger.error("%s", exc)
|
||||
return 1
|
||||
|
||||
dataset = build_dataset(dict(cfg["dataset"]))
|
||||
logger.info(
|
||||
"扫描到 %d 张图，来自 %d 个目录；簇数 %d，method=%s，device=%s",
|
||||
len(dataset),
|
||||
len({dataset.source_dir(path) for path in dataset.images}),
|
||||
num_clusters,
|
||||
cfg.get("method", "pretrained"),
|
||||
device,
|
||||
)
|
||||
|
||||
try:
|
||||
encoder = build_encoder(cfg, device, work_dir)
|
||||
except NoNvidiaGpuError as exc:
|
||||
logger.error("%s", exc)
|
||||
return 1
|
||||
|
||||
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())
|
||||
Reference in New Issue
Block a user