refactor: 按实际的五步工作流裁掉无关代码

实际工作流是：simclr 训练 -> 聚类 -> 人工归类 -> 分类训练 -> GradCAM/SHAP 生成
热力图与 box。据此移除不参与其中的检测/定位链路，共 31 个文件、约 2600 行：

- 模型：defect_net、postprocess、necks/、heads/center_head
- 数据：film_defect（YOLO 读取端）、target_encoder
- 引擎：trainer、evaluator、losses
- 入口：train.py、eval.py
- 脚本：convert_annotations、split_dataset、export_onnx、profile_inference
- 配置：configs/{train,eval,model,dataset}
- utils/labels（多类/单类标签模式，仅检测分类头使用）

共享模块按实际用量瘦身：engine/metrics 只留 AverageMeter（检测指标随之移除）；
datasets/yolo 只留写出侧（IMAGE_SUFFIXES 与 write_data_yaml）——第 5 步产出的
YOLO 标签是给外部检测训练用的交换格式，读取端不再需要。

裁剪过程中修掉一个真 bug：generate_boxes 的 --stage 默认 6（layer3）对 tiny_conv
这类非 nn.Sequential 骨架直接抛异常。现在自动退回骨架最终输出并警告，
另加 --stage -1 显式选择最终输出。

五步工作流已端到端实跑验证（simclr 1 轮 -> 聚类 -> 导出 -> 分类 1 轮 ->
gradcam 与 shap 两种方法各生成一遍）。

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

datasets/__init__.py

@@ -1,28 +1,20 @@
-"""数据集：抛出样本的 Dataset 类与相应的数据变换、标签编码。"""
+"""数据集：无标注（自监督/聚类用）与按目录组织的分类数据集。"""
 
-from datasets.build import build_dataloader, build_dataset
+from datasets.build import build_dataset
 from datasets.classification import FolderClassificationDataset, classification_collate
-from datasets.film_defect import FilmDefectDataset, collate_fn
 from datasets.registry import DATASETS
-from datasets.target_encoder import CenterTargetEncoder
 from datasets.transforms import build_transforms
 from datasets.unlabeled import UnlabeledImageDataset, unlabeled_collate
-from datasets.yolo import format_label_lines, load_classes, parse_label_file, write_data_yaml
+from datasets.yolo import IMAGE_SUFFIXES, write_data_yaml
 
 __all__ = [
     "build_dataset",
     "FolderClassificationDataset",
     "classification_collate",
-    "build_dataloader",
-    "FilmDefectDataset",
-    "collate_fn",
-    "DATASETS",
     "UnlabeledImageDataset",
     "unlabeled_collate",
-    "CenterTargetEncoder",
     "build_transforms",
-    "parse_label_file",
-    "format_label_lines",
-    "load_classes",
+    "DATASETS",
+    "IMAGE_SUFFIXES",
     "write_data_yaml",
 ]

datasets/build.py

@@ -1,39 +1,14 @@
-"""按配置构建 Dataset / DataLoader。"""
+"""按配置构建 Dataset。"""
 
 from __future__ import annotations
 
 from collections.abc import Mapping
 from typing import Any
 
-from torch.utils.data import DataLoader, Dataset
+from torch.utils.data import Dataset
 
-from datasets.film_defect import collate_fn
 from datasets.registry import DATASETS
-from utils.device import pin_memory_supported, resolve_device
 
 
 def build_dataset(cfg: Mapping[str, Any], **kwargs: Any) -> Dataset:
     return DATASETS.build(dict(cfg), **kwargs)
-
-
-def build_dataloader(
-    dataset: Dataset,
-    batch_size: int = 8,
-    shuffle: bool = False,
-    num_workers: int = 4,
-    pin_memory: bool | None = None,
-    drop_last: bool = False,
-) -> DataLoader:
-    """`pin_memory` 缺省时只在 CUDA 上开启（CPU/MPS 上无收益且会告警）。"""
-    if pin_memory is None:
-        pin_memory = pin_memory_supported(resolve_device("auto"))
-    return DataLoader(
-        dataset,
-        batch_size=batch_size,
-        shuffle=shuffle,
-        num_workers=num_workers,
-        pin_memory=pin_memory,
-        drop_last=drop_last,
-        collate_fn=collate_fn,
-        persistent_workers=num_workers > 0,
-    )

datasets/film_defect.py

@@ -1,131 +0,0 @@
-"""薄膜缺陷数据集（YOLO 风格标注）。
-
-目录与标注格式见 `datasets/yolo.py` 与 data/README.md：`images/` 下的每张图，
-在同层 `labels/` 下有一个同名 .txt，每行 `class_id cx cy w h`（相对宽高归一化）。
-缺失或空的 .txt 表示该图无缺陷。
-
-图像级标签的形式由配置里的 `label_mode` 决定：
-    multilabel（多类）——多热向量，一张图可同时含多种缺陷，全 0 表示 OK 片；
-    multiclass（单类）——单个类别下标，互斥；无缺陷时取背景下标 len(classes)。
-用 scripts/convert_annotations.py 从 labelme / VOC 转成该格式。
-"""
-
-from __future__ import annotations
-
-from collections.abc import Mapping, Sequence
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-import torch
-from PIL import Image
-from torch.utils.data import Dataset
-
-from datasets.registry import DATASETS
-from datasets.target_encoder import CenterTargetEncoder
-from datasets.transforms import build_transforms
-from datasets.yolo import denormalize, label_path_for, list_images, load_classes, parse_label_file
-from utils.labels import MULTILABEL, check_label_mode
-
-
-@DATASETS.register("film_defect")
-class FilmDefectDataset(Dataset):
-    """Args:
-    img_dir: 图片目录（YOLO 约定，如 `data/raw/film/images/train`），递归查找图片。
-    classes: 类别名列表，或指向 YOLO data.yaml 的路径（读其中的 `names`）。
-        列表顺序即类别下标顺序，必须与标注 txt 里的 class_id 对应。
-    label_dir: 标注目录；缺省按 YOLO 约定把路径里的 `images` 换成 `labels`。
-    list_file: 可选的图片清单（YOLO 的 train.txt 风格），给了就按它的顺序取图。
-    transforms: transform 配置列表，见 `datasets.transforms.build_transforms`。
-    stride: 模型输出下采样倍数，用于生成 heatmap 监督。
-    gray: 是否按单通道读取（薄膜线扫相机通常为灰度）。
-    max_objects: 单图最多编码的目标数。
-    label_mode: "multilabel"(多类，多热向量) 或 "multiclass"(单类，互斥下标)。
-        必须与模型配置的 label_mode 一致，train.py / eval.py 会统一下发。
-    """
-
-    def __init__(
-        self,
-        img_dir: str | Path,
-        classes: Sequence[str] | str | Path,
-        label_dir: str | Path | None = None,
-        list_file: str | Path | None = None,
-        transforms: Sequence[Mapping[str, Any]] = (),
-        stride: int = 4,
-        gray: bool = False,
-        max_objects: int = 64,
-        label_mode: str = "multilabel",
-    ) -> None:
-        self.label_mode = check_label_mode(label_mode)
-        self.img_dir = Path(img_dir)
-        self.label_dir = Path(label_dir) if label_dir is not None else None
-        self.classes = load_classes(classes)
-        self.class_to_idx = {name: i for i, name in enumerate(self.classes)}
-        self.transforms = build_transforms(transforms)
-        self.gray = gray
-        self.encoder = CenterTargetEncoder(
-            num_classes=len(self.classes), stride=stride, max_objects=max_objects
-        )
-
-        self.images = list_images(self.img_dir, Path(list_file) if list_file else None)
-        if not self.images:
-            raise FileNotFoundError(f"{self.img_dir} 下没有找到图片")
-
-    def label_file(self, image_path: Path) -> Path:
-        return label_path_for(image_path, self.img_dir, self.label_dir)
-
-    def __len__(self) -> int:
-        return len(self.images)
-
-    def __getitem__(self, index: int) -> dict:
-        path = self.images[index]
-        image = Image.open(path)
-        image = image.convert("L" if self.gray else "RGB")
-        width, height = image.size
-
-        norm_boxes, labels = parse_label_file(self.label_file(path), num_classes=len(self.classes))
-        boxes = denormalize(norm_boxes, width, height)
-
-        sample = self.transforms({"image": image, "boxes": boxes, "labels": labels})
-        image_tensor: torch.Tensor = sample["image"]
-        boxes = np.asarray(sample["boxes"], dtype=np.float32).reshape(-1, 4)
-        labels = np.asarray(sample["labels"], dtype=np.int64)
-
-        target = self.encoder(boxes, labels, image_tensor.shape[-2:])
-        image_label = self._image_label(labels)
-
-        return {
-            "image": image_tensor,
-            "target": target,
-            "boxes": torch.from_numpy(boxes),
-            "labels": torch.from_numpy(labels),
-            "image_label": image_label,
-            "meta": {"path": str(path), "index": index},
-        }
-
-    def _image_label(self, labels: np.ndarray) -> torch.Tensor:
-        """多类 -> 多热向量 (num_classes,)；单类 -> 标量下标（背景为 len(classes)）。"""
-        if self.label_mode == MULTILABEL:
-            image_label = torch.zeros(len(self.classes), dtype=torch.float32)
-            if len(labels):
-                image_label[torch.from_numpy(labels)] = 1.0
-            return image_label
-
-        if not len(labels):
-            return torch.tensor(len(self.classes), dtype=torch.long)
-        # 一图多种缺陷时取框数最多的类别（并列取下标最小的），保证结果确定。
-        counts = np.bincount(labels, minlength=len(self.classes))
-        return torch.tensor(int(counts.argmax()), dtype=torch.long)
-
-
-def collate_fn(batch: list[dict]) -> dict:
-    """框数量不定，按 list 保留；其余堆叠成 batch 张量。"""
-    target_keys = batch[0]["target"].keys()
-    return {
-        "images": torch.stack([item["image"] for item in batch]),
-        "target": {key: torch.stack([item["target"][key] for item in batch]) for key in target_keys},
-        "boxes": [item["boxes"] for item in batch],
-        "labels": [item["labels"] for item in batch],
-        "image_labels": torch.stack([item["image_label"] for item in batch]),
-        "meta": [item["meta"] for item in batch],
-    }

datasets/target_encoder.py

@@ -1,111 +0,0 @@
-"""把框标注编码成 CenterHead 需要的监督信号。"""
-
-from __future__ import annotations
-
-import numpy as np
-import torch
-
-
-def gaussian_radius(height: float, width: float, min_overlap: float = 0.7) -> float:
-    """CenterNet 的半径公式：保证半径内的框与 GT 的 IoU >= min_overlap。"""
-    a1 = 1.0
-    b1 = height + width
-    c1 = width * height * (1 - min_overlap) / (1 + min_overlap)
-    r1 = (b1 - np.sqrt(max(b1**2 - 4 * a1 * c1, 0.0))) / (2 * a1)
-
-    a2 = 4.0
-    b2 = 2 * (height + width)
-    c2 = (1 - min_overlap) * width * height
-    r2 = (b2 - np.sqrt(max(b2**2 - 4 * a2 * c2, 0.0))) / (2 * a2)
-
-    a3 = 4 * min_overlap
-    b3 = -2 * min_overlap * (height + width)
-    c3 = (min_overlap - 1) * width * height
-    r3 = (-b3 + np.sqrt(max(b3**2 - 4 * a3 * c3, 0.0))) / (2 * a3)
-
-    return max(0.0, min(r1, r2, r3))
-
-
-def draw_gaussian(heatmap: np.ndarray, center: tuple[int, int], radius: int) -> None:
-    """在 heatmap 上叠加二维高斯（取 max，避免相邻目标互相削弱）。"""
-    diameter = 2 * radius + 1
-    sigma = diameter / 6.0
-    grid = np.arange(-radius, radius + 1, dtype=np.float32)
-    gaussian = np.exp(-(grid[None, :] ** 2 + grid[:, None] ** 2) / (2 * sigma**2))
-
-    cx, cy = center
-    height, width = heatmap.shape
-    left, right = min(cx, radius), min(width - cx, radius + 1)
-    top, bottom = min(cy, radius), min(height - cy, radius + 1)
-    if right <= -left or bottom <= -top:
-        return
-
-    masked_heatmap = heatmap[cy - top : cy + bottom, cx - left : cx + right]
-    masked_gaussian = gaussian[radius - top : radius + bottom, radius - left : radius + right]
-    np.maximum(masked_heatmap, masked_gaussian, out=masked_heatmap)
-
-
-class CenterTargetEncoder:
-    """生成 heatmap / wh / offset / mask 四个监督张量。
-
-    Args:
-        num_classes: 类别数。
-        stride: 模型输出相对输入的下采样倍数。
-        max_objects: 单图最多编码的目标数。
-        min_overlap: 高斯半径的 IoU 阈值。
-    """
-
-    def __init__(
-        self,
-        num_classes: int,
-        stride: int,
-        max_objects: int = 64,
-        min_overlap: float = 0.7,
-    ) -> None:
-        self.num_classes = num_classes
-        self.stride = stride
-        self.max_objects = max_objects
-        self.min_overlap = min_overlap
-
-    def __call__(
-        self,
-        boxes: np.ndarray,
-        labels: np.ndarray,
-        image_size: tuple[int, int],
-    ) -> dict[str, torch.Tensor]:
-        img_h, img_w = image_size
-        out_h, out_w = img_h // self.stride, img_w // self.stride
-
-        heatmap = np.zeros((self.num_classes, out_h, out_w), dtype=np.float32)
-        wh = np.zeros((2, out_h, out_w), dtype=np.float32)
-        offset = np.zeros((2, out_h, out_w), dtype=np.float32)
-        mask = np.zeros((out_h, out_w), dtype=np.float32)
-
-        count = 0
-        for box, label in zip(np.asarray(boxes, dtype=np.float32), np.asarray(labels), strict=True):
-            if count >= self.max_objects:
-                break
-            x1, y1, x2, y2 = box / self.stride
-            box_w, box_h = x2 - x1, y2 - y1
-            if box_w <= 0 or box_h <= 0:
-                continue
-
-            cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
-            cx_int, cy_int = int(cx), int(cy)
-            if not (0 <= cx_int < out_w and 0 <= cy_int < out_h):
-                continue
-
-            radius = max(1, int(gaussian_radius(box_h, box_w, self.min_overlap)))
-            draw_gaussian(heatmap[int(label)], (cx_int, cy_int), radius)
-
-            wh[:, cy_int, cx_int] = (box_w, box_h)
-            offset[:, cy_int, cx_int] = (cx - cx_int, cy - cy_int)
-            mask[cy_int, cx_int] = 1.0
-            count += 1
-
-        return {
-            "heatmap": torch.from_numpy(heatmap),
-            "wh": torch.from_numpy(wh),
-            "offset": torch.from_numpy(offset),
-            "mask": torch.from_numpy(mask),
-        }

datasets/yolo.py

@@ -1,24 +1,13 @@
-"""YOLO 风格标注的读写。
+"""YOLO 数据集格式的共用定义。
 
-目录约定（与 ultralytics 一致，`images` 与 `labels` 同层对应）：
+工作流里只用到"写出"这一侧：scripts/generate_boxes.py 产出的
+`labels/<类别>/<图名>.txt` 与 `data.yaml` 是给外部检测训练用的交换格式。
 
-```
-data/raw/film/
-├── images/train/0001.png
-├── labels/train/0001.txt      # 与图片同名，后缀换成 .txt
-├── images/val/...
-├── labels/val/...
-└── data.yaml                  # path / train / val / names
-```
-
-标注文件每行一个目标，坐标是**相对图像宽高归一化**到 [0, 1] 的中心点格式：
+标注文件每行一个目标，坐标相对图像宽高归一化到 [0, 1]，中心点格式：
 
 ```
 <class_id> <cx> <cy> <w> <h>
 ```
-
-也接受 YOLO 分割格式的多边形（`<class_id> x1 y1 x2 y2 ...`，偶数个坐标），
-读取时取其外接矩形。空文件或缺失文件表示该图无缺陷（纯负样本）。
 """
 
 from __future__ import annotations
@@ -26,140 +15,9 @@ from __future__ import annotations
 from collections.abc import Sequence
 from pathlib import Path
 
-import numpy as np
-
 IMAGE_SUFFIXES = (".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp")
 
 
-def label_path_for(image_path: Path, img_dir: Path, label_dir: Path | None = None) -> Path:
-    """图片路径 -> 标注路径。
-
-    缺省按 YOLO 约定把路径里最后一个 `images` 目录名换成 `labels`；
-    显式给了 `label_dir` 时，保持图片相对 `img_dir` 的子目录结构。
-    """
-    relative = image_path.relative_to(img_dir).with_suffix(".txt")
-    if label_dir is not None:
-        return Path(label_dir) / relative
-
-    parts = list(img_dir.parts)
-    for i in range(len(parts) - 1, -1, -1):
-        if parts[i] == "images":
-            parts[i] = "labels"
-            return Path(*parts) / relative
-    # 没有 images 目录时退化为同级的 labels/
-    return img_dir.parent / "labels" / relative
-
-
-def list_images(img_dir: Path, list_file: Path | None = None) -> list[Path]:
-    """列出图片；给了 `list_file` 则按文件里的路径顺序（YOLO 的 train.txt 风格）。"""
-    if list_file is not None:
-        list_file = Path(list_file)
-        if not list_file.exists():
-            raise FileNotFoundError(f"图片清单不存在: {list_file}")
-        images = []
-        for line in list_file.read_text(encoding="utf-8").splitlines():
-            line = line.strip()
-            if not line or line.startswith("#"):
-                continue
-            path = Path(line)
-            images.append(path if path.is_absolute() else (list_file.parent / path).resolve())
-        return images
-
-    if not img_dir.is_dir():
-        raise FileNotFoundError(f"图片目录不存在: {img_dir}")
-    return sorted(p for p in img_dir.rglob("*") if p.suffix.lower() in IMAGE_SUFFIXES)
-
-
-def parse_label_file(path: Path, num_classes: int | None = None) -> tuple[np.ndarray, np.ndarray]:
-    """读取一个 .txt，返回归一化的 xyxy 框 (N, 4) 与类别下标 (N,)。
-
-    缺失文件按无缺陷处理（YOLO 的惯例），格式错误则带上 `文件:行号` 抛出。
-    """
-    boxes: list[list[float]] = []
-    labels: list[int] = []
-    if not path.exists():
-        return np.zeros((0, 4), dtype=np.float32), np.zeros((0,), dtype=np.int64)
-
-    for line_no, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1):
-        line = line.strip()
-        if not line or line.startswith("#"):
-            continue
-        parts = line.split()
-        try:
-            class_id = int(float(parts[0]))
-            values = [float(v) for v in parts[1:]]
-        except ValueError as exc:
-            raise ValueError(f"{path}:{line_no} 不是合法的 YOLO 标注行: {line!r}") from exc
-
-        if num_classes is not None and not 0 <= class_id < num_classes:
-            raise ValueError(
-                f"{path}:{line_no} 类别下标 {class_id} 越界，classes 只有 {num_classes} 个"
-            )
-
-        if len(values) == 4:
-            cx, cy, width, height = values
-            box = [cx - width / 2, cy - height / 2, cx + width / 2, cy + height / 2]
-        elif len(values) >= 6 and len(values) % 2 == 0:
-            # 分割格式的多边形，取外接矩形
-            xs, ys = values[0::2], values[1::2]
-            box = [min(xs), min(ys), max(xs), max(ys)]
-        else:
-            raise ValueError(
-                f"{path}:{line_no} 期望 `class cx cy w h` 或分割多边形，实际有 {len(values)} 个坐标"
-            )
-
-        box = [min(max(v, 0.0), 1.0) for v in box]
-        if box[2] <= box[0] or box[3] <= box[1]:
-            continue  # 退化框直接丢弃
-        boxes.append(box)
-        labels.append(class_id)
-
-    return (
-        np.asarray(boxes, dtype=np.float32).reshape(-1, 4),
-        np.asarray(labels, dtype=np.int64),
-    )
-
-
-def denormalize(boxes: np.ndarray, width: int, height: int) -> np.ndarray:
-    """归一化 xyxy -> 像素 xyxy。"""
-    if not len(boxes):
-        return np.zeros((0, 4), dtype=np.float32)
-    scale = np.array([width, height, width, height], dtype=np.float32)
-    return boxes.astype(np.float32) * scale
-
-
-def format_label_lines(boxes: np.ndarray, labels: Sequence[int], width: int, height: int) -> list[str]:
-    """像素 xyxy -> YOLO 文本行（归一化 cx cy w h，保留 6 位小数）。"""
-    lines = []
-    for (x1, y1, x2, y2), label in zip(np.asarray(boxes, dtype=np.float64), labels, strict=True):
-        x1, x2 = sorted((max(x1, 0.0), min(x2, float(width))))
-        y1, y2 = sorted((max(y1, 0.0), min(y2, float(height))))
-        box_w, box_h = (x2 - x1) / width, (y2 - y1) / height
-        if box_w <= 0 or box_h <= 0:
-            continue
-        cx, cy = (x1 + x2) / 2 / width, (y1 + y2) / 2 / height
-        lines.append(f"{int(label)} {cx:.6f} {cy:.6f} {box_w:.6f} {box_h:.6f}")
-    return lines
-
-
-def load_classes(spec: Sequence[str] | str | Path) -> list[str]:
-    """类别名可以直接写成列表，也可以指向 YOLO 的 data.yaml（读其中的 `names`）。"""
-    if not isinstance(spec, str | Path):
-        return list(spec)
-
-    import yaml
-
-    path = Path(spec)
-    if not path.exists():
-        raise FileNotFoundError(f"data.yaml 不存在: {path}")
-    names = (yaml.safe_load(path.read_text(encoding="utf-8")) or {}).get("names")
-    if names is None:
-        raise KeyError(f"{path} 里没有 `names` 字段")
-    if isinstance(names, dict):  # {0: scratch, 1: bubble}
-        return [names[key] for key in sorted(names, key=int)]
-    return list(names)
-
-
 def write_data_yaml(
     path: str | Path,
     classes: Sequence[str],
@@ -167,7 +25,7 @@ def write_data_yaml(
     train: str = "images/train",
     val: str = "images/val",
 ) -> Path:
-    """写一个 YOLO 风格的 data.yaml。"""
+    """写一个 YOLO 风格的 data.yaml，字段与 ultralytics 通用。"""
     import yaml
 
     path = Path(path)
@@ -179,7 +37,7 @@ def write_data_yaml(
                 "train": train,
                 "val": val,
                 "nc": len(classes),
-                "names": {i: name for i, name in enumerate(classes)},
+                "names": dict(enumerate(classes)),
             },
             allow_unicode=True,
             sort_keys=False,