Shaun ZhouandClaude Opus 5
5d395b7c25
原始数据从 data/<数据集> 移到了 data/raw/<数据集>。
- 聚类配置的 dataset.root：film / pearlescent 均改为 data/raw/ 下
- 检测配置的 img_dir 与 label_dir 注释同步
- README、data/README、算法文档、流程图（含重新渲染的 PNG）与各脚本
docstring 里的路径示例一并更新
- data/README 补充目录约定：原始数据只读放 data/raw/，
export_clusters.py 的派生数据集放在 data/ 下与 raw/ 平级
- 顺手把 unlabeled 数据集加进 datasets/__init__.py：此前它没在包初始化时注册，
build_dataset({"name": "unlabeled"}) 只有在别处先 import 过该模块时才可用
迁移后重跑全量聚类，结果与迁移前逐项一致（轮廓 0.3255、稳定性 0.9641、
簇大小 [4507, 3795, 4849, 1428]）。
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
132 lines
5.7 KiB
Python
132 lines
5.7 KiB
Python
"""薄膜缺陷数据集（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],
|
||
}
|