feat(datasets): YOLO 风格标注数据集与几何/像素增强

- yolo: images/ 与同层 labels/ 的路径约定、`class_id cx cy w h` 归一化行解析（兼容分割多边形取外接矩形）、data.yaml 读写；缺失或空 txt 即无缺陷样本 - film_defect: 抛出图像、框、图像级标签；label_mode 决定多热向量还是互斥下标 - transforms: 图框同步的 Resize / Letterbox / CenterCrop / RandomResizedCrop / 翻转 / ColorJitter / GaussianBlur / PerImageStandardize - target_encoder: 高斯半径与 heatmap / wh / offset / mask 监督信号编码 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-02 17:55:20 +08:00
parent 3a36aad6c3
commit af22dc639c
7 changed files with 811 additions and 0 deletions
@@ -0,0 +1,131 @@
+"""薄膜缺陷数据集（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/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],
+    }