Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
| from typing import Callable | |
| import torchvision.transforms.v2 as v2 | |
| from torchvision import tv_tensors | |
| import torch | |
| import random | |
| class RandomMask: | |
| def __init__(self, ratio=(0.0, 0.75), patch_size=16, p=0.5): | |
| self.ratio = ratio | |
| self.patch_size = patch_size | |
| self.p = p | |
| def __call__(self, tensor): | |
| if random.random() > self.p: | |
| return tensor | |
| _, height, width = tensor.shape | |
| ratio = random.uniform(*self.ratio) | |
| num_masks = int((height * width * ratio) / (self.patch_size ** 2)) | |
| mask = torch.ones((height, width), dtype=tensor.dtype, device=tensor.device) | |
| max_top = max((height // self.patch_size) - 1, 0) | |
| max_left = max((width // self.patch_size) - 1, 0) | |
| positions = set() | |
| while len(positions) < num_masks: | |
| top = random.randint(0, max_top) * self.patch_size | |
| left = random.randint(0, max_left) * self.patch_size | |
| positions.add((top, left)) | |
| for top, left in positions: | |
| mask[top:top + self.patch_size, left:left + self.patch_size] = 0 | |
| return tensor * mask.expand_as(tensor) | |
| class AIGITransformations: | |
| def __init__(self, transform_list: list[Callable]): | |
| self.transform_list = transform_list | |
| self.transforms = [self._create_transform(transform) for transform in self.transform_list] | |
| def _create_transform(self, transform_name: str): | |
| if transform_name == "resize_224": | |
| return v2.Resize((224, 224)) | |
| elif transform_name == "resize_short_256": | |
| return v2.Resize(256) | |
| elif transform_name == "resize_short_384": | |
| return v2.Resize(384) | |
| elif transform_name == "resize_short_448": | |
| return v2.Resize(448) | |
| elif transform_name == "center_crop_224": | |
| return v2.CenterCrop(224) | |
| elif transform_name == "center_crop_256": | |
| return v2.CenterCrop(256) | |
| elif transform_name == "center_crop_384": | |
| return v2.CenterCrop(384) | |
| elif transform_name == "center_crop_448": | |
| return v2.CenterCrop(448) | |
| elif transform_name == "to_tensor": | |
| return v2.ToTensor() | |
| elif transform_name == "to_dtype": | |
| return v2.ToDtype(torch.float32, scale=True) | |
| elif transform_name == "gaussian_blur": | |
| return v2.RandomApply( | |
| [v2.GaussianBlur(kernel_size=11, sigma=(1e-3, 3.0))], p=0.5 | |
| ) | |
| elif transform_name == "random_horizontal_flip": | |
| return v2.RandomHorizontalFlip(p=0.5) | |
| elif transform_name == "random_vertical_flip": | |
| return v2.RandomVerticalFlip(p=0.5) | |
| elif transform_name == "jpeg_compression": | |
| return v2.RandomApply([v2.JPEG(quality=(30, 100))], p=0.5) | |
| elif transform_name == "random_crop_224": | |
| return v2.RandomCrop(size=(224, 224)) | |
| elif transform_name == "random_crop_256": | |
| return v2.RandomCrop(size=(256, 256)) | |
| elif transform_name == "random_crop_256_pad": | |
| return v2.RandomCrop(size=(256, 256), pad_if_needed=True) | |
| elif transform_name == "random_crop_384": | |
| return v2.RandomCrop(size=(384, 384)) | |
| elif transform_name == "random_crop_448": | |
| return v2.RandomCrop(size=(448, 448)) | |
| elif transform_name == "random_crop_384_pad": | |
| return v2.RandomCrop(size=(384, 384), pad_if_needed=True) | |
| elif transform_name == "pad_random_crop_224": # PGC: pad-then-crop 224 (official PadRandomCrop) | |
| return v2.RandomCrop(size=(224, 224), pad_if_needed=True) | |
| elif transform_name == "random_rotation_180": | |
| return v2.RandomRotation(180) | |
| elif transform_name == "color_jitter_safe": | |
| return v2.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5) | |
| elif transform_name == "color_jitter_pgc": # PGC: official RandomApply(ColorJitter, p=0.2) | |
| return v2.RandomApply([v2.ColorJitter(brightness=0.15, contrast=0.15, saturation=0.15, hue=0.05)], p=0.2) | |
| elif transform_name == "random_mask_safe": | |
| return RandomMask(ratio=(0.0, 0.75), patch_size=16, p=0.5) | |
| elif transform_name == "aide_five_view": | |
| # AIDE native-resolution 5-view preprocessing: PIL/tensor -> [5,3,256,256] (DCT patch | |
| # selection at native res, then resize the selected patches). Must be the LAST transform. | |
| from models.image.aide.official_arch import AIDEFiveView | |
| return AIDEFiveView() | |
| elif transform_name == "aide_five_view_train": | |
| # Same, but caps the long side at 1024 before the DCT to bound dataloader-worker memory | |
| # on large reals (training only; inference uses the uncapped aide_five_view). | |
| from models.image.aide.official_arch import AIDEFiveView | |
| return AIDEFiveView(max_side=1024) | |
| elif transform_name == "normalize_dinov2": | |
| return v2.Normalize( | |
| mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225] | |
| ) | |
| elif transform_name == "normalize_clip": | |
| return v2.Normalize( | |
| mean=[0.48145466, 0.4578275, 0.40821073], | |
| std=[0.26862954, 0.26130258, 0.27577711], | |
| ) | |
| elif transform_name == "normalize_dinov3": | |
| return v2.Normalize( | |
| mean=(0.430, 0.411, 0.296), | |
| std=(0.213, 0.156, 0.143), | |
| ) | |
| else: | |
| raise ValueError(f"Invalid transform name: {transform_name}") | |
| def __call__(self, image): | |
| # A 4D [T, C, H, W] tensor is a video clip. Wrap it as tv_tensors.Video so the v2 | |
| # transforms (crops, flips, normalize, ...) apply with the SAME sampled parameters to | |
| # every frame (temporal consistency) instead of treating the frames as an independent | |
| # batch. Single images ([C, H, W] tensors or PIL) pass through unchanged. | |
| if isinstance(image, torch.Tensor) and image.ndim == 4: | |
| image = tv_tensors.Video(image) | |
| for transform in self.transforms: | |
| image = transform(image) | |
| return image | |