Spaces:
Runtime error
Runtime error
| # Copyright (c) Facebook, Inc. and its affiliates. | |
| # Copyright (c) Meta Platforms, Inc. All Rights Reserved | |
| import math | |
| import numbers | |
| import numpy as np | |
| from detectron2.data.transforms.augmentation import Augmentation | |
| from detectron2.data.transforms.transform import ( | |
| CropTransform, | |
| ResizeTransform, | |
| TransformList, | |
| ) | |
| from PIL import Image | |
| from fvcore.transforms.transform import PadTransform | |
| def mask2box(mask: np.ndarray): | |
| # use naive way | |
| row = np.nonzero(mask.sum(axis=0))[0] | |
| if len(row) == 0: | |
| return None | |
| x1 = row.min() | |
| x2 = row.max() | |
| col = np.nonzero(mask.sum(axis=1))[0] | |
| y1 = col.min() | |
| y2 = col.max() | |
| return x1, y1, x2 + 1 - x1, y2 + 1 - y1 | |
| def expand_box(x, y, w, h, expand_ratio=1.0, max_h=None, max_w=None): | |
| cx = x + 0.5 * w | |
| cy = y + 0.5 * h | |
| w = w * expand_ratio | |
| h = h * expand_ratio | |
| box = [cx - 0.5 * w, cy - 0.5 * h, cx + 0.5 * w, cy + 0.5 * h] | |
| if max_h is not None: | |
| box[1] = max(0, box[1]) | |
| box[3] = min(max_h - 1, box[3]) | |
| if max_w is not None: | |
| box[0] = max(0, box[0]) | |
| box[2] = min(max_w - 1, box[2]) | |
| box[2] = box[2] - box[0] | |
| box[3] = box[3] - box[1] | |
| return [int(b) for b in box] | |
| class CropImageWithMask(Augmentation): | |
| def __init__(self, expand_ratio=1.0, mode="choice"): | |
| if isinstance(expand_ratio, numbers.Number): | |
| expand_ratio = (expand_ratio, expand_ratio) | |
| self.mode = mode | |
| self.expand_ratio = expand_ratio | |
| if self.mode == "range": | |
| assert len(expand_ratio) == 2 and expand_ratio[0] < expand_ratio[1] | |
| def get_transform(self, image, sem_seg, category_id): | |
| input_size = image.shape[:2] | |
| bin_mask = sem_seg == category_id | |
| x, y, w, h = mask2box(bin_mask) | |
| if self.mode == "choice": | |
| expand_ratio = np.random.choice(self.expand_ratio) | |
| else: | |
| expand_ratio = np.random.uniform(self.expand_ratio[0], self.expand_ratio[1]) | |
| x, y, w, h = expand_box(x, y, w, h, expand_ratio, *input_size) | |
| w = max(w, 1) | |
| h = max(h, 1) | |
| return CropTransform(x, y, w, h, input_size[1], input_size[0]) | |
| class CropImageWithBox(Augmentation): | |
| def __init__(self, expand_ratio=1.0, mode="choice"): | |
| if isinstance(expand_ratio, numbers.Number): | |
| expand_ratio = (expand_ratio, expand_ratio) | |
| self.mode = mode | |
| self.expand_ratio = expand_ratio | |
| if self.mode == "range": | |
| assert len(expand_ratio) == 2 and expand_ratio[0] < expand_ratio[1] | |
| def get_transform(self, image, boxes): | |
| input_size = image.shape[:2] | |
| x, y, x2, y2 = boxes[0] | |
| w = x2 - x + 1 | |
| h = y2 - y + 1 | |
| if self.mode == "choice": | |
| expand_ratio = np.random.choice(self.expand_ratio) | |
| else: | |
| expand_ratio = np.random.uniform(self.expand_ratio[0], self.expand_ratio[1]) | |
| x, y, w, h = expand_box(x, y, w, h, expand_ratio, *input_size) | |
| w = max(w, 1) | |
| h = max(h, 1) | |
| return CropTransform(x, y, w, h, input_size[1], input_size[0]) | |
| class RandomResizedCrop(Augmentation): | |
| def __init__( | |
| self, | |
| size, | |
| scale=(0.08, 1.0), | |
| ratio=(3.0 / 4.0, 4.0 / 3.0), | |
| interpolation=Image.BILINEAR, | |
| ): | |
| if isinstance(size, int): | |
| size = (size, size) | |
| else: | |
| assert isinstance(size, (tuple, list)) and len(size) == 2 | |
| self.size = size | |
| self.scale = scale | |
| self.ratio = ratio | |
| self.interpolation = interpolation | |
| def get_transform(self, image): | |
| height, width = image.shape[:2] | |
| area = height * width | |
| log_ratio = np.log(np.array(self.ratio)) | |
| is_success = False | |
| for _ in range(10): | |
| target_area = area * np.random.uniform(self.scale[0], self.scale[1]) | |
| aspect_ratio = np.exp(np.random.uniform(log_ratio[0], log_ratio[1])) | |
| w = int(round(math.sqrt(target_area * aspect_ratio))) | |
| h = int(round(math.sqrt(target_area / aspect_ratio))) | |
| if 0 < w <= width and 0 < h <= height: | |
| i = np.random.randint(0, width - w + 1) | |
| j = np.random.randint(0, height - h + 1) | |
| is_success = True | |
| break | |
| if not is_success: | |
| # Fallback to central crop | |
| in_ratio = float(width) / float(height) | |
| if in_ratio < min(self.ratio): | |
| w = width | |
| h = int(round(w / min(self.ratio))) | |
| elif in_ratio > max(self.ratio): | |
| h = height | |
| w = int(round(h * max(self.ratio))) | |
| else: # whole image | |
| w = width | |
| h = height | |
| i = (width - w) // 2 | |
| j = (height - h) // 2 | |
| return TransformList( | |
| [ | |
| CropTransform(i, j, w, h, width, height), | |
| ResizeTransform( | |
| h, w, self.size[1], self.size[0], interp=self.interpolation | |
| ), | |
| ] | |
| ) | |
| class CenterCrop(Augmentation): | |
| def __init__(self, size, seg_ignore_label): | |
| if isinstance(size, numbers.Number): | |
| size = (int(size), int(size)) | |
| elif isinstance(size, (tuple, list)) and len(size) == 1: | |
| size = (size[0], size[0]) | |
| self.size = size | |
| self.seg_ignore_label = seg_ignore_label | |
| def get_transform(self, image): | |
| image_height, image_width = image.shape[:2] | |
| crop_height, crop_width = self.size | |
| transforms = [] | |
| if crop_width > image_width or crop_height > image_height: | |
| padding_ltrb = [ | |
| (crop_width - image_width) // 2 if crop_width > image_width else 0, | |
| (crop_height - image_height) // 2 if crop_height > image_height else 0, | |
| (crop_width - image_width + 1) // 2 if crop_width > image_width else 0, | |
| (crop_height - image_height + 1) // 2 | |
| if crop_height > image_height | |
| else 0, | |
| ] | |
| transforms.append( | |
| PadTransform( | |
| *padding_ltrb, | |
| orig_w=image_width, | |
| orig_h=image_height, | |
| seg_pad_value=self.seg_ignore_label | |
| ) | |
| ) | |
| image_width, image_height = ( | |
| image_width + padding_ltrb[0] + padding_ltrb[2], | |
| image_height + padding_ltrb[1] + padding_ltrb[3], | |
| ) | |
| crop_top = int(round((image_height - crop_height) / 2.0)) | |
| crop_left = int(round((image_width - crop_width) / 2.0)) | |
| transforms.append( | |
| CropTransform( | |
| crop_left, crop_top, crop_width, crop_height, image_width, image_height | |
| ) | |
| ) | |
| return TransformList(transforms) | |