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:
co-authored by
Claude Opus 5
parent
4b1147ed8f
commit
f433e5cf59
@@ -0,0 +1,123 @@
|
||||
"""无标注图像数据集,用于无监督聚类。
|
||||
|
||||
递归扫描一个目录下的所有图片,不需要任何标注。
|
||||
|
||||
样本还会带出形状特征(宽、高、宽高比、面积的对数),但**默认不参与聚类**:
|
||||
这些尺寸是上游 ROI 裁剪的产物,不是缺陷本身的属性,拿它做特征只会让聚类
|
||||
去区分采集方式。保留它只是为了写进 assignments.csv 供人工核对、以及做对照实验。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from collections.abc import Mapping, Sequence
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
from datasets.registry import DATASETS
|
||||
from datasets.transforms import build_transforms
|
||||
from datasets.yolo import IMAGE_SUFFIXES
|
||||
|
||||
SHAPE_FEATURE_DIM = 4
|
||||
|
||||
|
||||
def shape_features(width: int, height: int) -> torch.Tensor:
|
||||
"""[log 宽, log 高, log 宽高比, log 面积],取对数是因为尺寸跨了两个数量级。"""
|
||||
return torch.tensor(
|
||||
[
|
||||
math.log(width),
|
||||
math.log(height),
|
||||
math.log(width / height),
|
||||
math.log(width * height),
|
||||
],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
|
||||
@DATASETS.register("unlabeled")
|
||||
class UnlabeledImageDataset(Dataset):
|
||||
"""Args:
|
||||
root: 图片根目录,递归扫描。
|
||||
transforms: transform 配置列表;聚类默认用 Letterbox + ToTensor + Normalize。
|
||||
gray: 是否转单通道(RGBA / 调色板图会先统一到 RGB 或 L)。
|
||||
views: 每个样本返回几个增强视图。1 = 推理取特征;2 = SimCLR 对比学习。
|
||||
view_transforms: `views > 1` 时每个视图各自过一遍的随机增强。
|
||||
exclude: 需要跳过的目录名(默认跳过已划分好的 images/labels 与临时目录)。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
root: str | Path,
|
||||
transforms: Sequence[Mapping[str, Any]] = (),
|
||||
gray: bool = False,
|
||||
views: int = 1,
|
||||
view_transforms: Sequence[Mapping[str, Any]] = (),
|
||||
exclude: Sequence[str] = (),
|
||||
) -> None:
|
||||
self.root = Path(root)
|
||||
if not self.root.is_dir():
|
||||
raise FileNotFoundError(f"目录不存在: {self.root}")
|
||||
self.gray = gray
|
||||
self.views = int(views)
|
||||
self.transforms = build_transforms(transforms)
|
||||
self.view_transforms = build_transforms(view_transforms) if view_transforms else None
|
||||
self.exclude = set(exclude)
|
||||
|
||||
self.images = sorted(
|
||||
path
|
||||
for path in self.root.rglob("*")
|
||||
if path.suffix.lower() in IMAGE_SUFFIXES
|
||||
and not path.name.startswith(".")
|
||||
and not self.exclude & set(path.relative_to(self.root).parts)
|
||||
)
|
||||
if not self.images:
|
||||
raise FileNotFoundError(f"{self.root} 下没有找到图片")
|
||||
|
||||
def source_dir(self, path: Path) -> str:
|
||||
"""图片相对 root 的父目录,用来和聚类结果做交叉表(人工已分好的部分即是弱标签)。"""
|
||||
relative = path.relative_to(self.root).parent
|
||||
return str(relative) if str(relative) != "." else "(root)"
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.images)
|
||||
|
||||
def __getitem__(self, index: int) -> dict:
|
||||
path = self.images[index]
|
||||
with Image.open(path) as raw:
|
||||
image = raw.convert("L" if self.gray else "RGB") # RGBA / P 一并统一掉
|
||||
width, height = image.size
|
||||
|
||||
if self.views == 1:
|
||||
views = [self.transforms({"image": image.copy()})["image"]]
|
||||
else:
|
||||
views = []
|
||||
for _ in range(self.views):
|
||||
sample = {"image": image.copy()}
|
||||
if self.view_transforms is not None:
|
||||
sample = self.view_transforms(sample)
|
||||
views.append(self.transforms(sample)["image"])
|
||||
|
||||
return {
|
||||
"views": views,
|
||||
"shape": shape_features(width, height),
|
||||
"meta": {
|
||||
"path": str(path),
|
||||
"index": index,
|
||||
"source_dir": self.source_dir(path),
|
||||
"size": (width, height),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def unlabeled_collate(batch: list[dict]) -> dict:
|
||||
"""把每个视图各自堆成一个 batch:views[i] 形状 (B, C, H, W)。"""
|
||||
num_views = len(batch[0]["views"])
|
||||
return {
|
||||
"views": [torch.stack([item["views"][i] for item in batch]) for i in range(num_views)],
|
||||
"shape": torch.stack([item["shape"] for item in batch]),
|
||||
"meta": [item["meta"] for item in batch],
|
||||
}
|
||||
Reference in New Issue
Block a user