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>
This commit is contained in:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
3a36aad6c3
commit
af22dc639c
@@ -0,0 +1,22 @@
|
||||
"""数据集：抛出样本的 Dataset 类与相应的数据变换、标签编码。"""
|
||||
|
||||
from datasets.build import build_dataloader, build_dataset
|
||||
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.yolo import format_label_lines, load_classes, parse_label_file, write_data_yaml
|
||||
|
||||
__all__ = [
|
||||
"build_dataset",
|
||||
"build_dataloader",
|
||||
"FilmDefectDataset",
|
||||
"collate_fn",
|
||||
"DATASETS",
|
||||
"CenterTargetEncoder",
|
||||
"build_transforms",
|
||||
"parse_label_file",
|
||||
"format_label_lines",
|
||||
"load_classes",
|
||||
"write_data_yaml",
|
||||
]
|
||||
@@ -0,0 +1,39 @@
|
||||
"""按配置构建 Dataset / DataLoader。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
|
||||
from datasets.film_defect import collate_fn
|
||||
from datasets.registry import DATASETS
|
||||
|
||||
|
||||
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 = torch.cuda.is_available()
|
||||
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,
|
||||
)
|
||||
@@ -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],
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
from utils.registry import Registry
|
||||
|
||||
DATASETS = Registry("datasets")
|
||||
|
||||
__all__ = ["DATASETS"]
|
||||
@@ -0,0 +1,111 @@
|
||||
"""把框标注编码成 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),
|
||||
}
|
||||
@@ -0,0 +1,314 @@
|
||||
"""同时作用于图像与标注框的几何/像素变换。
|
||||
|
||||
样本约定为 dict:
|
||||
image: PIL.Image 或 torch.Tensor(CHW)
|
||||
boxes: np.ndarray (N, 4) xyxy，原图像素坐标
|
||||
labels: np.ndarray (N,)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import random
|
||||
from collections.abc import Callable, Iterable, Mapping, Sequence
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
Sample = dict
|
||||
|
||||
|
||||
class Compose:
|
||||
def __init__(self, transforms: Sequence[Callable[[Sample], Sample]]) -> None:
|
||||
self.transforms = list(transforms)
|
||||
|
||||
def __call__(self, sample: Sample) -> Sample:
|
||||
for transform in self.transforms:
|
||||
sample = transform(sample)
|
||||
return sample
|
||||
|
||||
|
||||
class Resize:
|
||||
"""按目标尺寸缩放（不保持长宽比），同步缩放框。"""
|
||||
|
||||
def __init__(self, size: Sequence[int]) -> None:
|
||||
self.height, self.width = int(size[0]), int(size[1])
|
||||
|
||||
def __call__(self, sample: Sample) -> Sample:
|
||||
image: Image.Image = sample["image"]
|
||||
src_w, src_h = image.size
|
||||
sample["image"] = image.resize((self.width, self.height), Image.BILINEAR)
|
||||
boxes = sample.get("boxes")
|
||||
if boxes is not None and len(boxes):
|
||||
scale = np.array([self.width / src_w, self.height / src_h] * 2, dtype=np.float32)
|
||||
sample["boxes"] = boxes.astype(np.float32) * scale
|
||||
return sample
|
||||
|
||||
|
||||
class Letterbox:
|
||||
"""保长宽比缩放到目标尺寸，不足的边用 `fill` 填充（居中）。
|
||||
|
||||
薄膜缺陷图的宽高比跨度极大（长条划痕 vs 团块），直接拉伸会把长条压扁到不可辨认；
|
||||
letterbox 保住形状，代价是丢掉绝对尺度——需要尺度信息时用 shape features 补回来。
|
||||
"""
|
||||
|
||||
def __init__(self, size: Sequence[int], fill: int = 114) -> None:
|
||||
self.height, self.width = int(size[0]), int(size[1])
|
||||
self.fill = fill
|
||||
|
||||
def __call__(self, sample: Sample) -> Sample:
|
||||
image: Image.Image = sample["image"]
|
||||
src_w, src_h = image.size
|
||||
scale = min(self.width / src_w, self.height / src_h)
|
||||
new_w, new_h = max(1, round(src_w * scale)), max(1, round(src_h * scale))
|
||||
pad_x, pad_y = (self.width - new_w) // 2, (self.height - new_h) // 2
|
||||
|
||||
fill = self.fill if image.mode != "L" else int(self.fill)
|
||||
canvas = Image.new(image.mode, (self.width, self.height), fill)
|
||||
canvas.paste(image.resize((new_w, new_h), Image.BILINEAR), (pad_x, pad_y))
|
||||
sample["image"] = canvas
|
||||
|
||||
boxes = sample.get("boxes")
|
||||
if boxes is not None and len(boxes):
|
||||
sample["boxes"] = boxes.astype(np.float32) * scale + np.array(
|
||||
[pad_x, pad_y, pad_x, pad_y], dtype=np.float32
|
||||
)
|
||||
sample["letterbox"] = {"scale": scale, "pad": (pad_x, pad_y), "src_size": (src_w, src_h)}
|
||||
return sample
|
||||
|
||||
|
||||
class RandomHorizontalFlip:
|
||||
def __init__(self, p: float = 0.5) -> None:
|
||||
self.p = p
|
||||
|
||||
def __call__(self, sample: Sample) -> Sample:
|
||||
if random.random() >= self.p:
|
||||
return sample
|
||||
image: Image.Image = sample["image"]
|
||||
width = image.size[0]
|
||||
sample["image"] = image.transpose(Image.FLIP_LEFT_RIGHT)
|
||||
boxes = sample.get("boxes")
|
||||
if boxes is not None and len(boxes):
|
||||
flipped = boxes.copy()
|
||||
flipped[:, 0] = width - boxes[:, 2]
|
||||
flipped[:, 2] = width - boxes[:, 0]
|
||||
sample["boxes"] = flipped
|
||||
return sample
|
||||
|
||||
|
||||
class RandomVerticalFlip:
|
||||
def __init__(self, p: float = 0.5) -> None:
|
||||
self.p = p
|
||||
|
||||
def __call__(self, sample: Sample) -> Sample:
|
||||
if random.random() >= self.p:
|
||||
return sample
|
||||
image: Image.Image = sample["image"]
|
||||
height = image.size[1]
|
||||
sample["image"] = image.transpose(Image.FLIP_TOP_BOTTOM)
|
||||
boxes = sample.get("boxes")
|
||||
if boxes is not None and len(boxes):
|
||||
flipped = boxes.copy()
|
||||
flipped[:, 1] = height - boxes[:, 3]
|
||||
flipped[:, 3] = height - boxes[:, 1]
|
||||
sample["boxes"] = flipped
|
||||
return sample
|
||||
|
||||
|
||||
class CenterCrop:
|
||||
"""在原始分辨率上取中心区域，不缩放。
|
||||
|
||||
Resize 的备选：完全不重采样，纹理保持原始尺度，也没有填充边可供网络识别宽高比。
|
||||
图比目标小时用 symmetric 方式补齐（镜像延拓，比灰边少引入人造边缘）。
|
||||
"""
|
||||
|
||||
def __init__(self, size: Sequence[int]) -> None:
|
||||
self.height, self.width = int(size[0]), int(size[1])
|
||||
|
||||
def __call__(self, sample: Sample) -> Sample:
|
||||
array = np.asarray(sample["image"])
|
||||
squeeze = array.ndim == 2
|
||||
if squeeze:
|
||||
array = array[:, :, None]
|
||||
src_h, src_w = array.shape[:2]
|
||||
|
||||
pad_y = max(0, self.height - src_h)
|
||||
pad_x = max(0, self.width - src_w)
|
||||
if pad_y or pad_x:
|
||||
array = np.pad(
|
||||
array,
|
||||
((pad_y // 2, pad_y - pad_y // 2), (pad_x // 2, pad_x - pad_x // 2), (0, 0)),
|
||||
mode="symmetric",
|
||||
)
|
||||
src_h, src_w = array.shape[:2]
|
||||
y0 = (src_h - self.height) // 2
|
||||
x0 = (src_w - self.width) // 2
|
||||
cropped = array[y0 : y0 + self.height, x0 : x0 + self.width]
|
||||
|
||||
sample["image"] = Image.fromarray(cropped[:, :, 0] if squeeze else cropped)
|
||||
boxes = sample.get("boxes")
|
||||
if boxes is not None and len(boxes):
|
||||
shifted = boxes.astype(np.float32) - np.array(
|
||||
[x0 - pad_x // 2, y0 - pad_y // 2] * 2, dtype=np.float32
|
||||
)
|
||||
shifted[:, 0::2] = shifted[:, 0::2].clip(0, self.width)
|
||||
shifted[:, 1::2] = shifted[:, 1::2].clip(0, self.height)
|
||||
keep = (shifted[:, 2] > shifted[:, 0]) & (shifted[:, 3] > shifted[:, 1])
|
||||
sample["boxes"] = shifted[keep]
|
||||
labels = sample.get("labels")
|
||||
if labels is not None and len(labels):
|
||||
sample["labels"] = np.asarray(labels)[keep]
|
||||
return sample
|
||||
|
||||
|
||||
class RandomResizedCrop:
|
||||
"""随机裁一块再缩放回目标尺寸；SimCLR 生成正样本对的主力增强。
|
||||
|
||||
对缺陷小图不能裁太狠——`scale` 下限别低于 0.4，否则容易裁到纯背景，
|
||||
把"背景 vs 背景"当成正样本对，学出来的特征就废了。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
size: Sequence[int],
|
||||
scale: Sequence[float] = (0.5, 1.0),
|
||||
ratio: Sequence[float] = (0.75, 1.333),
|
||||
) -> None:
|
||||
self.height, self.width = int(size[0]), int(size[1])
|
||||
self.scale = (float(scale[0]), float(scale[1]))
|
||||
self.ratio = (float(ratio[0]), float(ratio[1]))
|
||||
|
||||
def __call__(self, sample: Sample) -> Sample:
|
||||
image: Image.Image = sample["image"]
|
||||
src_w, src_h = image.size
|
||||
|
||||
for _ in range(10):
|
||||
area = src_w * src_h * random.uniform(*self.scale)
|
||||
aspect = math.exp(random.uniform(math.log(self.ratio[0]), math.log(self.ratio[1])))
|
||||
crop_w, crop_h = round(math.sqrt(area * aspect)), round(math.sqrt(area / aspect))
|
||||
if 0 < crop_w <= src_w and 0 < crop_h <= src_h:
|
||||
break
|
||||
else: # 退化为中心裁剪
|
||||
crop_w, crop_h = min(src_w, src_h), min(src_w, src_h)
|
||||
|
||||
x0 = random.randint(0, src_w - crop_w)
|
||||
y0 = random.randint(0, src_h - crop_h)
|
||||
sample["image"] = image.crop((x0, y0, x0 + crop_w, y0 + crop_h)).resize(
|
||||
(self.width, self.height), Image.BILINEAR
|
||||
)
|
||||
|
||||
boxes = sample.get("boxes")
|
||||
if boxes is not None and len(boxes):
|
||||
shifted = boxes.astype(np.float32) - np.array([x0, y0, x0, y0], dtype=np.float32)
|
||||
shifted[:, 0::2] = shifted[:, 0::2].clip(0, crop_w) * (self.width / crop_w)
|
||||
shifted[:, 1::2] = shifted[:, 1::2].clip(0, crop_h) * (self.height / crop_h)
|
||||
keep = (shifted[:, 2] > shifted[:, 0]) & (shifted[:, 3] > shifted[:, 1])
|
||||
sample["boxes"] = shifted[keep]
|
||||
labels = sample.get("labels")
|
||||
if labels is not None and len(labels):
|
||||
sample["labels"] = np.asarray(labels)[keep]
|
||||
return sample
|
||||
|
||||
|
||||
class GaussianBlur:
|
||||
"""SimCLR 的标准增强之一，防止网络靠高频噪声区分两个视图。"""
|
||||
|
||||
def __init__(self, p: float = 0.5, radius: Sequence[float] = (0.1, 1.5)) -> None:
|
||||
self.p = p
|
||||
self.radius = (float(radius[0]), float(radius[1]))
|
||||
|
||||
def __call__(self, sample: Sample) -> Sample:
|
||||
if random.random() < self.p:
|
||||
from PIL import ImageFilter
|
||||
|
||||
sample["image"] = sample["image"].filter(
|
||||
ImageFilter.GaussianBlur(random.uniform(*self.radius))
|
||||
)
|
||||
return sample
|
||||
|
||||
|
||||
class ColorJitter:
|
||||
"""亮度/对比度扰动；薄膜图像的光照波动是主要域偏移来源。"""
|
||||
|
||||
def __init__(self, brightness: float = 0.2, contrast: float = 0.2) -> None:
|
||||
self.brightness = brightness
|
||||
self.contrast = contrast
|
||||
|
||||
def __call__(self, sample: Sample) -> Sample:
|
||||
from PIL import ImageEnhance
|
||||
|
||||
image = sample["image"]
|
||||
if self.brightness > 0:
|
||||
factor = 1.0 + random.uniform(-self.brightness, self.brightness)
|
||||
image = ImageEnhance.Brightness(image).enhance(factor)
|
||||
if self.contrast > 0:
|
||||
factor = 1.0 + random.uniform(-self.contrast, self.contrast)
|
||||
image = ImageEnhance.Contrast(image).enhance(factor)
|
||||
sample["image"] = image
|
||||
return sample
|
||||
|
||||
|
||||
class ToTensor:
|
||||
def __call__(self, sample: Sample) -> Sample:
|
||||
array = np.asarray(sample["image"], dtype=np.float32) / 255.0
|
||||
if array.ndim == 2:
|
||||
array = array[:, :, None]
|
||||
sample["image"] = torch.from_numpy(array.transpose(2, 0, 1)).contiguous()
|
||||
return sample
|
||||
|
||||
|
||||
class Normalize:
|
||||
def __init__(self, mean: Sequence[float], std: Sequence[float]) -> None:
|
||||
self.mean = torch.tensor(mean, dtype=torch.float32).view(-1, 1, 1)
|
||||
self.std = torch.tensor(std, dtype=torch.float32).view(-1, 1, 1)
|
||||
|
||||
def __call__(self, sample: Sample) -> Sample:
|
||||
sample["image"] = (sample["image"] - self.mean) / self.std
|
||||
return sample
|
||||
|
||||
|
||||
class PerImageStandardize:
|
||||
"""逐图 z-score：减自身均值、除自身标准差。
|
||||
|
||||
薄膜图之间的整体亮度差异来自光照和相机增益，不是缺陷属性。用固定的
|
||||
ImageNet mean/std 归一化保不住这一点——实测聚类会把大半的簇花在
|
||||
"背景偏亮/偏暗"上，只剩一个簇装真正有缺陷的图。逐图标准化把亮度和
|
||||
对比度这两个全局因子直接除掉，逼特征去看局部结构。
|
||||
"""
|
||||
|
||||
def __init__(self, eps: float = 1e-5) -> None:
|
||||
self.eps = eps
|
||||
|
||||
def __call__(self, sample: Sample) -> Sample:
|
||||
image = sample["image"]
|
||||
sample["image"] = (image - image.mean()) / (image.std() + self.eps)
|
||||
return sample
|
||||
|
||||
|
||||
_TRANSFORMS: dict[str, Callable[..., Any]] = {
|
||||
"Resize": Resize,
|
||||
"Letterbox": Letterbox,
|
||||
"CenterCrop": CenterCrop,
|
||||
"RandomResizedCrop": RandomResizedCrop,
|
||||
"GaussianBlur": GaussianBlur,
|
||||
"RandomHorizontalFlip": RandomHorizontalFlip,
|
||||
"RandomVerticalFlip": RandomVerticalFlip,
|
||||
"ColorJitter": ColorJitter,
|
||||
"ToTensor": ToTensor,
|
||||
"Normalize": Normalize,
|
||||
"PerImageStandardize": PerImageStandardize,
|
||||
}
|
||||
|
||||
|
||||
def build_transforms(cfgs: Iterable[Mapping[str, Any]]) -> Compose:
|
||||
"""由 yaml 列表构建，如 `- {name: Resize, size: [512, 512]}`。"""
|
||||
transforms = []
|
||||
for cfg in cfgs:
|
||||
cfg = dict(cfg)
|
||||
name = cfg.pop("name")
|
||||
if name not in _TRANSFORMS:
|
||||
raise KeyError(f"未知 transform: {name}，可选 {sorted(_TRANSFORMS)}")
|
||||
transforms.append(_TRANSFORMS[name](**cfg))
|
||||
return Compose(transforms)
|
||||
@@ -0,0 +1,189 @@
|
||||
"""YOLO 风格标注的读写。
|
||||
|
||||
目录约定（与 ultralytics 一致，`images` 与 `labels` 同层对应）：
|
||||
|
||||
```
|
||||
data/film/
|
||||
├── images/train/0001.png
|
||||
├── labels/train/0001.txt # 与图片同名，后缀换成 .txt
|
||||
├── images/val/...
|
||||
├── labels/val/...
|
||||
└── data.yaml # path / train / val / names
|
||||
```
|
||||
|
||||
标注文件每行一个目标，坐标是**相对图像宽高归一化**到 [0, 1] 的中心点格式：
|
||||
|
||||
```
|
||||
<class_id> <cx> <cy> <w> <h>
|
||||
```
|
||||
|
||||
也接受 YOLO 分割格式的多边形（`<class_id> x1 y1 x2 y2 ...`，偶数个坐标），
|
||||
读取时取其外接矩形。空文件或缺失文件表示该图无缺陷（纯负样本）。
|
||||
"""
|
||||
|
||||
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],
|
||||
root: str | Path = ".",
|
||||
train: str = "images/train",
|
||||
val: str = "images/val",
|
||||
) -> Path:
|
||||
"""写一个 YOLO 风格的 data.yaml。"""
|
||||
import yaml
|
||||
|
||||
path = Path(path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
"path": str(root),
|
||||
"train": train,
|
||||
"val": val,
|
||||
"nc": len(classes),
|
||||
"names": {i: name for i, name in enumerate(classes)},
|
||||
},
|
||||
allow_unicode=True,
|
||||
sort_keys=False,
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
return path
|
||||
Reference in New Issue
Block a user