Spaces:
Runtime error
Runtime error
File size: 2,147 Bytes
78856c0 | 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 | """Image Augmentation and Transformation Pipelines using Albumentations."""
from typing import Tuple
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 _make_random_resized_crop(image_size: int, scale: Tuple[float, float] = (0.8, 1.0)):
"""Version-agnostic RandomResizedCrop creator."""
try:
return A.RandomResizedCrop(size=(image_size, image_size), scale=scale)
except Exception:
return A.RandomResizedCrop(height=image_size, width=image_size, scale=scale)
def _make_resize(image_size: int):
"""Version-agnostic Resize creator."""
try:
return A.Resize(size=(image_size, image_size))
except Exception:
return A.Resize(height=image_size, width=image_size)
def _make_coarse_dropout(p: float = 0.3):
"""Version-agnostic CoarseDropout creator."""
try:
return A.CoarseDropout(
num_holes_range=(1, 8),
hole_height_range=(8, 32),
hole_width_range=(8, 32),
p=p,
)
except Exception:
try:
return A.CoarseDropout(max_holes=8, max_height=32, max_width=32, min_holes=1, p=p)
except Exception:
return A.Dropout(p=p)
def get_train_transforms(image_size: int = 380) -> A.Compose:
return A.Compose(
[
_make_random_resized_crop(image_size=image_size, scale=(0.8, 1.0)),
A.HorizontalFlip(p=0.5),
A.VerticalFlip(p=0.5),
A.RandomRotate90(p=0.5),
A.ColorJitter(
brightness=0.2,
contrast=0.2,
saturation=0.1,
hue=0.05,
p=0.5,
),
_make_coarse_dropout(p=0.3),
A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
ToTensorV2(),
]
)
def get_val_transforms(image_size: int = 380) -> A.Compose:
return A.Compose(
[
_make_resize(image_size=image_size),
A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
ToTensorV2(),
]
)
|