Spaces:
Running on Zero
Running on Zero
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| # | |
| from logging import getLogger | |
| from math import floor | |
| from PIL import ImageFilter | |
| import torch | |
| import torchvision.transforms as transforms | |
| import torchvision.transforms.functional as F | |
| _GLOBAL_SEED = 0 | |
| logger = getLogger() | |
| def make_transforms( | |
| crop_size=224, | |
| crop_scale=(0.3, 1.0), | |
| color_jitter=1.0, | |
| random_resize_crop=True, | |
| horizontal_flip=False, | |
| color_distortion=False, | |
| gaussian_blur=False, | |
| normalization=((0.485, 0.456, 0.406), | |
| (0.229, 0.224, 0.225)), | |
| preserve_aspect_ratio=True, | |
| max_patches=14*14, | |
| patch_size=16 | |
| ): | |
| logger.info('making data transforms') | |
| def get_color_distortion(s=1.0): | |
| # s is the strength of color distortion. | |
| color_jitter = transforms.ColorJitter(0.8*s, 0.8*s, 0.8*s, 0.2*s) | |
| rnd_color_jitter = transforms.RandomApply([color_jitter], p=0.8) | |
| rnd_gray = transforms.RandomGrayscale(p=0.2) | |
| color_distort = transforms.Compose([ | |
| rnd_color_jitter, | |
| rnd_gray]) | |
| return color_distort | |
| transform_list = [] | |
| if preserve_aspect_ratio: | |
| if random_resize_crop: | |
| transform_list.append(RandomResizedCropARP(scale=crop_scale)) | |
| transform_list.append( | |
| ResizeToFixedPatches(max_patches=max_patches, patch_size=patch_size) | |
| ) | |
| else: | |
| if random_resize_crop: | |
| transform_list.append(transforms.RandomResizedCrop(crop_size, scale=crop_scale)) | |
| else: | |
| transform_list.append(transforms.Resize((crop_size, crop_size))) | |
| if horizontal_flip: | |
| transform_list.append(transforms.RandomHorizontalFlip()) | |
| if color_distortion: | |
| transform_list.append(get_color_distortion(s=color_jitter)) | |
| if gaussian_blur: | |
| transform_list.append(GaussianBlur(p=0.5)) | |
| mean, std = normalization | |
| transform_list.extend([transforms.ToTensor(), transforms.Normalize(mean, std)]) | |
| transform = transforms.Compose(transform_list) | |
| return transform | |
| class GaussianBlur(object): | |
| """ | |
| Apply Gaussian blur with a given probability. | |
| :param p: Probability of applying blur | |
| :param radius_min: Minimum blur radius | |
| :param radius_max: Maximum blur radius | |
| """ | |
| def __init__(self, p=0.5, radius_min=0.1, radius_max=2.): | |
| self.prob = p | |
| self.radius_min = radius_min | |
| self.radius_max = radius_max | |
| def __call__(self, img): | |
| if torch.bernoulli(torch.tensor(self.prob)) == 0: | |
| return img | |
| radius = self.radius_min + torch.rand(1) * (self.radius_max - self.radius_min) | |
| return img.filter(ImageFilter.GaussianBlur(radius=radius)) | |
| class ResizeToFixedPatches(object): | |
| """ | |
| Resize an image so that its patch grid fits within max_patches. | |
| :param max_patches: Maximum number of patches (H*W) | |
| :param patch_size: Patch size in pixels | |
| """ | |
| def __init__(self, max_patches=14*14, patch_size=16): | |
| self.max_patches = int(max_patches) | |
| self.patch_size = int(patch_size) | |
| def __call__(self, img): | |
| image_width, image_height = img.size | |
| scale = (self.max_patches * (self.patch_size / image_height) * (self.patch_size / image_width)) ** 0.5 | |
| scale += 1e-6 | |
| num_feasible_rows = max(min(floor(scale * image_height / self.patch_size), self.max_patches), 1) | |
| num_feasible_cols = max(min(floor(scale * image_width / self.patch_size), self.max_patches), 1) | |
| resized_height = max(int(num_feasible_rows * self.patch_size), 1) | |
| resized_width = max(int(num_feasible_cols * self.patch_size), 1) | |
| return F.resize( | |
| img, | |
| (resized_height, resized_width), | |
| interpolation=transforms.InterpolationMode.BICUBIC, | |
| antialias=True | |
| ) | |
| class RandomResizedCropARP(object): | |
| """ | |
| Random resized crop that preserves the original aspect ratio. | |
| Uses torchvision's RandomResizedCrop sampling logic with ratio fixed | |
| to the input image's aspect ratio. | |
| :param scale: Range of area scaling for the crop | |
| """ | |
| def __init__(self, scale=(0.3, 1.0)): | |
| self.scale = scale | |
| def __call__(self, img): | |
| image_width, image_height = img.size | |
| if image_width <= 0 or image_height <= 0: | |
| return img | |
| ratio = image_width / image_height | |
| i, j, h, w = transforms.RandomResizedCrop.get_params( | |
| img, | |
| scale=self.scale, | |
| ratio=(ratio, ratio) | |
| ) | |
| return F.crop(img, i, j, h, w) | |