Spaces:
Runtime error
Runtime error
| # ------------------------------------------------------------------------ | |
| # Copyright (c) 2021 megvii-model. All Rights Reserved. | |
| # ------------------------------------------------------------------------ | |
| # Modified from Deformable DETR (https://github.com/fundamentalvision/Deformable-DETR) | |
| # Copyright (c) 2020 SenseTime. All Rights Reserved. | |
| # ------------------------------------------------------------------------ | |
| # Modified from DETR (https://github.com/facebookresearch/detr) | |
| # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved | |
| # ------------------------------------------------------------------------ | |
| """ | |
| Transforms and data augmentation for both image + bbox. | |
| """ | |
| import copy | |
| import random | |
| import PIL | |
| import torch | |
| import torchvision.transforms as T | |
| import torchvision.transforms.functional as F | |
| from PIL import Image, ImageDraw | |
| from util.box_ops import box_xyxy_to_cxcywh | |
| from util.misc import interpolate | |
| import numpy as np | |
| import os | |
| def crop_mot(image, target, region): | |
| cropped_image = F.crop(image, *region) | |
| target = target.copy() | |
| i, j, h, w = region | |
| # should we do something wrt the original size? | |
| target["size"] = torch.tensor([h, w]) | |
| fields = ["labels", "area", "iscrowd"] | |
| if 'obj_ids' in target: | |
| fields.append('obj_ids') | |
| if "boxes" in target: | |
| boxes = target["boxes"] | |
| max_size = torch.as_tensor([w, h], dtype=torch.float32) | |
| cropped_boxes = boxes - torch.as_tensor([j, i, j, i]) | |
| for i, box in enumerate(cropped_boxes): | |
| l, t, r, b = box | |
| # if l < 0: | |
| # l = 0 | |
| # if r < 0: | |
| # r = 0 | |
| # if l > w: | |
| # l = w | |
| # if r > w: | |
| # r = w | |
| # if t < 0: | |
| # t = 0 | |
| # if b < 0: | |
| # b = 0 | |
| # if t > h: | |
| # t = h | |
| # if b > h: | |
| # b = h | |
| if l < 0 and r < 0: | |
| l = r = 0 | |
| if l > w and r > w: | |
| l = r = w | |
| if t < 0 and b < 0: | |
| t = b = 0 | |
| if t > h and b > h: | |
| t = b = h | |
| cropped_boxes[i] = torch.tensor([l, t, r, b], dtype=box.dtype) | |
| cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size) | |
| cropped_boxes = cropped_boxes.clamp(min=0) | |
| area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1) | |
| target["boxes"] = cropped_boxes.reshape(-1, 4) | |
| target["area"] = area | |
| fields.append("boxes") | |
| if "masks" in target: | |
| # FIXME should we update the area here if there are no boxes? | |
| target['masks'] = target['masks'][:, i:i + h, j:j + w] | |
| fields.append("masks") | |
| # remove elements for which the boxes or masks that have zero area | |
| if "boxes" in target or "masks" in target: | |
| # favor boxes selection when defining which elements to keep | |
| # this is compatible with previous implementation | |
| if "boxes" in target: | |
| cropped_boxes = target['boxes'].reshape(-1, 2, 2) | |
| keep = torch.all(cropped_boxes[:, 1, :] > cropped_boxes[:, 0, :], dim=1) | |
| else: | |
| keep = target['masks'].flatten(1).any(1) | |
| for field in fields: | |
| target[field] = target[field][keep] | |
| return cropped_image, target | |
| def random_shift(image, target, region, sizes): | |
| oh, ow = sizes | |
| # step 1, shift crop and re-scale image firstly | |
| cropped_image = F.crop(image, *region) | |
| cropped_image = F.resize(cropped_image, sizes) | |
| target = target.copy() | |
| i, j, h, w = region | |
| # should we do something wrt the original size? | |
| target["size"] = torch.tensor([h, w]) | |
| fields = ["labels", "area", "iscrowd"] | |
| if 'obj_ids' in target: | |
| fields.append('obj_ids') | |
| if "boxes" in target: | |
| boxes = target["boxes"] | |
| max_size = torch.as_tensor([w, h], dtype=torch.float32) | |
| cropped_boxes = boxes - torch.as_tensor([j, i, j, i]) | |
| for i, box in enumerate(cropped_boxes): | |
| l, t, r, b = box | |
| if l < 0: | |
| l = 0 | |
| if r < 0: | |
| r = 0 | |
| if l > w: | |
| l = w | |
| if r > w: | |
| r = w | |
| if t < 0: | |
| t = 0 | |
| if b < 0: | |
| b = 0 | |
| if t > h: | |
| t = h | |
| if b > h: | |
| b = h | |
| # step 2, re-scale coords secondly | |
| ratio_h = 1.0 * oh / h | |
| ratio_w = 1.0 * ow / w | |
| cropped_boxes[i] = torch.tensor([ratio_w * l, ratio_h * t, ratio_w * r, ratio_h * b], dtype=box.dtype) | |
| cropped_boxes = cropped_boxes.reshape(-1, 2, 2) | |
| area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1) | |
| target["boxes"] = cropped_boxes.reshape(-1, 4) | |
| target["area"] = area | |
| fields.append("boxes") | |
| if "masks" in target: | |
| # FIXME should we update the area here if there are no boxes? | |
| target['masks'] = target['masks'][:, i:i + h, j:j + w] | |
| fields.append("masks") | |
| # remove elements for which the boxes or masks that have zero area | |
| if "boxes" in target or "masks" in target: | |
| # favor boxes selection when defining which elements to keep | |
| # this is compatible with previous implementation | |
| if "boxes" in target: | |
| cropped_boxes = target['boxes'].reshape(-1, 2, 2) | |
| keep = torch.all(cropped_boxes[:, 1, :] > cropped_boxes[:, 0, :], dim=1) | |
| else: | |
| keep = target['masks'].flatten(1).any(1) | |
| for field in fields: | |
| target[field] = target[field][keep] | |
| return cropped_image, target | |
| def crop(image, target, region): | |
| cropped_image = F.crop(image, *region) | |
| target = target.copy() | |
| i, j, h, w = region | |
| # should we do something wrt the original size? | |
| target["size"] = torch.tensor([h, w]) | |
| fields = ["labels", "area", "iscrowd"] | |
| if 'obj_ids' in target: | |
| fields.append('obj_ids') | |
| if "boxes" in target: | |
| boxes = target["boxes"] | |
| max_size = torch.as_tensor([w, h], dtype=torch.float32) | |
| cropped_boxes = boxes - torch.as_tensor([j, i, j, i]) | |
| cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size) | |
| cropped_boxes = cropped_boxes.clamp(min=0) | |
| area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1) | |
| target["boxes"] = cropped_boxes.reshape(-1, 4) | |
| target["area"] = area | |
| fields.append("boxes") | |
| if "masks" in target: | |
| # FIXME should we update the area here if there are no boxes? | |
| target['masks'] = target['masks'][:, i:i + h, j:j + w] | |
| fields.append("masks") | |
| # remove elements for which the boxes or masks that have zero area | |
| if "boxes" in target or "masks" in target: | |
| # favor boxes selection when defining which elements to keep | |
| # this is compatible with previous implementation | |
| if "boxes" in target: | |
| cropped_boxes = target['boxes'].reshape(-1, 2, 2) | |
| keep = torch.all(cropped_boxes[:, 1, :] > cropped_boxes[:, 0, :], dim=1) | |
| else: | |
| keep = target['masks'].flatten(1).any(1) | |
| for field in fields: | |
| target[field] = target[field][keep] | |
| return cropped_image, target | |
| def hflip(image, target): | |
| flipped_image = F.hflip(image) | |
| w, h = image.size | |
| target = target.copy() | |
| if "boxes" in target: | |
| boxes = target["boxes"] | |
| boxes = boxes[:, [2, 1, 0, 3]] * torch.as_tensor([-1, 1, -1, 1]) + torch.as_tensor([w, 0, w, 0]) | |
| target["boxes"] = boxes | |
| if "masks" in target: | |
| target['masks'] = target['masks'].flip(-1) | |
| return flipped_image, target | |
| def resize(image, target, size, max_size=None): | |
| # size can be min_size (scalar) or (w, h) tuple | |
| def get_size_with_aspect_ratio(image_size, size, max_size=None): | |
| w, h = image_size | |
| if max_size is not None: | |
| min_original_size = float(min((w, h))) | |
| max_original_size = float(max((w, h))) | |
| if max_original_size / min_original_size * size > max_size: | |
| size = int(round(max_size * min_original_size / max_original_size)) | |
| if (w <= h and w == size) or (h <= w and h == size): | |
| return (h, w) | |
| if w < h: | |
| ow = size | |
| oh = int(size * h / w) | |
| else: | |
| oh = size | |
| ow = int(size * w / h) | |
| return (oh, ow) | |
| def get_size(image_size, size, max_size=None): | |
| if isinstance(size, (list, tuple)): | |
| return size[::-1] | |
| else: | |
| return get_size_with_aspect_ratio(image_size, size, max_size) | |
| size = get_size(image.size, size, max_size) | |
| rescaled_image = F.resize(image, size) | |
| if target is None: | |
| return rescaled_image, None | |
| ratios = tuple(float(s) / float(s_orig) for s, s_orig in zip(rescaled_image.size, image.size)) | |
| ratio_width, ratio_height = ratios | |
| target = target.copy() | |
| if "boxes" in target: | |
| boxes = target["boxes"] | |
| scaled_boxes = boxes * torch.as_tensor([ratio_width, ratio_height, ratio_width, ratio_height]) | |
| target["boxes"] = scaled_boxes | |
| if "area" in target: | |
| area = target["area"] | |
| scaled_area = area * (ratio_width * ratio_height) | |
| target["area"] = scaled_area | |
| h, w = size | |
| target["size"] = torch.tensor([h, w]) | |
| if "masks" in target: | |
| target['masks'] = interpolate( | |
| target['masks'][:, None].float(), size, mode="nearest")[:, 0] > 0.5 | |
| return rescaled_image, target | |
| def pad(image, target, padding): | |
| # assumes that we only pad on the bottom right corners | |
| padded_image = F.pad(image, (0, 0, padding[0], padding[1])) | |
| if target is None: | |
| return padded_image, None | |
| target = target.copy() | |
| # should we do something wrt the original size? | |
| target["size"] = torch.tensor(padded_image[::-1]) | |
| if "masks" in target: | |
| target['masks'] = torch.nn.functional.pad(target['masks'], (0, padding[0], 0, padding[1])) | |
| return padded_image, target | |
| class RandomCrop(object): | |
| def __init__(self, size): | |
| self.size = size | |
| def __call__(self, img, target): | |
| region = T.RandomCrop.get_params(img, self.size) | |
| return crop(img, target, region) | |
| class MotRandomCrop(RandomCrop): | |
| def __call__(self, imgs: list, targets: list): | |
| ret_imgs = [] | |
| ret_targets = [] | |
| region = T.RandomCrop.get_params(imgs[0], self.size) | |
| for img_i, targets_i in zip(imgs, targets): | |
| img_i, targets_i = crop(img_i, targets_i, region) | |
| ret_imgs.append(img_i) | |
| ret_targets.append(targets_i) | |
| return ret_imgs, ret_targets | |
| class FixedMotRandomCrop(object): | |
| def __init__(self, min_size: int, max_size: int): | |
| self.min_size = min_size | |
| self.max_size = max_size | |
| def __call__(self, imgs: list, targets: list): | |
| ret_imgs = [] | |
| ret_targets = [] | |
| w = random.randint(self.min_size, min(imgs[0].width, self.max_size)) | |
| h = random.randint(self.min_size, min(imgs[0].height, self.max_size)) | |
| region = T.RandomCrop.get_params(imgs[0], [h, w]) | |
| for img_i, targets_i in zip(imgs, targets): | |
| img_i, targets_i = crop_mot(img_i, targets_i, region) | |
| ret_imgs.append(img_i) | |
| ret_targets.append(targets_i) | |
| return ret_imgs, ret_targets | |
| class MotRandomShift(object): | |
| def __init__(self, bs=1): | |
| self.bs = bs | |
| def __call__(self, imgs: list, targets: list): | |
| ret_imgs = copy.deepcopy(imgs) | |
| ret_targets = copy.deepcopy(targets) | |
| n_frames = len(imgs) | |
| select_i = random.choice(list(range(n_frames))) | |
| w, h = imgs[select_i].size | |
| xshift = (100 * torch.rand(self.bs)).int() | |
| xshift *= (torch.randn(self.bs) > 0.0).int() * 2 - 1 | |
| yshift = (100 * torch.rand(self.bs)).int() | |
| yshift *= (torch.randn(self.bs) > 0.0).int() * 2 - 1 | |
| ymin = max(0, -yshift[0]) | |
| ymax = min(h, h - yshift[0]) | |
| xmin = max(0, -xshift[0]) | |
| xmax = min(w, w - xshift[0]) | |
| region = (int(ymin), int(xmin), int(ymax-ymin), int(xmax-xmin)) | |
| ret_imgs[select_i], ret_targets[select_i] = random_shift(imgs[select_i], targets[select_i], region, (h,w)) | |
| return ret_imgs, ret_targets | |
| class FixedMotRandomShift(object): | |
| def __init__(self, bs=1, padding=50): | |
| self.bs = bs | |
| self.padding = padding | |
| def __call__(self, imgs: list, targets: list): | |
| ret_imgs = [] | |
| ret_targets = [] | |
| n_frames = len(imgs) | |
| w, h = imgs[0].size | |
| xshift = (self.padding * torch.rand(self.bs)).int() + 1 | |
| xshift *= (torch.randn(self.bs) > 0.0).int() * 2 - 1 | |
| yshift = (self.padding * torch.rand(self.bs)).int() + 1 | |
| yshift *= (torch.randn(self.bs) > 0.0).int() * 2 - 1 | |
| ret_imgs.append(imgs[0]) | |
| ret_targets.append(targets[0]) | |
| for i in range(1, n_frames): | |
| ymin = max(0, -yshift[0]) | |
| ymax = min(h, h - yshift[0]) | |
| xmin = max(0, -xshift[0]) | |
| xmax = min(w, w - xshift[0]) | |
| prev_img = ret_imgs[i-1].copy() | |
| prev_target = copy.deepcopy(ret_targets[i-1]) | |
| region = (int(ymin), int(xmin), int(ymax - ymin), int(xmax - xmin)) | |
| img_i, target_i = random_shift(prev_img, prev_target, region, (h, w)) | |
| ret_imgs.append(img_i) | |
| ret_targets.append(target_i) | |
| return ret_imgs, ret_targets | |
| class RandomSizeCrop(object): | |
| def __init__(self, min_size: int, max_size: int): | |
| self.min_size = min_size | |
| self.max_size = max_size | |
| def __call__(self, img: PIL.Image.Image, target: dict): | |
| w = random.randint(self.min_size, min(img.width, self.max_size)) | |
| h = random.randint(self.min_size, min(img.height, self.max_size)) | |
| region = T.RandomCrop.get_params(img, [h, w]) | |
| return crop(img, target, region) | |
| class MotRandomSizeCrop(RandomSizeCrop): | |
| def __call__(self, imgs, targets): | |
| w = random.randint(self.min_size, min(imgs[0].width, self.max_size)) | |
| h = random.randint(self.min_size, min(imgs[0].height, self.max_size)) | |
| region = T.RandomCrop.get_params(imgs[0], [h, w]) | |
| ret_imgs = [] | |
| ret_targets = [] | |
| for img_i, targets_i in zip(imgs, targets): | |
| img_i, targets_i = crop(img_i, targets_i, region) | |
| ret_imgs.append(img_i) | |
| ret_targets.append(targets_i) | |
| return ret_imgs, ret_targets | |
| class CenterCrop(object): | |
| def __init__(self, size): | |
| self.size = size | |
| def __call__(self, img, target): | |
| image_width, image_height = img.size | |
| crop_height, crop_width = self.size | |
| crop_top = int(round((image_height - crop_height) / 2.)) | |
| crop_left = int(round((image_width - crop_width) / 2.)) | |
| return crop(img, target, (crop_top, crop_left, crop_height, crop_width)) | |
| class MotCenterCrop(CenterCrop): | |
| def __call__(self, imgs, targets): | |
| image_width, image_height = imgs[0].size | |
| crop_height, crop_width = self.size | |
| crop_top = int(round((image_height - crop_height) / 2.)) | |
| crop_left = int(round((image_width - crop_width) / 2.)) | |
| ret_imgs = [] | |
| ret_targets = [] | |
| for img_i, targets_i in zip(imgs, targets): | |
| img_i, targets_i = crop(img_i, targets_i, (crop_top, crop_left, crop_height, crop_width)) | |
| ret_imgs.append(img_i) | |
| ret_targets.append(targets_i) | |
| return ret_imgs, ret_targets | |
| class RandomHorizontalFlip(object): | |
| def __init__(self, p=0.5): | |
| self.p = p | |
| def __call__(self, img, target): | |
| if random.random() < self.p: | |
| return hflip(img, target) | |
| return img, target | |
| class MotRandomHorizontalFlip(RandomHorizontalFlip): | |
| def __call__(self, imgs, targets): | |
| if random.random() < self.p: | |
| ret_imgs = [] | |
| ret_targets = [] | |
| for img_i, targets_i in zip(imgs, targets): | |
| img_i, targets_i = hflip(img_i, targets_i) | |
| ret_imgs.append(img_i) | |
| ret_targets.append(targets_i) | |
| return ret_imgs, ret_targets | |
| return imgs, targets | |
| class RandomResize(object): | |
| def __init__(self, sizes, max_size=None): | |
| assert isinstance(sizes, (list, tuple)) | |
| self.sizes = sizes | |
| self.max_size = max_size | |
| def __call__(self, img, target=None): | |
| size = random.choice(self.sizes) | |
| return resize(img, target, size, self.max_size) | |
| class MotRandomResize(RandomResize): | |
| def __call__(self, imgs, targets): | |
| size = random.choice(self.sizes) | |
| ret_imgs = [] | |
| ret_targets = [] | |
| for img_i, targets_i in zip(imgs, targets): | |
| img_i, targets_i = resize(img_i, targets_i, size, self.max_size) | |
| ret_imgs.append(img_i) | |
| ret_targets.append(targets_i) | |
| return ret_imgs, ret_targets | |
| class RandomPad(object): | |
| def __init__(self, max_pad): | |
| self.max_pad = max_pad | |
| def __call__(self, img, target): | |
| pad_x = random.randint(0, self.max_pad) | |
| pad_y = random.randint(0, self.max_pad) | |
| return pad(img, target, (pad_x, pad_y)) | |
| class MotRandomPad(RandomPad): | |
| def __call__(self, imgs, targets): | |
| pad_x = random.randint(0, self.max_pad) | |
| pad_y = random.randint(0, self.max_pad) | |
| ret_imgs = [] | |
| ret_targets = [] | |
| for img_i, targets_i in zip(imgs, targets): | |
| img_i, target_i = pad(img_i, targets_i, (pad_x, pad_y)) | |
| ret_imgs.append(img_i) | |
| ret_targets.append(targets_i) | |
| return ret_imgs, ret_targets | |
| class RandomSelect(object): | |
| """ | |
| Randomly selects between transforms1 and transforms2, | |
| with probability p for transforms1 and (1 - p) for transforms2 | |
| """ | |
| def __init__(self, transforms1, transforms2, p=0.5): | |
| self.transforms1 = transforms1 | |
| self.transforms2 = transforms2 | |
| self.p = p | |
| def __call__(self, img, target): | |
| if random.random() < self.p: | |
| return self.transforms1(img, target) | |
| return self.transforms2(img, target) | |
| class MotRandomSelect(RandomSelect): | |
| """ | |
| Randomly selects between transforms1 and transforms2, | |
| with probability p for transforms1 and (1 - p) for transforms2 | |
| """ | |
| def __call__(self, imgs, targets): | |
| if random.random() < self.p: | |
| return self.transforms1(imgs, targets) | |
| return self.transforms2(imgs, targets) | |
| class ToTensor(object): | |
| def __call__(self, img, target): | |
| return F.to_tensor(img), target | |
| class MotToTensor(ToTensor): | |
| def __call__(self, imgs, targets): | |
| ret_imgs = [] | |
| for img in imgs: | |
| ret_imgs.append(F.to_tensor(img)) | |
| return ret_imgs, targets | |
| class RandomErasing(object): | |
| def __init__(self, *args, **kwargs): | |
| self.eraser = T.RandomErasing(*args, **kwargs) | |
| def __call__(self, img, target): | |
| return self.eraser(img), target | |
| class MotRandomErasing(RandomErasing): | |
| def __call__(self, imgs, targets): | |
| # TODO: Rewrite this part to ensure the data augmentation is same to each image. | |
| ret_imgs = [] | |
| for img_i, targets_i in zip(imgs, targets): | |
| ret_imgs.append(self.eraser(img_i)) | |
| return ret_imgs, targets | |
| class MoTColorJitter(T.ColorJitter): | |
| def __call__(self, imgs, targets): | |
| transform = self.get_params(self.brightness, self.contrast, | |
| self.saturation, self.hue) | |
| ret_imgs = [] | |
| for img_i, targets_i in zip(imgs, targets): | |
| ret_imgs.append(transform(img_i)) | |
| return ret_imgs, targets | |
| class Normalize(object): | |
| def __init__(self, mean, std): | |
| self.mean = mean | |
| self.std = std | |
| def __call__(self, image, target=None): | |
| if target is not None: | |
| target['ori_img'] = image.clone() | |
| image = F.normalize(image, mean=self.mean, std=self.std) | |
| if target is None: | |
| return image, None | |
| target = target.copy() | |
| h, w = image.shape[-2:] | |
| if "boxes" in target: | |
| boxes = target["boxes"] | |
| boxes = box_xyxy_to_cxcywh(boxes) | |
| boxes = boxes / torch.tensor([w, h, w, h], dtype=torch.float32) | |
| target["boxes"] = boxes | |
| return image, target | |
| class MotNormalize(Normalize): | |
| def __call__(self, imgs, targets=None): | |
| ret_imgs = [] | |
| ret_targets = [] | |
| for i in range(len(imgs)): | |
| img_i = imgs[i] | |
| targets_i = targets[i] if targets is not None else None | |
| img_i, targets_i = super().__call__(img_i, targets_i) | |
| ret_imgs.append(img_i) | |
| ret_targets.append(targets_i) | |
| return ret_imgs, ret_targets | |
| class Compose(object): | |
| def __init__(self, transforms): | |
| self.transforms = transforms | |
| def __call__(self, image, target): | |
| for t in self.transforms: | |
| image, target = t(image, target) | |
| return image, target | |
| def __repr__(self): | |
| format_string = self.__class__.__name__ + "(" | |
| for t in self.transforms: | |
| format_string += "\n" | |
| format_string += " {0}".format(t) | |
| format_string += "\n)" | |
| return format_string | |
| class MotCompose(Compose): | |
| def __call__(self, imgs, targets): | |
| for t in self.transforms: | |
| imgs, targets = t(imgs, targets) | |
| return imgs, targets | |