Spaces:
Runtime error
Runtime error
| #!/usr/bin/env python3 | |
| # -*- coding:utf-8 -*- | |
| # Copyright (c) Megvii, Inc. and its affiliates. | |
| import cv2 | |
| import numpy as np | |
| from yolox.utils import adjust_box_anns | |
| import random | |
| from ..data_augment import box_candidates, random_perspective, augment_hsv | |
| from .datasets_wrapper import Dataset | |
| def get_mosaic_coordinate(mosaic_image, mosaic_index, xc, yc, w, h, input_h, input_w): | |
| # TODO update doc | |
| # index0 to top left part of image | |
| if mosaic_index == 0: | |
| x1, y1, x2, y2 = max(xc - w, 0), max(yc - h, 0), xc, yc | |
| small_coord = w - (x2 - x1), h - (y2 - y1), w, h | |
| # index1 to top right part of image | |
| elif mosaic_index == 1: | |
| x1, y1, x2, y2 = xc, max(yc - h, 0), min(xc + w, input_w * 2), yc | |
| small_coord = 0, h - (y2 - y1), min(w, x2 - x1), h | |
| # index2 to bottom left part of image | |
| elif mosaic_index == 2: | |
| x1, y1, x2, y2 = max(xc - w, 0), yc, xc, min(input_h * 2, yc + h) | |
| small_coord = w - (x2 - x1), 0, w, min(y2 - y1, h) | |
| # index2 to bottom right part of image | |
| elif mosaic_index == 3: | |
| x1, y1, x2, y2 = xc, yc, min(xc + w, input_w * 2), min(input_h * 2, yc + h) # noqa | |
| small_coord = 0, 0, min(w, x2 - x1), min(y2 - y1, h) | |
| return (x1, y1, x2, y2), small_coord | |
| class MosaicDetection(Dataset): | |
| """Detection dataset wrapper that performs mixup for normal dataset.""" | |
| def __init__( | |
| self, dataset, img_size, mosaic=True, preproc=None, | |
| degrees=10.0, translate=0.1, scale=(0.5, 1.5), mscale=(0.5, 1.5), | |
| shear=2.0, perspective=0.0, enable_mixup=True, *args | |
| ): | |
| """ | |
| Args: | |
| dataset(Dataset) : Pytorch dataset object. | |
| img_size (tuple): | |
| mosaic (bool): enable mosaic augmentation or not. | |
| preproc (func): | |
| degrees (float): | |
| translate (float): | |
| scale (tuple): | |
| mscale (tuple): | |
| shear (float): | |
| perspective (float): | |
| enable_mixup (bool): | |
| *args(tuple) : Additional arguments for mixup random sampler. | |
| """ | |
| super().__init__(img_size, mosaic=mosaic) | |
| self._dataset = dataset | |
| self.preproc = preproc | |
| self.degrees = degrees | |
| self.translate = translate | |
| self.scale = scale | |
| self.shear = shear | |
| self.perspective = perspective | |
| self.mixup_scale = mscale | |
| self.enable_mosaic = mosaic | |
| self.enable_mixup = enable_mixup | |
| def __len__(self): | |
| return len(self._dataset) | |
| def __getitem__(self, idx): | |
| if self.enable_mosaic: | |
| mosaic_labels = [] | |
| input_dim = self._dataset.input_dim | |
| input_h, input_w = input_dim[0], input_dim[1] | |
| # yc, xc = s, s # mosaic center x, y | |
| yc = int(random.uniform(0.5 * input_h, 1.5 * input_h)) | |
| xc = int(random.uniform(0.5 * input_w, 1.5 * input_w)) | |
| # 3 additional image indices | |
| indices = [idx] + [random.randint(0, len(self._dataset) - 1) for _ in range(3)] | |
| for i_mosaic, index in enumerate(indices): | |
| img, _labels, _, _ = self._dataset.pull_item(index) | |
| h0, w0 = img.shape[:2] # orig hw | |
| scale = min(1. * input_h / h0, 1. * input_w / w0) | |
| img = cv2.resize( | |
| img, (int(w0 * scale), int(h0 * scale)), interpolation=cv2.INTER_LINEAR | |
| ) | |
| # generate output mosaic image | |
| (h, w, c) = img.shape[:3] | |
| if i_mosaic == 0: | |
| mosaic_img = np.full((input_h * 2, input_w * 2, c), 114, dtype=np.uint8) | |
| # suffix l means large image, while s means small image in mosaic aug. | |
| (l_x1, l_y1, l_x2, l_y2), (s_x1, s_y1, s_x2, s_y2) = get_mosaic_coordinate( | |
| mosaic_img, i_mosaic, xc, yc, w, h, input_h, input_w | |
| ) | |
| mosaic_img[l_y1:l_y2, l_x1:l_x2] = img[s_y1:s_y2, s_x1:s_x2] | |
| padw, padh = l_x1 - s_x1, l_y1 - s_y1 | |
| labels = _labels.copy() | |
| # Normalized xywh to pixel xyxy format | |
| if _labels.size > 0: | |
| labels[:, 0] = scale * _labels[:, 0] + padw | |
| labels[:, 1] = scale * _labels[:, 1] + padh | |
| labels[:, 2] = scale * _labels[:, 2] + padw | |
| labels[:, 3] = scale * _labels[:, 3] + padh | |
| mosaic_labels.append(labels) | |
| if len(mosaic_labels): | |
| mosaic_labels = np.concatenate(mosaic_labels, 0) | |
| ''' | |
| np.clip(mosaic_labels[:, 0], 0, 2 * input_w, out=mosaic_labels[:, 0]) | |
| np.clip(mosaic_labels[:, 1], 0, 2 * input_h, out=mosaic_labels[:, 1]) | |
| np.clip(mosaic_labels[:, 2], 0, 2 * input_w, out=mosaic_labels[:, 2]) | |
| np.clip(mosaic_labels[:, 3], 0, 2 * input_h, out=mosaic_labels[:, 3]) | |
| ''' | |
| mosaic_labels = mosaic_labels[mosaic_labels[:, 0] < 2 * input_w] | |
| mosaic_labels = mosaic_labels[mosaic_labels[:, 2] > 0] | |
| mosaic_labels = mosaic_labels[mosaic_labels[:, 1] < 2 * input_h] | |
| mosaic_labels = mosaic_labels[mosaic_labels[:, 3] > 0] | |
| #augment_hsv(mosaic_img) | |
| mosaic_img, mosaic_labels = random_perspective( | |
| mosaic_img, | |
| mosaic_labels, | |
| degrees=self.degrees, | |
| translate=self.translate, | |
| scale=self.scale, | |
| shear=self.shear, | |
| perspective=self.perspective, | |
| border=[-input_h // 2, -input_w // 2], | |
| ) # border to remove | |
| # ----------------------------------------------------------------- | |
| # CopyPaste: https://arxiv.org/abs/2012.07177 | |
| # ----------------------------------------------------------------- | |
| if self.enable_mixup and not len(mosaic_labels) == 0: | |
| mosaic_img, mosaic_labels = self.mixup(mosaic_img, mosaic_labels, self.input_dim) | |
| mix_img, padded_labels = self.preproc(mosaic_img, mosaic_labels, self.input_dim) | |
| img_info = (mix_img.shape[1], mix_img.shape[0]) | |
| return mix_img, padded_labels, img_info, np.array([idx]) | |
| else: | |
| self._dataset._input_dim = self.input_dim | |
| img, label, img_info, id_ = self._dataset.pull_item(idx) | |
| img, label = self.preproc(img, label, self.input_dim) | |
| return img, label, img_info, id_ | |
| def mixup(self, origin_img, origin_labels, input_dim): | |
| jit_factor = random.uniform(*self.mixup_scale) | |
| FLIP = random.uniform(0, 1) > 0.5 | |
| cp_labels = [] | |
| while len(cp_labels) == 0: | |
| cp_index = random.randint(0, self.__len__() - 1) | |
| cp_labels = self._dataset.load_anno(cp_index) | |
| img, cp_labels, _, _ = self._dataset.pull_item(cp_index) | |
| if len(img.shape) == 3: | |
| cp_img = np.ones((input_dim[0], input_dim[1], 3)) * 114.0 | |
| else: | |
| cp_img = np.ones(input_dim) * 114.0 | |
| cp_scale_ratio = min(input_dim[0] / img.shape[0], input_dim[1] / img.shape[1]) | |
| resized_img = cv2.resize( | |
| img, | |
| (int(img.shape[1] * cp_scale_ratio), int(img.shape[0] * cp_scale_ratio)), | |
| interpolation=cv2.INTER_LINEAR, | |
| ).astype(np.float32) | |
| cp_img[ | |
| : int(img.shape[0] * cp_scale_ratio), : int(img.shape[1] * cp_scale_ratio) | |
| ] = resized_img | |
| cp_img = cv2.resize( | |
| cp_img, | |
| (int(cp_img.shape[1] * jit_factor), int(cp_img.shape[0] * jit_factor)), | |
| ) | |
| cp_scale_ratio *= jit_factor | |
| if FLIP: | |
| cp_img = cp_img[:, ::-1, :] | |
| origin_h, origin_w = cp_img.shape[:2] | |
| target_h, target_w = origin_img.shape[:2] | |
| padded_img = np.zeros( | |
| (max(origin_h, target_h), max(origin_w, target_w), 3) | |
| ).astype(np.uint8) | |
| padded_img[:origin_h, :origin_w] = cp_img | |
| x_offset, y_offset = 0, 0 | |
| if padded_img.shape[0] > target_h: | |
| y_offset = random.randint(0, padded_img.shape[0] - target_h - 1) | |
| if padded_img.shape[1] > target_w: | |
| x_offset = random.randint(0, padded_img.shape[1] - target_w - 1) | |
| padded_cropped_img = padded_img[ | |
| y_offset: y_offset + target_h, x_offset: x_offset + target_w | |
| ] | |
| cp_bboxes_origin_np = adjust_box_anns( | |
| cp_labels[:, :4].copy(), cp_scale_ratio, 0, 0, origin_w, origin_h | |
| ) | |
| if FLIP: | |
| cp_bboxes_origin_np[:, 0::2] = ( | |
| origin_w - cp_bboxes_origin_np[:, 0::2][:, ::-1] | |
| ) | |
| cp_bboxes_transformed_np = cp_bboxes_origin_np.copy() | |
| ''' | |
| cp_bboxes_transformed_np[:, 0::2] = np.clip( | |
| cp_bboxes_transformed_np[:, 0::2] - x_offset, 0, target_w | |
| ) | |
| cp_bboxes_transformed_np[:, 1::2] = np.clip( | |
| cp_bboxes_transformed_np[:, 1::2] - y_offset, 0, target_h | |
| ) | |
| ''' | |
| cp_bboxes_transformed_np[:, 0::2] = cp_bboxes_transformed_np[:, 0::2] - x_offset | |
| cp_bboxes_transformed_np[:, 1::2] = cp_bboxes_transformed_np[:, 1::2] - y_offset | |
| keep_list = box_candidates(cp_bboxes_origin_np.T, cp_bboxes_transformed_np.T, 5) | |
| if keep_list.sum() >= 1.0: | |
| cls_labels = cp_labels[keep_list, 4:5].copy() | |
| id_labels = cp_labels[keep_list, 5:6].copy() | |
| box_labels = cp_bboxes_transformed_np[keep_list] | |
| labels = np.hstack((box_labels, cls_labels, id_labels)) | |
| # remove outside bbox | |
| labels = labels[labels[:, 0] < target_w] | |
| labels = labels[labels[:, 2] > 0] | |
| labels = labels[labels[:, 1] < target_h] | |
| labels = labels[labels[:, 3] > 0] | |
| origin_labels = np.vstack((origin_labels, labels)) | |
| origin_img = origin_img.astype(np.float32) | |
| origin_img = 0.5 * origin_img + 0.5 * padded_cropped_img.astype(np.float32) | |
| return origin_img, origin_labels | |