Spaces:
Sleeping
Sleeping
File size: 8,708 Bytes
7ee2717 0705c62 7ee2717 | 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 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 | from random import randrange
import torchvision.transforms.functional as TF
from typing import List, Callable, Union
from PIL.Image import Image as PILImage
from training.aug_utils.distortions import *
distortion_groups = {
"blur": ["gaublur", "lensblur", "motionblur"],
"color_distortion": ["colordiff", "colorshift", "colorsat1", "colorsat2"],
"jpeg": ["jpeg2000", "jpeg"],
"noise": ["whitenoise", "whitenoiseCC", "impulsenoise", "multnoise"],
"brightness_change": ["brighten", "darken", "meanshift"],
"spatial_distortion": ["jitter", "noneccpatch", "pixelate", "quantization", "colorblock"],
"sharpness_contrast": ["highsharpen", "lincontrchange", "nonlincontrchange"],
}
distortion_groups_mapping = {
"gaublur": "blur",
"lensblur": "blur",
"motionblur": "blur",
"colordiff": "color_distortion",
"colorshift": "color_distortion",
"colorsat1": "color_distortion",
"colorsat2": "color_distortion",
"jpeg2000": "jpeg",
"jpeg": "jpeg",
"whitenoise": "noise",
"whitenoiseCC": "noise",
"impulsenoise": "noise",
"multnoise": "noise",
"brighten": "brightness_change",
"darken": "brightness_change",
"meanshift": "brightness_change",
"jitter": "spatial_distortion",
"noneccpatch": "spatial_distortion",
"pixelate": "spatial_distortion",
"quantization": "spatial_distortion",
"colorblock": "spatial_distortion",
"highsharpen": "sharpness_contrast",
"lincontrchange": "sharpness_contrast",
"nonlincontrchange": "sharpness_contrast",
}
distortion_range = {
"gaublur": [0.1, 0.5, 1, 2, 5],
"lensblur": [1, 2, 4, 6, 8],
"motionblur": [1, 2, 4, 6, 10],
"colordiff": [1, 3, 6, 8, 12],
"colorshift": [1, 3, 6, 8, 12],
"colorsat1": [0.4, 0.2, 0.1, 0, -0.4],
"colorsat2": [1, 2, 3, 6, 9],
"jpeg2000": [16, 32, 45, 120, 170],
"jpeg": [43, 36, 24, 7, 4],
"whitenoise": [0.001, 0.002, 0.003, 0.005, 0.01],
"whitenoiseCC": [0.0001, 0.0005, 0.001, 0.002, 0.003],
"impulsenoise": [0.001, 0.005, 0.01, 0.02, 0.03],
"multnoise": [0.001, 0.005, 0.01, 0.02, 0.05],
"brighten": [0.1, 0.2, 0.4, 0.7, 1.1],
"darken": [0.05, 0.1, 0.2, 0.4, 0.8],
"meanshift": [0, 0.08, -0.08, 0.15, -0.15],
"jitter": [0.05, 0.1, 0.2, 0.5, 1],
"noneccpatch": [20, 40, 60, 80, 100],
"pixelate": [0.01, 0.05, 0.1, 0.2, 0.5],
"quantization": [20, 16, 13, 10, 7],
"colorblock": [2, 4, 6, 8, 10],
"highsharpen": [1, 2, 3, 6, 12],
"lincontrchange": [0., 0.15, -0.4, 0.3, -0.6],
"nonlincontrchange": [0.4, 0.3, 0.2, 0.1, 0.05],
}
distortion_functions = {
"gaublur": gaussian_blur,
"lensblur": lens_blur,
"motionblur": motion_blur,
"colordiff": color_diffusion,
"colorshift": color_shift,
"colorsat1": color_saturation1,
"colorsat2": color_saturation2,
"jpeg2000": jpeg2000,
"jpeg": jpeg,
"whitenoise": white_noise,
"whitenoiseCC": white_noise_cc,
"impulsenoise": impulse_noise,
"multnoise": multiplicative_noise,
"brighten": brighten,
"darken": darken,
"meanshift": mean_shift,
"jitter": jitter,
"noneccpatch": non_eccentricity_patch,
"pixelate": pixelate,
"quantization": quantization,
"colorblock": color_block,
"highsharpen": high_sharpen,
"lincontrchange": linear_contrast_change,
"nonlincontrchange": non_linear_contrast_change,
}
def distort_images(image: torch.Tensor, distort_functions: list = None, distort_values: list = None,
max_distortions: int = 4, num_levels: int = 5) -> torch.Tensor:
"""
Distorts an image using the distortion composition obtained with the image degradation model proposed in the paper
https://arxiv.org/abs/2310.14918.
Args:
image (Tensor): image to distort
distort_functions (list): list of the distortion functions to apply to the image. If None, the functions are randomly chosen.
distort_values (list): list of the values of the distortion functions to apply to the image. If None, the values are randomly chosen.
max_distortions (int): maximum number of distortions to apply to the image
num_levels (int): number of levels of distortion that can be applied to the image
Returns:
image (Tensor): distorted image
distort_functions (list): list of the distortion functions applied to the image
distort_values (list): list of the values of the distortion functions applied to the image
"""
if distort_functions is None or distort_values is None:
distort_functions, distort_values = get_distortions_composition(max_distortions, num_levels)
for distortion, value in zip(distort_functions, distort_values):
image = distortion(image, value)
image = image.to(torch.float32)
image = torch.clip(image, 0, 1)
return image, distort_functions, distort_values
def get_distortions_composition(max_distortions: int = 7, num_levels: int = 5) -> (List[Callable], List[Union[int, float]]):
"""
Image Degradation model proposed in the paper https://arxiv.org/abs/2310.14918. Returns a randomly assembled ordered
sequence of distortion functions and their values.
Args:
max_distortions (int): maximum number of distortions to apply to the image
num_levels (int): number of levels of distortion that can be applied to the image
Returns:
distort_functions (list): list of the distortion functions to apply to the image
distort_values (list): list of the values of the distortion functions to apply to the image
"""
MEAN = 0
STD = 2.5
num_distortions = random.randint(1, max_distortions)
groups = random.sample(list(distortion_groups.keys()), num_distortions)
distortions = [random.choice(distortion_groups[group]) for group in groups]
distort_functions = [distortion_functions[dist] for dist in distortions]
probabilities = [1 / (STD * np.sqrt(2 * np.pi)) * np.exp(-((i - MEAN) ** 2) / (2 * STD ** 2))
for i in range(num_levels)] # probabilities according to a gaussian distribution
normalized_probabilities = [prob / sum(probabilities)
for prob in probabilities] # normalize probabilities
distort_values = [np.random.choice(distortion_range[dist][:num_levels], p=normalized_probabilities) for dist
in distortions]
return distort_functions, distort_values
def resize_crop(img: PILImage, crop_size: int = 224, downscale_factor: int = 1) -> PILImage:
"""
Resize the image with the desired downscale factor and optionally crop it to the desired size. The crop is randomly
sampled from the image. If crop_size is None, no crop is applied. If the crop is out of bounds, the image is
automatically padded with zeros.
Args:
img (PIL Image): image to resize and crop
crop_size (int): size of the crop. If None, no crop is applied
downscale_factor (int): downscale factor to apply to the image
Returns:
img (PIL Image): resized and/or cropped image
"""
w, h = img.size
if downscale_factor > 1:
img = img.resize((w // downscale_factor, h // downscale_factor))
w, h = img.size
if crop_size is not None:
top = randrange(0, max(1, h - crop_size))
left = randrange(0, max(1, w - crop_size))
img = TF.crop(img, top, left, crop_size, crop_size) # Automatically pad with zeros if the crop is out of bounds
return img
def center_corners_crop(img: PILImage, crop_size: int = 224) -> List[PILImage]:
"""
Return the center crop and the four corners of the image.
Args:
img (PIL.Image): image to crop
crop_size (int): size of each crop
Returns:
crops (List[PIL.Image]): list of the five crops
"""
width, height = img.size
# Calculate the coordinates for the center crop and the four corners
cx = width // 2
cy = height // 2
crops = [
TF.crop(img, cy - crop_size // 2, cx - crop_size // 2, crop_size, crop_size), # Center
TF.crop(img, 0, 0, crop_size, crop_size), # Top-left corner
TF.crop(img, height - crop_size, 0, crop_size, crop_size), # Bottom-left corner
TF.crop(img, 0, width - crop_size, crop_size, crop_size), # Top-right corner
TF.crop(img, height - crop_size, width - crop_size, crop_size, crop_size) # Bottom-right corner
]
return crops
|