Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
File size: 6,199 Bytes
cd24b4a | 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 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | 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
|