| """ |
| Copyright (c) Meta Platforms, Inc. and affiliates. |
| """ |
|
|
| import numbers |
| import random |
|
|
| import numpy as np |
|
|
|
|
| def get_co_transforms(aug_args): |
| transforms = [] |
| if aug_args.swap: |
| transforms.append(RandomTemporalSwap()) |
| if aug_args.hflip: |
| transforms.append(RandomHorizontalFlip()) |
| if aug_args.crop: |
| transforms.append(RandomCrop(aug_args.para_crop)) |
| return Compose(transforms) |
|
|
|
|
| class Compose: |
| def __init__(self, co_transforms): |
| self.co_transforms = co_transforms |
|
|
| def __call__(self, imgs, full_segs, key_objs, target): |
| for t in self.co_transforms: |
| imgs, full_segs, key_objs, target = t(imgs, full_segs, key_objs, target) |
| return imgs, full_segs, key_objs, target |
|
|
|
|
| class RandomCrop: |
| """Crops the given PIL.Image at a random location to have a region of |
| the given size. size can be a tuple (target_height, target_width) |
| or an integer, in which case the target will be of a square shape (size, size) |
| """ |
|
|
| def __init__(self, size): |
| if isinstance(size, numbers.Number): |
| self.size = (int(size), int(size)) |
| else: |
| self.size = size |
|
|
| def __call__(self, imgs, full_segs, key_objs, target): |
| h, w, _ = imgs[0].shape |
| th, tw = self.size |
| if w == tw and h == th: |
| return imgs, target |
|
|
| x1 = random.randint(0, w - tw) |
| y1 = random.randint(0, h - th) |
| imgs = [img[y1 : y1 + th, x1 : x1 + tw] for img in imgs] |
| full_segs = [full_seg[y1 : y1 + th, x1 : x1 + tw] for full_seg in full_segs] |
| key_objs = [key_obj[y1 : y1 + th, x1 : x1 + tw] for key_obj in key_objs] |
|
|
| if target != {}: |
| raise NotImplementedError( |
| "RandomCrop currently does not take ground-truth labels" |
| ) |
|
|
| return imgs, full_segs, key_objs, target |
|
|
|
|
| class RandomTemporalSwap: |
| """Randomly swap first and second frames""" |
|
|
| def __call__(self, imgs, full_segs, key_objs, target): |
|
|
| if random.random() < 0.5: |
| imgs = imgs[::-1] |
| full_segs = full_segs[::-1] |
| key_objs = key_objs[::-1] |
|
|
| if target != {}: |
| raise NotImplementedError( |
| "RandomTemporalSwap currently does not take ground-truth labels" |
| ) |
|
|
| return imgs, full_segs, key_objs, target |
|
|
|
|
| class RandomHorizontalFlip: |
| """Randomly horizontally flips the given PIL.Image with a probability of 0.5""" |
|
|
| def __call__(self, imgs, full_segs, key_objs, target): |
| if random.random() < 0.5: |
| imgs = [np.copy(np.fliplr(im)) for im in imgs] |
| full_segs = [np.copy(np.fliplr(full_seg)) for full_seg in full_segs] |
| key_objs = [np.copy(np.fliplr(key_obj)) for key_obj in key_objs] |
|
|
| if target != {}: |
| raise NotImplementedError( |
| "RandomHorizontalFlip currently does not take ground-truth labels" |
| ) |
| return imgs, full_segs, key_objs, target |
|
|
