Spaces:
Runtime error
Runtime error
| """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(), | |
| ] | |
| ) | |