Shaun ZhouandClaude Opus 5
af22dc639c
- 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>
315 lines
12 KiB
Python
315 lines
12 KiB
Python
"""同时作用于图像与标注框的几何/像素变换。
|
|
|
|
样本约定为 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)
|