Shaun ZhouandClaude Opus 5
f433e5cf59
把 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>
124 lines
4.4 KiB
Python
124 lines
4.4 KiB
Python
"""无标注图像数据集，用于无监督聚类。
|
||
|
||
递归扫描一个目录下的所有图片，不需要任何标注。
|
||
|
||
样本还会带出形状特征（宽、高、宽高比、面积的对数），但**默认不参与聚类**：
|
||
这些尺寸是上游 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],
|
||
}
|