File size: 2,916 Bytes
b8fae22 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | """Albumentations-based transform builder (the conventional-augmentation tier).
Key correctness guarantees for segmentation:
* masks always use NEAREST interpolation (integer class ids never blended);
every geometric transform sets mask_interpolation=cv2.INTER_NEAREST.
* image and mask receive the SAME random spatial parameters (Albumentations
applies one transform jointly to image= and mask=).
`aug` presets: none (resize+normalize only) | standard | strong.
Returns a callable: (image HWC uint8, mask HW int) -> (FloatTensor[C,H,W], LongTensor[H,W]).
"""
from __future__ import annotations
from typing import Tuple
import cv2
import numpy as np
import torch
import albumentations as A
from albumentations.pytorch import ToTensorV2
_IMAGENET_MEAN = (0.485, 0.456, 0.406)
_IMAGENET_STD = (0.229, 0.224, 0.225)
def _normalize(in_channels: int, normalize: str) -> A.Normalize:
if normalize == "none":
mean = (0.0,) * in_channels
std = (1.0,) * in_channels
elif normalize == "imagenet" and in_channels == 3:
mean, std = _IMAGENET_MEAN, _IMAGENET_STD
else: # auto
if in_channels == 3:
mean, std = _IMAGENET_MEAN, _IMAGENET_STD
else:
mean, std = (0.5,) * in_channels, (0.5,) * in_channels
return A.Normalize(mean=mean, std=std, max_pixel_value=255.0)
def build_transform(img_size: int, in_channels: int, train: bool,
aug: str = "standard", normalize: str = "auto"):
N = cv2.INTER_NEAREST
L = cv2.INTER_LINEAR
ops = []
if train and aug != "none":
ops += [
A.Resize(img_size, img_size, interpolation=L, mask_interpolation=N),
A.HorizontalFlip(p=0.5),
A.VerticalFlip(p=0.5),
A.Affine(scale=(0.9, 1.1), translate_percent=(0.0, 0.05),
rotate=(-15, 15), interpolation=L, mask_interpolation=N, p=0.5),
]
if aug == "strong":
ops += [
A.ElasticTransform(alpha=30, sigma=6, interpolation=L,
mask_interpolation=N, p=0.3),
A.GridDistortion(num_steps=5, distort_limit=0.2, interpolation=L,
mask_interpolation=N, p=0.3),
A.RandomBrightnessContrast(p=0.5),
A.GaussNoise(p=0.2),
]
if in_channels == 3:
ops.append(A.CLAHE(p=0.2))
else:
ops.append(A.Resize(img_size, img_size, interpolation=L, mask_interpolation=N))
ops += [_normalize(in_channels, normalize), ToTensorV2()]
compose = A.Compose(ops)
def _apply(image: np.ndarray, mask: np.ndarray) -> Tuple[torch.Tensor, torch.Tensor]:
out = compose(image=image, mask=mask)
img = out["image"].float() # C,H,W
msk = out["mask"].long() # H,W
return img, msk
return _apply
|