ICML2022/OFA
18
1import random2 3import torch4import torchvision.transforms as T5import torchvision.transforms.functional as F6import numpy as np7from PIL import Image8 9 10def crop(image, target, region, delete=True):11 cropped_image = F.crop(image, *region)12 13 target = target.copy()14 i, j, h, w = region15 16 # should we do something wrt the original size?17 target["size"] = torch.tensor([h, w])18 19 fields = ["labels", "area"]20 21 if "boxes" in target:22 boxes = target["boxes"]23 max_size = torch.as_tensor([w, h], dtype=torch.float32)24 cropped_boxes = boxes - torch.as_tensor([j, i, j, i])25 cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size)26 cropped_boxes = cropped_boxes.clamp(min=0)27 area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1)28 target["boxes"] = cropped_boxes.reshape(-1, 4)29 target["area"] = area30 fields.append("boxes")31 32 if "polygons" in target:33 polygons = target["polygons"]34 num_polygons = polygons.shape[0]35 max_size = torch.as_tensor([w, h], dtype=torch.float32)36 start_coord = torch.cat([torch.tensor([j, i], dtype=torch.float32)37 for _ in range(polygons.shape[1] // 2)], dim=0)38 cropped_boxes = polygons - start_coord39 cropped_boxes = torch.min(cropped_boxes.reshape(num_polygons, -1, 2), max_size)40 cropped_boxes = cropped_boxes.clamp(min=0)41 target["polygons"] = cropped_boxes.reshape(num_polygons, -1)42 fields.append("polygons")43 44 if "masks" in target:45 # FIXME should we update the area here if there are no boxes?46 target['masks'] = target['masks'][:, i:i + h, j:j + w]47 fields.append("masks")48 49 # remove elements for which the boxes or masks that have zero area50 if delete and ("boxes" in target or "masks" in target):51 # favor boxes selection when defining which elements to keep52 # this is compatible with previous implementation53 if "boxes" in target:54 cropped_boxes = target['boxes'].reshape(-1, 2, 2)55 keep = torch.all(cropped_boxes[:, 1, :] > cropped_boxes[:, 0, :], dim=1)56 else:57 keep = target['masks'].flatten(1).any(1)58 59 for field in fields:60 target[field] = target[field][keep.tolist()]61 62 return cropped_image, target63 64 65def hflip(image, target):66 flipped_image = F.hflip(image)67 68 w, h = image.size69 70 target = target.copy()71 if "boxes" in target:72 boxes = target["boxes"]73 boxes = boxes[:, [2, 1, 0, 3]] * torch.as_tensor([-1, 1, -1, 1]) + torch.as_tensor([w, 0, w, 0])74 target["boxes"] = boxes75 76 if "polygons" in target:77 polygons = target["polygons"]78 num_polygons = polygons.shape[0]79 polygons = polygons.reshape(num_polygons, -1, 2) * torch.as_tensor([-1, 1]) + torch.as_tensor([w, 0])80 target["polygons"] = polygons81 82 if "masks" in target:83 target['masks'] = target['masks'].flip(-1)84 85 return flipped_image, target86 87 88def resize(image, target, size, max_size=None):89 # size can be min_size (scalar) or (w, h) tuple90 91 def get_size_with_aspect_ratio(image_size, size, max_size=None):92 w, h = image_size93 94 if (w <= h and w == size) or (h <= w and h == size):95 if max_size is not None:96 max_size = int(max_size)97 h = min(h, max_size)98 w = min(w, max_size)99 return (h, w)100 101 if w < h:102 ow = size103 oh = int(size * h / w)104 else:105 oh = size106 ow = int(size * w / h)107 108 if max_size is not None:109 max_size = int(max_size)110 oh = min(oh, max_size)111 ow = min(ow, max_size)112 113 return (oh, ow)114 115 def get_size(image_size, size, max_size=None):116 if isinstance(size, (list, tuple)):117 return size[::-1]118 else:119 return get_size_with_aspect_ratio(image_size, size, max_size)120 121 size = get_size(image.size, size, max_size)122 rescaled_image = F.resize(image, size, interpolation=Image.BICUBIC)123 124 if target is None:125 return rescaled_image126 127 ratios = tuple(float(s) / float(s_orig) for s, s_orig in zip(rescaled_image.size, image.size))128 ratio_width, ratio_height = ratios129 130 target = target.copy()131 if "boxes" in target:132 boxes = target["boxes"]133 scaled_boxes = boxes * torch.as_tensor([ratio_width, ratio_height, ratio_width, ratio_height])134 target["boxes"] = scaled_boxes135 136 if "polygons" in target:137 polygons = target["polygons"]138 scaled_ratio = torch.cat([torch.tensor([ratio_width, ratio_height])139 for _ in range(polygons.shape[1] // 2)], dim=0)140 scaled_polygons = polygons * scaled_ratio141 target["polygons"] = scaled_polygons142 143 if "area" in target:144 area = target["area"]145 scaled_area = area * (ratio_width * ratio_height)146 target["area"] = scaled_area147 148 h, w = size149 target["size"] = torch.tensor([h, w])150 151 if "masks" in target:152 assert False153 # target['masks'] = interpolate(154 # target['masks'][:, None].float(), size, mode="nearest")[:, 0] > 0.5155 156 return rescaled_image, target157 158 159class CenterCrop(object):160 def __init__(self, size):161 self.size = size162 163 def __call__(self, img, target):164 image_width, image_height = img.size165 crop_height, crop_width = self.size166 crop_top = int(round((image_height - crop_height) / 2.))167 crop_left = int(round((image_width - crop_width) / 2.))168 return crop(img, target, (crop_top, crop_left, crop_height, crop_width))169 170 171class ObjectCenterCrop(object):172 def __init__(self, size):173 self.size = size174 175 def __call__(self, img, target):176 image_width, image_height = img.size177 crop_height, crop_width = self.size178 179 x0 = float(target['boxes'][0][0])180 y0 = float(target['boxes'][0][1])181 x1 = float(target['boxes'][0][2])182 y1 = float(target['boxes'][0][3])183 184 center_x = (x0 + x1) / 2185 center_y = (y0 + y1) / 2186 crop_left = max(center_x-crop_width/2 + min(image_width-center_x-crop_width/2, 0), 0)187 crop_top = max(center_y-crop_height/2 + min(image_height-center_y-crop_height/2, 0), 0)188 189 return crop(img, target, (crop_top, crop_left, crop_height, crop_width), delete=False)190 191 192class RandomHorizontalFlip(object):193 def __init__(self, p=0.5):194 self.p = p195 196 def __call__(self, img, target):197 if random.random() < self.p:198 return hflip(img, target)199 return img, target200 201 202class RandomResize(object):203 def __init__(self, sizes, max_size=None, equal=False):204 assert isinstance(sizes, (list, tuple))205 self.sizes = sizes206 self.max_size = max_size207 self.equal = equal208 209 def __call__(self, img, target=None):210 size = random.choice(self.sizes)211 if self.equal:212 return resize(img, target, size, size)213 else:214 return resize(img, target, size, self.max_size)215 216 217class ToTensor(object):218 def __call__(self, img, target):219 return F.to_tensor(img), target220 221 222class Normalize(object):223 def __init__(self, mean, std, max_image_size=512):224 self.mean = mean225 self.std = std226 self.max_image_size = max_image_size227 228 def __call__(self, image, target=None):229 image = F.normalize(image, mean=self.mean, std=self.std)230 if target is None:231 return image, None232 target = target.copy()233 # h, w = image.shape[-2:]234 h, w = target["size"][0], target["size"][1]235 if "boxes" in target:236 boxes = target["boxes"]237 boxes = boxes / self.max_image_size238 target["boxes"] = boxes239 if "polygons" in target:240 polygons = target["polygons"]241 scale = torch.cat([torch.tensor([w, h], dtype=torch.float32)242 for _ in range(polygons.shape[1] // 2)], dim=0)243 polygons = polygons / scale244 target["polygons"] = polygons245 return image, target246 247 248class Compose(object):249 def __init__(self, transforms):250 self.transforms = transforms251 252 def __call__(self, image, target):253 for t in self.transforms:254 image, target = t(image, target)255 return image, target256 257 def __repr__(self):258 format_string = self.__class__.__name__ + "("259 for t in self.transforms:260 format_string += "\n"261 format_string += " {0}".format(t)262 format_string += "\n)"263 return format_string264 265 266class LargeScaleJitter(object):267 """268 implementation of large scale jitter from copy_paste269 """270 271 def __init__(self, output_size=512, aug_scale_min=0.3, aug_scale_max=2.0):272 self.desired_size = torch.tensor([output_size])273 self.aug_scale_min = aug_scale_min274 self.aug_scale_max = aug_scale_max275 276 def rescale_target(self, scaled_size, image_size, target):277 # compute rescaled targets278 image_scale = scaled_size / image_size279 ratio_height, ratio_width = image_scale280 281 target = target.copy()282 target["size"] = scaled_size283 284 if "boxes" in target:285 boxes = target["boxes"]286 scaled_boxes = boxes * torch.as_tensor([ratio_width, ratio_height, ratio_width, ratio_height])287 target["boxes"] = scaled_boxes288 289 if "area" in target:290 area = target["area"]291 scaled_area = area * (ratio_width * ratio_height)292 target["area"] = scaled_area293 294 if "masks" in target:295 assert False296 masks = target['masks']297 # masks = interpolate(298 # masks[:, None].float(), scaled_size, mode="nearest")[:, 0] > 0.5299 target['masks'] = masks300 return target301 302 def crop_target(self, region, target):303 i, j, h, w = region304 fields = ["labels", "area"]305 306 target = target.copy()307 target["size"] = torch.tensor([h, w])308 309 if "boxes" in target:310 boxes = target["boxes"]311 max_size = torch.as_tensor([w, h], dtype=torch.float32)312 cropped_boxes = boxes - torch.as_tensor([j, i, j, i])313 cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size)314 cropped_boxes = cropped_boxes.clamp(min=0)315 area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1)316 target["boxes"] = cropped_boxes.reshape(-1, 4)317 target["area"] = area318 fields.append("boxes")319 320 if "masks" in target:321 # FIXME should we update the area here if there are no boxes?322 target['masks'] = target['masks'][:, i:i + h, j:j + w]323 fields.append("masks")324 325 # remove elements for which the boxes or masks that have zero area326 if "boxes" in target or "masks" in target:327 # favor boxes selection when defining which elements to keep328 # this is compatible with previous implementation329 if "boxes" in target:330 cropped_boxes = target['boxes'].reshape(-1, 2, 2)331 keep = torch.all(cropped_boxes[:, 1, :] > cropped_boxes[:, 0, :], dim=1)332 else:333 keep = target['masks'].flatten(1).any(1)334 335 for field in fields:336 target[field] = target[field][keep.tolist()]337 return target338 339 def pad_target(self, padding, target):340 target = target.copy()341 if "masks" in target:342 target['masks'] = torch.nn.functional.pad(target['masks'], (0, padding[1], 0, padding[0]))343 return target344 345 def __call__(self, image, target=None):346 image_size = image.size347 image_size = torch.tensor(image_size[::-1])348 349 random_scale = torch.rand(1) * (self.aug_scale_max - self.aug_scale_min) + self.aug_scale_min350 scaled_size = (random_scale * self.desired_size).round()351 352 scale = torch.maximum(scaled_size / image_size[0], scaled_size / image_size[1])353 scaled_size = (image_size * scale).round().int()354 355 scaled_image = F.resize(image, scaled_size.tolist(), interpolation=Image.BICUBIC)356 357 if target is not None:358 target = self.rescale_target(scaled_size, image_size, target)359 360 # randomly crop or pad images361 if random_scale >= 1:362 # Selects non-zero random offset (x, y) if scaled image is larger than desired_size.363 max_offset = scaled_size - self.desired_size364 offset = (max_offset * torch.rand(2)).floor().int()365 region = (offset[0].item(), offset[1].item(),366 self.desired_size[0].item(), self.desired_size[0].item())367 output_image = F.crop(scaled_image, *region)368 if target is not None:369 target = self.crop_target(region, target)370 else:371 assert False372 padding = self.desired_size - scaled_size373 output_image = F.pad(scaled_image, [0, 0, padding[1].item(), padding[0].item()])374 if target is not None:375 target = self.pad_target(padding, target)376 377 return output_image, target378 379 380class OriginLargeScaleJitter(object):381 """382 implementation of large scale jitter from copy_paste383 """384 385 def __init__(self, output_size=512, aug_scale_min=0.3, aug_scale_max=2.0):386 self.desired_size = torch.tensor(output_size)387 self.aug_scale_min = aug_scale_min388 self.aug_scale_max = aug_scale_max389 390 def rescale_target(self, scaled_size, image_size, target):391 # compute rescaled targets392 image_scale = scaled_size / image_size393 ratio_height, ratio_width = image_scale394 395 target = target.copy()396 target["size"] = scaled_size397 398 if "boxes" in target:399 boxes = target["boxes"]400 scaled_boxes = boxes * torch.as_tensor([ratio_width, ratio_height, ratio_width, ratio_height])401 target["boxes"] = scaled_boxes402 403 if "area" in target:404 area = target["area"]405 scaled_area = area * (ratio_width * ratio_height)406 target["area"] = scaled_area407 408 if "masks" in target:409 assert False410 masks = target['masks']411 # masks = interpolate(412 # masks[:, None].float(), scaled_size, mode="nearest")[:, 0] > 0.5413 target['masks'] = masks414 return target415 416 def crop_target(self, region, target):417 i, j, h, w = region418 fields = ["labels", "area"]419 420 target = target.copy()421 target["size"] = torch.tensor([h, w])422 423 if "boxes" in target:424 boxes = target["boxes"]425 max_size = torch.as_tensor([w, h], dtype=torch.float32)426 cropped_boxes = boxes - torch.as_tensor([j, i, j, i])427 cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size)428 cropped_boxes = cropped_boxes.clamp(min=0)429 area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1)430 target["boxes"] = cropped_boxes.reshape(-1, 4)431 target["area"] = area432 fields.append("boxes")433 434 if "masks" in target:435 # FIXME should we update the area here if there are no boxes?436 target['masks'] = target['masks'][:, i:i + h, j:j + w]437 fields.append("masks")438 439 # remove elements for which the boxes or masks that have zero area440 if "boxes" in target or "masks" in target:441 # favor boxes selection when defining which elements to keep442 # this is compatible with previous implementation443 if "boxes" in target:444 cropped_boxes = target['boxes'].reshape(-1, 2, 2)445 keep = torch.all(cropped_boxes[:, 1, :] > cropped_boxes[:, 0, :], dim=1)446 else:447 keep = target['masks'].flatten(1).any(1)448 449 for field in fields:450 target[field] = target[field][keep.tolist()]451 return target452 453 def pad_target(self, padding, target):454 target = target.copy()455 if "masks" in target:456 target['masks'] = torch.nn.functional.pad(target['masks'], (0, padding[1], 0, padding[0]))457 return target458 459 def __call__(self, image, target=None):460 image_size = image.size461 image_size = torch.tensor(image_size[::-1])462 463 out_desired_size = (self.desired_size * image_size / max(image_size)).round().int()464 465 random_scale = torch.rand(1) * (self.aug_scale_max - self.aug_scale_min) + self.aug_scale_min466 scaled_size = (random_scale * self.desired_size).round()467 468 scale = torch.minimum(scaled_size / image_size[0], scaled_size / image_size[1])469 scaled_size = (image_size * scale).round().int()470 471 scaled_image = F.resize(image, scaled_size.tolist())472 473 if target is not None:474 target = self.rescale_target(scaled_size, image_size, target)475 476 # randomly crop or pad images477 if random_scale > 1:478 # Selects non-zero random offset (x, y) if scaled image is larger than desired_size.479 max_offset = scaled_size - out_desired_size480 offset = (max_offset * torch.rand(2)).floor().int()481 region = (offset[0].item(), offset[1].item(),482 out_desired_size[0].item(), out_desired_size[1].item())483 output_image = F.crop(scaled_image, *region)484 if target is not None:485 target = self.crop_target(region, target)486 else:487 padding = out_desired_size - scaled_size488 output_image = F.pad(scaled_image, [0, 0, padding[1].item(), padding[0].item()])489 if target is not None:490 target = self.pad_target(padding, target)491 492 return output_image, target493 494 495class RandomDistortion(object):496 """497 Distort image w.r.t hue, saturation and exposure.498 """499 500 def __init__(self, brightness=0, contrast=0, saturation=0, hue=0, prob=0.5):501 self.prob = prob502 self.tfm = T.ColorJitter(brightness, contrast, saturation, hue)503 504 def __call__(self, img, target=None):505 if np.random.random() < self.prob:506 return self.tfm(img), target507 else:508 return img, target509