naver/PUMP
1
1# Copyright 2022-present NAVER Corp.2# CC BY-NC-SA 4.03# Available only for non-commercial use4 5from pdb import set_trace as bb6import warnings7 8import numpy as np9from PIL import Image, ImageOps10 11import torch12import torch.nn as nn13from torchvision import transforms as tvf14 15from . import transforms_tools as F16from .utils import DatasetWithRng17 18'''19Example command to try out some transformation chain:20 21python -m pytools.transforms --trfs "Scale(384), ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.1), RandomRotation(10), RandomTilting(0.5, 'all'), RandomScale(240,320), RandomCrop(224)"22'''23 24def instanciate_transforms(transforms, use_gpu=False, rng=None, compose=True):25 ''' Instanciate a sequence of transformations.26 27 transforms: (str, list) 28 Comma-separated list of transformations.29 Ex: "Rotate(10), Scale(256)"30 '''31 try:32 transforms = transforms or '[]'33 34 if isinstance(transforms, str):35 if transforms.lstrip()[0] not in '[(': transforms = f'[{transforms}]'36 if compose: transforms = f'Compose({transforms})'37 transforms = eval(transforms)38 39 if isinstance(transforms, list) and transforms and isinstance(transforms[0], str):40 transforms = [eval(trf) for trf in transforms]41 if compose: transforms = Compose(transforms)42 43 if use_gpu and not isinstance(transforms, nn.Module):44 while hasattr(transforms,'transforms') or hasattr(transforms,'transform'): 45 transforms = getattr(transforms,'transforms',getattr(transforms,'transform',None))46 transforms = [trf for trf in transforms if isinstance(trf, nn.Module)]47 transforms = nn.Sequential(*transforms) if compose else nn.ModuleList(transforms)48 49 if transforms and rng: 50 for trf in transforms.transforms: 51 assert hasattr(trf, 'rng'), f"Transformation {trf} has no self.rng"52 trf.rng = rng53 54 if isinstance(transforms, Compose) and len(transforms.transforms) == 1:55 transforms = transforms.transforms[0]56 return transforms57 58 except Exception as e:59 print("\nError: Cannot interpret this transform list: %s\n" % transforms)60 raise e61 62 63 64class Compose (DatasetWithRng):65 def __init__(self, transforms, **rng_seed):66 super().__init__(**rng_seed)67 self.transforms = [self.with_same_rng(trf) for trf in transforms]68 69 def __call__(self, data):70 for trf in self.transforms:71 data = trf(data)72 return data73 74 75class Scale (DatasetWithRng):76 """ Rescale the input PIL.Image to a given size.77 Copied from https://github.com/pytorch in torchvision/transforms/transforms.py78 79 The smallest dimension of the resulting image will be = size.80 81 if largest == True: same behaviour for the largest dimension.82 83 if not can_upscale: don't upscale84 if not can_downscale: don't downscale85 """86 def __init__(self, size, interpolation=Image.BILINEAR, largest=False, 87 can_upscale=True, can_downscale=True, **rng_seed):88 super().__init__(**rng_seed)89 assert isinstance(size, int) or (len(size) == 2)90 self.size = size91 self.interpolation = interpolation92 self.largest = largest93 self.can_upscale = can_upscale94 self.can_downscale = can_downscale95 96 def __repr__(self):97 fmt_str = "RandomScale(%s" % str(self.size)98 if self.largest: fmt_str += ', largest=True'99 if not self.can_upscale: fmt_str += ', can_upscale=False'100 if not self.can_downscale: fmt_str += ', can_downscale=False'101 return fmt_str+')'102 103 def get_params(self, imsize):104 w,h = imsize105 if isinstance(self.size, int):106 cmp = lambda a,b: (a>=b) if self.largest else (a<=b)107 if (cmp(w, h) and w == self.size) or (cmp(h, w) and h == self.size):108 ow, oh = w, h109 elif cmp(w, h):110 ow = self.size111 oh = int(self.size * h / w)112 else:113 oh = self.size114 ow = int(self.size * w / h)115 else:116 ow, oh = self.size117 return ow, oh118 119 def __call__(self, inp):120 img = F.grab(inp,'img')121 w, h = img.size122 123 size2 = ow, oh = self.get_params(img.size)124 125 if size2 != img.size:126 a1, a2 = img.size, size2127 if (self.can_upscale and min(a1) < min(a2)) or (self.can_downscale and min(a1) > min(a2)):128 img = img.resize(size2, self.interpolation)129 130 return F.update(inp, img=img, homography=np.diag((ow/w,oh/h,1)))131 132 133 134class RandomScale (Scale):135 """Rescale the input PIL.Image to a random size.136 Copied from https://github.com/pytorch in torchvision/transforms/transforms.py137 138 Args:139 min_size (int): min size of the smaller edge of the picture.140 max_size (int): max size of the smaller edge of the picture.141 142 ar (float or tuple):143 max change of aspect ratio (width/height).144 145 interpolation (int, optional): Desired interpolation. Default is146 ``PIL.Image.BILINEAR``147 """148 149 def __init__(self, min_size, max_size, ar=1, larger=False,150 can_upscale=False, can_downscale=True, interpolation=Image.BILINEAR):151 Scale.__init__(self, (min_size,max_size), can_upscale=can_upscale, can_downscale=can_downscale, interpolation=interpolation)152 assert type(min_size) == type(max_size), 'min_size and max_size can only be 2 ints or 2 floats'153 assert isinstance(min_size, int) and min_size >= 1 or isinstance(min_size, float) and min_size>0154 assert isinstance(max_size, (int,float)) and min_size <= max_size155 self.min_size = min_size156 self.max_size = max_size157 if type(ar) in (float,int): ar = (min(1/ar,ar),max(1/ar,ar))158 assert 0.2 < ar[0] <= ar[1] < 5159 self.ar = ar160 self.larger = larger161 162 def get_params(self, imsize):163 w,h = imsize164 if isinstance(self.min_size, float): min_size = int(self.min_size*min(w,h) + 0.5)165 if isinstance(self.max_size, float): max_size = int(self.max_size*min(w,h) + 0.5)166 if isinstance(self.min_size, int): min_size = self.min_size167 if isinstance(self.max_size, int): max_size = self.max_size168 169 if not(self.can_upscale) and not(self.larger):170 max_size = min(max_size,min(w,h))171 172 size = int(0.5 + F.rand_log_uniform(self.rng, min_size, max_size))173 if not(self.can_upscale) and self.larger:174 size = min(size, min(w,h))175 176 ar = F.rand_log_uniform(self.rng, *self.ar) # change of aspect ratio177 178 if w < h: # image is taller179 ow = size180 oh = int(0.5 + size * h / w / ar)181 if oh < min_size:182 ow,oh = int(0.5 + ow*float(min_size)/oh),min_size183 else: # image is wider184 oh = size185 ow = int(0.5 + size * w / h * ar)186 if ow < min_size:187 ow,oh = min_size,int(0.5 + oh*float(min_size)/ow)188 189 assert ow >= min_size, 'image too small (width=%d < min_size=%d)' % (ow, min_size)190 assert oh >= min_size, 'image too small (height=%d < min_size=%d)' % (oh, min_size)191 return ow, oh192 193 194 195class RandomCrop (DatasetWithRng):196 """Crop the given PIL Image at a random location.197 Copied from https://github.com/pytorch in torchvision/transforms/transforms.py198 199 Args:200 size (sequence or int): Desired output size of the crop. If size is an201 int instead of sequence like (h, w), a square crop (size, size) is202 made.203 padding (int or sequence, optional): Optional padding on each border204 of the image. Default is 0, i.e no padding. If a sequence of length205 4 is provided, it is used to pad left, top, right, bottom borders206 respectively.207 """208 209 def __init__(self, size, padding=0, **rng_seed):210 super().__init__(**rng_seed)211 if isinstance(size, int):212 self.size = (int(size), int(size))213 else:214 self.size = size215 self.padding = padding216 217 def __repr__(self):218 return "RandomCrop(%s)" % str(self.size)219 220 def get_params(self, img, output_size):221 w, h = img.size222 th, tw = output_size223 assert h >= th and w >= tw, "Image of %dx%d is too small for crop %dx%d" % (w,h,tw,th)224 225 y = self.rng.integers(0, h - th) if h > th else 0226 x = self.rng.integers(0, w - tw) if w > tw else 0227 return x, y, tw, th228 229 def __call__(self, inp):230 img = F.grab(inp,'img')231 232 padl = padt = 0233 if self.padding:234 if F.is_pil_image(img):235 img = ImageOps.expand(img, border=self.padding, fill=0)236 else:237 assert isinstance(img, F.DummyImg)238 img = img.expand(border=self.padding)239 if isinstance(self.padding, int):240 padl = padt = self.padding241 else:242 padl, padt = self.padding[0:2]243 244 i, j, tw, th = self.get_params(img, self.size)245 img = img.crop((i, j, i+tw, j+th))246 247 return F.update(inp, img=img, homography=np.float32(((1,0,padl-i),(0,1,padt-j),(0,0,1))))248 249 250class CenterCrop (RandomCrop):251 """Crops the given PIL Image at the center.252 Copied from https://github.com/pytorch in torchvision/transforms/transforms.py253 254 Args:255 size (sequence or int): Desired output size of the crop. If size is an256 int instead of sequence like (h, w), a square crop (size, size) is257 made.258 """259 @staticmethod260 def get_params(img, output_size):261 w, h = img.size262 th, tw = output_size263 y = int(0.5 +((h - th) / 2.))264 x = int(0.5 +((w - tw) / 2.))265 return x, y, tw, th266 267 268class RandomRotation (DatasetWithRng):269 """Rescale the input PIL.Image to a random size.270 Copied from https://github.com/pytorch in torchvision/transforms/transforms.py271 272 Args:273 degrees (float):274 rotation angle.275 276 interpolation (int, optional): Desired interpolation. Default is277 ``PIL.Image.BILINEAR``278 """279 280 def __init__(self, degrees, interpolation=Image.BILINEAR, **rng_seed):281 super().__init__(**rng_seed)282 self.degrees = degrees283 self.interpolation = interpolation284 285 def __repr__(self):286 return f"RandomRotation({self.degrees})"287 288 def __call__(self, inp):289 img = F.grab(inp,'img')290 w, h = img.size291 292 angle = self.rng.uniform(-self.degrees, self.degrees)293 294 img = img.rotate(angle, resample=self.interpolation)295 w2, h2 = img.size296 297 trf = F.translate(w2/2,h2/2) @ F.rotate(-angle * np.pi/180) @ F.translate(-w/2,-h/2)298 return F.update(inp, img=img, homography=trf)299 300 301class RandomTilting (DatasetWithRng):302 """Apply a random tilting (left, right, up, down) to the input PIL.Image303 Copied from https://github.com/pytorch in torchvision/transforms/transforms.py304 305 Args:306 maginitude (float):307 maximum magnitude of the random skew (value between 0 and 1)308 directions (string):309 tilting directions allowed (all, left, right, up, down)310 examples: "all", "left,right", "up-down-right"311 """312 313 def __init__(self, magnitude, directions='all', **rng_seed):314 super().__init__(**rng_seed)315 self.magnitude = magnitude316 self.directions = directions.lower().replace(',',' ').replace('-',' ')317 318 def __repr__(self):319 return "RandomTilt(%g, '%s')" % (self.magnitude,self.directions)320 321 def __call__(self, inp):322 img = F.grab(inp,'img')323 w, h = img.size324 325 x1,y1,x2,y2 = 0,0,h,w326 original_plane = [(y1, x1), (y2, x1), (y2, x2), (y1, x2)]327 328 max_skew_amount = max(w, h)329 max_skew_amount = int(np.ceil(max_skew_amount * self.magnitude))330 skew_amount = self.rng.integers(1, max_skew_amount)331 332 if self.directions == 'all':333 choices = [0,1,2,3]334 else:335 dirs = ['left', 'right', 'up', 'down']336 choices = []337 for d in self.directions.split():338 try:339 choices.append(dirs.index(d))340 except:341 raise ValueError('Tilting direction %s not recognized' % d)342 343 skew_direction = self.rng.choice(choices)344 345 # print('randomtitlting: ', skew_amount, skew_direction) # to debug random346 347 if skew_direction == 0:348 # Left Tilt349 new_plane = [(y1, x1 - skew_amount), # Top Left350 (y2, x1), # Top Right351 (y2, x2), # Bottom Right352 (y1, x2 + skew_amount)] # Bottom Left353 elif skew_direction == 1:354 # Right Tilt355 new_plane = [(y1, x1), # Top Left356 (y2, x1 - skew_amount), # Top Right357 (y2, x2 + skew_amount), # Bottom Right358 (y1, x2)] # Bottom Left359 elif skew_direction == 2:360 # Forward Tilt361 new_plane = [(y1 - skew_amount, x1), # Top Left362 (y2 + skew_amount, x1), # Top Right363 (y2, x2), # Bottom Right364 (y1, x2)] # Bottom Left365 elif skew_direction == 3:366 # Backward Tilt367 new_plane = [(y1, x1), # Top Left368 (y2, x1), # Top Right369 (y2 + skew_amount, x2), # Bottom Right370 (y1 - skew_amount, x2)] # Bottom Left371 372 # To calculate the coefficients required by PIL for the perspective skew,373 # see the following Stack Overflow discussion: https://goo.gl/sSgJdj374 homography = F.homography_from_4pts(original_plane, new_plane)375 img = img.transform(img.size, Image.PERSPECTIVE, homography, resample=Image.BICUBIC)376 377 homography = np.linalg.pinv(np.float32(homography+(1,)).reshape(3,3))378 return F.update(inp, img=img, homography=homography)379 380 381RandomHomography = RandomTilt = RandomTilting # redefinition382 383 384class Homography(object):385 """Apply a known tilting to an image386 """387 def __init__(self, *homography):388 assert len(homography) == 8389 self.homography = homography390 391 def __call__(self, inp):392 img = F.grab(inp, 'img')393 homography = self.homography394 395 img = img.transform(img.size, Image.PERSPECTIVE, homography, resample=Image.BICUBIC)396 397 homography = np.linalg.pinv(np.float32(list(homography)+[1]).reshape(3,3))398 return F.update(inp, img=img, homography=homography)399 400 401 402class StillTransform (DatasetWithRng):403 """ Takes and return an image, without changing its shape or geometry.404 """405 def _transform(self, img):406 raise NotImplementedError()407 408 def __call__(self, inp):409 img = F.grab(inp,'img')410 411 # transform the image (size should not change)412 try:413 img = self._transform(img)414 except TypeError:415 pass416 417 return F.update(inp, img=img)418 419 420 421class PixelNoise (StillTransform):422 """ Takes an image, and add random white noise.423 """424 def __init__(self, ampl=20, **rng_seed):425 super().__init__(**rng_seed)426 assert 0 <= ampl < 255427 self.ampl = ampl428 429 def __repr__(self):430 return "PixelNoise(%g)" % self.ampl431 432 def _transform(self, img):433 img = np.float32(img)434 img += self.rng.uniform(0.5-self.ampl/2, 0.5+self.ampl/2, size=img.shape)435 return Image.fromarray(np.uint8(img.clip(0,255)))436 437 438 439class ColorJitter (StillTransform):440 """Randomly change the brightness, contrast and saturation of an image.441 Copied from https://github.com/pytorch in torchvision/transforms/transforms.py442 443 Args:444 brightness (float): How much to jitter brightness. brightness_factor445 is chosen uniformly from [max(0, 1 - brightness), 1 + brightness].446 contrast (float): How much to jitter contrast. contrast_factor447 is chosen uniformly from [max(0, 1 - contrast), 1 + contrast].448 saturation (float): How much to jitter saturation. saturation_factor449 is chosen uniformly from [max(0, 1 - saturation), 1 + saturation].450 hue(float): How much to jitter hue. hue_factor is chosen uniformly from451 [-hue, hue]. Should be >=0 and <= 0.5.452 """453 def __init__(self, brightness=0, contrast=0, saturation=0, hue=0):454 self.brightness = brightness455 self.contrast = contrast456 self.saturation = saturation457 self.hue = hue458 459 def __repr__(self):460 return "ColorJitter(%g,%g,%g,%g)" % (461 self.brightness, self.contrast, self.saturation, self.hue)462 463 def get_params(self, brightness, contrast, saturation, hue):464 """Get a randomized transform to be applied on image.465 Arguments are same as that of __init__.466 Returns:467 Transform which randomly adjusts brightness, contrast and468 saturation in a random order.469 """470 transforms = []471 if brightness > 0:472 brightness_factor = self.rng.uniform(max(0, 1 - brightness), 1 + brightness)473 transforms.append(tvf.Lambda(lambda img: F.adjust_brightness(img, brightness_factor)))474 475 if contrast > 0:476 contrast_factor = self.rng.uniform(max(0, 1 - contrast), 1 + contrast)477 transforms.append(tvf.Lambda(lambda img: F.adjust_contrast(img, contrast_factor)))478 479 if saturation > 0:480 saturation_factor = self.rng.uniform(max(0, 1 - saturation), 1 + saturation)481 transforms.append(tvf.Lambda(lambda img: F.adjust_saturation(img, saturation_factor)))482 483 if hue > 0:484 hue_factor = self.rng.uniform(-hue, hue)485 transforms.append(tvf.Lambda(lambda img: F.adjust_hue(img, hue_factor)))486 487 # print('colorjitter: ', brightness_factor, contrast_factor, saturation_factor, hue_factor) # to debug random seed488 self.rng.shuffle(transforms)489 transform = tvf.Compose(transforms)490 return transform491 492 def _transform(self, img):493 transform = self.get_params(self.brightness, self.contrast, self.saturation, self.hue)494 return transform(img)495 496 497def pil_loader(path, mode='RGB'):498 with warnings.catch_warnings():499 warnings.simplefilter("ignore")500 # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)501 with (path if hasattr(path,'read') else open(path, 'rb')) as f:502 img = Image.open(f)503 return img.convert(mode)504 505def torchvision_loader(path, mode='RGB'):506 from torchvision.io import read_file, decode_image, read_image, image507 return read_image(getattr(path,'name',path), mode=getattr(image.ImageReadMode,mode))508 509 510 511if __name__ == '__main__':512 from matplotlib import pyplot as pl513 import argparse514 515 parser = argparse.ArgumentParser("Script to try out and visualize transformations")516 parser.add_argument('--img', type=str, default='imgs/test.png', help='input image')517 parser.add_argument('--trfs', type=str, required=True, help='list of transformations')518 parser.add_argument('--layout', type=int, nargs=2, default=(3,3), help='nb of rows,cols')519 args = parser.parse_args()520 521 img = dict(img=pil_loader(args.img))522 523 trfs = instanciate_transforms(args.trfs)524 525 pl.subplots_adjust(0,0,1,1)526 nr,nc = args.layout527 528 while True:529 t0 = now()530 imgs2 = [trfs(img) for _ in range(nr*nc)]531 532 for j in range(nr):533 for i in range(nc):534 pl.subplot(nr,nc,i+j*nc+1)535 img2 = img if i==j==0 else imgs2.pop() #trfs(img)536 img2 = img2['img']537 pl.imshow(img2)538 pl.xlabel("%d x %d" % img2.size)539 print(f'Took {now() - t0:.2f} seconds')540 pl.show()541