| import cv2 | |
| import numpy as np | |
| import albumentations as albu | |
| from albumentations.augmentations.geometric.resize import LongestMaxSize | |
| def round_pixel_dim(dimension: float) -> int: | |
| """Rounds pixel dimensions consistently.""" | |
| if abs(round(dimension) - dimension) == 0.5: | |
| return int(2.0 * round(dimension / 2.0)) | |
| return int(round(dimension)) | |
| def resize_with_padding(image, target_size, stride=32, interpolation=cv2.INTER_LINEAR): | |
| """Resizes image maintaining aspect ratio and ensures dimensions are stride-compatible.""" | |
| height, width = image.shape[:2] | |
| max_dimension = max(height, width) | |
| if ((height % stride == 0) and (width % stride == 0) and | |
| (max_dimension <= target_size)): | |
| return image | |
| scale = target_size / float(max(width, height)) | |
| new_dims = tuple(round_pixel_dim(dim * scale) for dim in (height, width)) | |
| new_height, new_width = new_dims | |
| new_height = ((new_height // stride + 1) * stride | |
| if new_height % stride != 0 else new_height) | |
| new_width = ((new_width // stride + 1) * stride | |
| if new_width % stride != 0 else new_width) | |
| return cv2.resize(image, (new_width, new_height), interpolation=interpolation) | |
| class PaddedResize(LongestMaxSize): | |
| def apply(self, img: np.ndarray, target_size: int = 1024, | |
| interpolation: int = cv2.INTER_LINEAR, **params) -> np.ndarray: | |
| return resize_with_padding(img, target_size=target_size, interpolation=interpolation) | |
| def get_training_augmentations(width=768, height=576): | |
| """Configures training-time augmentations.""" | |
| target_size = max([width, height]) | |
| transforms = [ | |
| albu.HorizontalFlip(p=0.5), | |
| albu.ShiftScaleRotate( | |
| scale_limit=0.5, rotate_limit=90, shift_limit=0.1, p=0.5, border_mode=0), | |
| albu.PadIfNeeded(min_height=target_size, min_width=target_size, always_apply=True), | |
| albu.RandomCrop(height=target_size, width=target_size, always_apply=True), | |
| albu.GaussNoise(p=0.2), | |
| albu.Perspective(p=0.2), | |
| albu.OneOf([albu.CLAHE(p=1), albu.RandomGamma(p=1)], p=0.33), | |
| albu.OneOf([ | |
| albu.Sharpen(p=1), | |
| albu.Blur(blur_limit=3, p=1), | |
| albu.MotionBlur(blur_limit=3, p=1)], p=0.33), | |
| albu.OneOf([ | |
| albu.RandomBrightnessContrast(p=1), | |
| albu.HueSaturationValue(p=1)], p=0.33), | |
| ] | |
| return albu.Compose(transforms) | |
| def get_validation_augmentations(width=1920, height=1440, fixed_size=True): | |
| """Configures validation/inference-time augmentations.""" | |
| if fixed_size: | |
| transforms = [albu.Resize(height=height, width=width, always_apply=True)] | |
| return albu.Compose(transforms) | |
| target_size = max(width, height) | |
| transforms = [PaddedResize(max_size=target_size, always_apply=True)] | |
| return albu.Compose(transforms) | |
| def convert_to_tensor(x, **kwargs): | |
| """Converts image array to PyTorch tensor format.""" | |
| if x.ndim == 2: | |
| x = np.expand_dims(x, axis=-1) | |
| return x.transpose(2, 0, 1).astype('float32') | |
| def get_preprocessing_pipeline(preprocessing_fn): | |
| """Builds preprocessing pipeline including normalization and tensor conversion.""" | |
| transforms = [ | |
| albu.Lambda(image=preprocessing_fn), | |
| albu.Lambda(image=convert_to_tensor, mask=convert_to_tensor), | |
| ] | |
| return albu.Compose(transforms) |