File size: 5,692 Bytes
d9bb75c | 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 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | # Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This software may be used and distributed in accordance with
# the terms of the DINOv3 License Agreement.
import logging
from typing import Sequence
import torch
from torchvision.transforms import v2
logger = logging.getLogger("dinov3")
def make_interpolation_mode(mode_str: str) -> v2.InterpolationMode:
return {mode.value: mode for mode in v2.InterpolationMode}[mode_str]
class GaussianBlur(v2.RandomApply):
"""
Apply Gaussian Blur to the PIL image.
"""
def __init__(self, *, p: float = 0.5, radius_min: float = 0.1, radius_max: float = 2.0):
# NOTE: torchvision is applying 1 - probability to return the original image
keep_p = 1 - p
transform = v2.GaussianBlur(kernel_size=9, sigma=(radius_min, radius_max))
super().__init__(transforms=[transform], p=keep_p)
# Use timm's names
IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
CROP_DEFAULT_SIZE = 224
RESIZE_DEFAULT_SIZE = int(256 * CROP_DEFAULT_SIZE / 224)
def make_normalize_transform(
mean: Sequence[float] = IMAGENET_DEFAULT_MEAN,
std: Sequence[float] = IMAGENET_DEFAULT_STD,
) -> v2.Normalize:
return v2.Normalize(mean=mean, std=std)
def make_base_transform(
mean: Sequence[float] = IMAGENET_DEFAULT_MEAN,
std: Sequence[float] = IMAGENET_DEFAULT_STD,
) -> v2.Normalize:
return v2.Compose(
[
v2.ToDtype(torch.float32, scale=True),
make_normalize_transform(mean=mean, std=std),
]
)
# This roughly matches torchvision's preset for classification training:
# https://github.com/pytorch/vision/blob/main/references/classification/presets.py#L6-L44
def make_classification_train_transform(
*,
crop_size: int = CROP_DEFAULT_SIZE,
interpolation=v2.InterpolationMode.BICUBIC,
hflip_prob: float = 0.5,
mean: Sequence[float] = IMAGENET_DEFAULT_MEAN,
std: Sequence[float] = IMAGENET_DEFAULT_STD,
):
transforms_list = [v2.ToImage(), v2.RandomResizedCrop(crop_size, interpolation=interpolation)]
if hflip_prob > 0.0:
transforms_list.append(v2.RandomHorizontalFlip(hflip_prob))
transforms_list.append(make_base_transform(mean, std))
transform = v2.Compose(transforms_list)
logger.info(f"Built classification train transform\n{transform}")
return transform
def make_resize_transform(
*,
resize_size: int,
resize_square: bool = False,
resize_large_side: bool = False, # Set the larger side to resize_size instead of the smaller
interpolation: v2.InterpolationMode = v2.InterpolationMode.BICUBIC,
):
assert not (resize_square and resize_large_side), "These two options can not be set together"
if resize_square:
logger.info("resizing image as a square")
size = (resize_size, resize_size)
transform = v2.Resize(size=size, interpolation=interpolation)
return transform
elif resize_large_side:
logger.info("resizing based on large side")
transform = v2.Resize(size=None, max_size=resize_size, interpolation=interpolation)
return transform
else:
transform = v2.Resize(resize_size, interpolation=interpolation)
return transform
# Derived from make_classification_eval_transform() with more control over resize and crop
def make_eval_transform(
*,
resize_size: int = RESIZE_DEFAULT_SIZE,
crop_size: int = CROP_DEFAULT_SIZE,
resize_square: bool = False,
resize_large_side: bool = False, # Set the larger side to resize_size instead of the smaller
interpolation: v2.InterpolationMode = v2.InterpolationMode.BICUBIC,
mean: Sequence[float] = IMAGENET_DEFAULT_MEAN,
std: Sequence[float] = IMAGENET_DEFAULT_STD,
) -> v2.Compose:
transforms_list = [v2.ToImage()]
resize_transform = make_resize_transform(
resize_size=resize_size,
resize_square=resize_square,
resize_large_side=resize_large_side,
interpolation=interpolation,
)
transforms_list.append(resize_transform)
if crop_size:
transforms_list.append(v2.CenterCrop(crop_size))
transforms_list.append(make_base_transform(mean, std))
transform = v2.Compose(transforms_list)
logger.info(f"Built eval transform\n{transform}")
return transform
# This matches (roughly) torchvision's preset for classification evaluation:
# https://github.com/pytorch/vision/blob/main/references/classification/presets.py#L47-L69
def make_classification_eval_transform(
*,
resize_size: int = RESIZE_DEFAULT_SIZE,
crop_size: int = CROP_DEFAULT_SIZE,
interpolation=v2.InterpolationMode.BICUBIC,
mean: Sequence[float] = IMAGENET_DEFAULT_MEAN,
std: Sequence[float] = IMAGENET_DEFAULT_STD,
) -> v2.Compose:
return make_eval_transform(
resize_size=resize_size,
crop_size=crop_size,
interpolation=interpolation,
mean=mean,
std=std,
resize_square=False,
resize_large_side=False,
)
def voc2007_classification_target_transform(label, n_categories=20):
one_hot = torch.zeros(n_categories, dtype=int)
for instance in label.instances:
one_hot[instance.category_id] = True
return one_hot
def imaterialist_classification_target_transform(label, n_categories=294):
one_hot = torch.zeros(n_categories, dtype=int)
one_hot[label.attributes] = True
return one_hot
def get_target_transform(dataset_str):
if "VOC2007" in dataset_str:
return voc2007_classification_target_transform
elif "IMaterialist" in dataset_str:
return imaterialist_classification_target_transform
return None
|