|
| 1 | +""" |
| 2 | +data augmentations for CALD method, including horizontal_flip, rotate(5'), cutout |
| 3 | +official code: https://github.com/we1pingyu/CALD/blob/master/cald/cald_helper.py |
| 4 | +""" |
| 5 | +import random |
| 6 | +from typing import Any, List, Tuple |
| 7 | + |
| 8 | +import cv2 |
| 9 | +import numpy as np |
| 10 | +from nptyping import NDArray |
| 11 | + |
| 12 | +from utils.ymir_yolov5 import BBOX, CV_IMAGE |
| 13 | + |
| 14 | + |
| 15 | +def intersect(boxes1: BBOX, boxes2: BBOX) -> NDArray: |
| 16 | + ''' |
| 17 | + Find intersection of every box combination between two sets of box |
| 18 | + boxes1: bounding boxes 1, a tensor of dimensions (n1, 4) |
| 19 | + boxes2: bounding boxes 2, a tensor of dimensions (n2, 4) |
| 20 | +
|
| 21 | + Out: Intersection each of boxes1 with respect to each of boxes2, |
| 22 | + a tensor of dimensions (n1, n2) |
| 23 | + ''' |
| 24 | + n1 = boxes1.shape[0] |
| 25 | + n2 = boxes2.shape[0] |
| 26 | + max_xy = np.minimum( |
| 27 | + np.expand_dims(boxes1[:, 2:], axis=1).repeat(n2, axis=1), |
| 28 | + np.expand_dims(boxes2[:, 2:], axis=0).repeat(n1, axis=0)) |
| 29 | + |
| 30 | + min_xy = np.maximum( |
| 31 | + np.expand_dims(boxes1[:, :2], axis=1).repeat(n2, axis=1), |
| 32 | + np.expand_dims(boxes2[:, :2], axis=0).repeat(n1, axis=0)) |
| 33 | + inter = np.clip(max_xy - min_xy, a_min=0, a_max=None) # (n1, n2, 2) |
| 34 | + return inter[:, :, 0] * inter[:, :, 1] # (n1, n2) |
| 35 | + |
| 36 | + |
| 37 | +def horizontal_flip(image: CV_IMAGE, bbox: BBOX) \ |
| 38 | + -> Tuple[CV_IMAGE, BBOX]: |
| 39 | + """ |
| 40 | + image: opencv image, [height,width,channels] |
| 41 | + bbox: numpy.ndarray, [N,4] --> [x1,y1,x2,y2] |
| 42 | + """ |
| 43 | + image = image.copy() |
| 44 | + |
| 45 | + width = image.shape[1] |
| 46 | + # Flip image horizontally |
| 47 | + image = image[:, ::-1, :] |
| 48 | + if len(bbox) > 0: |
| 49 | + bbox = bbox.copy() |
| 50 | + # Flip bbox horizontally |
| 51 | + bbox[:, [0, 2]] = width - bbox[:, [2, 0]] |
| 52 | + return image, bbox |
| 53 | + |
| 54 | + |
| 55 | +def cutout(image: CV_IMAGE, |
| 56 | + bbox: BBOX, |
| 57 | + cut_num: int = 2, |
| 58 | + fill_val: int = 0, |
| 59 | + bbox_remove_thres: float = 0.4, |
| 60 | + bbox_min_thres: float = 0.1) -> Tuple[CV_IMAGE, BBOX]: |
| 61 | + ''' |
| 62 | + Cutout augmentation |
| 63 | + image: A PIL image |
| 64 | + boxes: bounding boxes, a tensor of dimensions (#objects, 4) |
| 65 | + labels: labels of object, a tensor of dimensions (#objects) |
| 66 | + fill_val: Value filled in cut out |
| 67 | + bbox_remove_thres: Theshold to remove bbox cut by cutout |
| 68 | +
|
| 69 | + Out: new image, new_boxes, new_labels |
| 70 | + ''' |
| 71 | + image = image.copy() |
| 72 | + bbox = bbox.copy() |
| 73 | + |
| 74 | + if len(bbox) == 0: |
| 75 | + return image, bbox |
| 76 | + |
| 77 | + original_h, original_w, original_channel = image.shape |
| 78 | + count = 0 |
| 79 | + for _ in range(50): |
| 80 | + # Random cutout size: [0.15, 0.5] of original dimension |
| 81 | + cutout_size_h = random.uniform(0.05 * original_h, 0.2 * original_h) |
| 82 | + cutout_size_w = random.uniform(0.05 * original_w, 0.2 * original_w) |
| 83 | + |
| 84 | + # Random position for cutout |
| 85 | + left = random.uniform(0, original_w - cutout_size_w) |
| 86 | + right = left + cutout_size_w |
| 87 | + top = random.uniform(0, original_h - cutout_size_h) |
| 88 | + bottom = top + cutout_size_h |
| 89 | + cutout = np.array([[float(left), float(top), float(right), float(bottom)]]) |
| 90 | + |
| 91 | + # Calculate intersect between cutout and bounding boxes |
| 92 | + overlap_size = intersect(cutout, bbox) |
| 93 | + area_boxes = (bbox[:, 2] - bbox[:, 0]) * (bbox[:, 3] - bbox[:, 1]) |
| 94 | + ratio = overlap_size / (area_boxes + 1e-14) |
| 95 | + # If all boxes have Iou greater than bbox_remove_thres, try again |
| 96 | + if ratio.max() > bbox_remove_thres or ratio.max() < bbox_min_thres: |
| 97 | + continue |
| 98 | + |
| 99 | + image[int(top):int(bottom), int(left):int(right), :] = fill_val |
| 100 | + count += 1 |
| 101 | + if count >= cut_num: |
| 102 | + break |
| 103 | + return image, bbox |
| 104 | + |
| 105 | + |
| 106 | +def rotate(image: CV_IMAGE, bbox: BBOX, rot: float = 5) -> Tuple[CV_IMAGE, BBOX]: |
| 107 | + image = image.copy() |
| 108 | + bbox = bbox.copy() |
| 109 | + h, w, c = image.shape |
| 110 | + center = np.array([w / 2.0, h / 2.0]) |
| 111 | + s = max(h, w) * 1.0 |
| 112 | + trans = get_affine_transform(center, s, rot, [w, h]) |
| 113 | + if len(bbox) > 0: |
| 114 | + for i in range(bbox.shape[0]): |
| 115 | + x1, y1 = affine_transform(bbox[i, :2], trans) |
| 116 | + x2, y2 = affine_transform(bbox[i, 2:], trans) |
| 117 | + x3, y3 = affine_transform(bbox[i, [2, 1]], trans) |
| 118 | + x4, y4 = affine_transform(bbox[i, [0, 3]], trans) |
| 119 | + bbox[i, :2] = [min(x1, x2, x3, x4), min(y1, y2, y3, y4)] |
| 120 | + bbox[i, 2:] = [max(x1, x2, x3, x4), max(y1, y2, y3, y4)] |
| 121 | + image = cv2.warpAffine(image, trans, (w, h), flags=cv2.INTER_LINEAR) |
| 122 | + return image, bbox |
| 123 | + |
| 124 | + |
| 125 | +def get_3rd_point(a: NDArray, b: NDArray) -> NDArray: |
| 126 | + direct = a - b |
| 127 | + return b + np.array([-direct[1], direct[0]], dtype=np.float32) |
| 128 | + |
| 129 | + |
| 130 | +def get_dir(src_point: NDArray, rot_rad: float) -> List: |
| 131 | + sn, cs = np.sin(rot_rad), np.cos(rot_rad) |
| 132 | + |
| 133 | + src_result = [0, 0] |
| 134 | + src_result[0] = src_point[0] * cs - src_point[1] * sn |
| 135 | + src_result[1] = src_point[0] * sn + src_point[1] * cs |
| 136 | + |
| 137 | + return src_result |
| 138 | + |
| 139 | + |
| 140 | +def transform_preds(coords: NDArray, center: NDArray, scale: Any, rot: float, output_size: List) -> NDArray: |
| 141 | + trans = get_affine_transform(center, scale, rot, output_size, inv=True) |
| 142 | + target_coords = affine_transform(coords, trans) |
| 143 | + return target_coords |
| 144 | + |
| 145 | + |
| 146 | +def get_affine_transform(center: NDArray, |
| 147 | + scale: Any, |
| 148 | + rot: float, |
| 149 | + output_size: List, |
| 150 | + shift: NDArray = np.array([0, 0], dtype=np.float32), |
| 151 | + inv: bool = False) -> NDArray: |
| 152 | + if not isinstance(scale, np.ndarray) and not isinstance(scale, list): |
| 153 | + scale = np.array([scale, scale], dtype=np.float32) |
| 154 | + |
| 155 | + scale_tmp = scale |
| 156 | + src_w = scale_tmp[0] |
| 157 | + dst_w = output_size[0] |
| 158 | + dst_h = output_size[1] |
| 159 | + |
| 160 | + rot_rad = np.pi * rot / 180 |
| 161 | + src_dir = get_dir([0, src_w * -0.5], rot_rad) |
| 162 | + dst_dir = np.array([0, dst_w * -0.5], np.float32) |
| 163 | + |
| 164 | + src = np.zeros((3, 2), dtype=np.float32) |
| 165 | + dst = np.zeros((3, 2), dtype=np.float32) |
| 166 | + src[0, :] = center + scale_tmp * shift |
| 167 | + src[1, :] = center + src_dir + scale_tmp * shift |
| 168 | + dst[0, :] = [dst_w * 0.5, dst_h * 0.5] |
| 169 | + dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5], np.float32) + dst_dir |
| 170 | + |
| 171 | + src[2:, :] = get_3rd_point(src[0, :], src[1, :]) |
| 172 | + dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :]) |
| 173 | + |
| 174 | + if inv: |
| 175 | + trans = cv2.getAffineTransform(np.float32(dst), np.float32(src)) |
| 176 | + else: |
| 177 | + trans = cv2.getAffineTransform(np.float32(src), np.float32(dst)) |
| 178 | + |
| 179 | + return trans |
| 180 | + |
| 181 | + |
| 182 | +def affine_transform(pt: NDArray, t: NDArray) -> NDArray: |
| 183 | + new_pt = np.array([pt[0], pt[1], 1.], dtype=np.float32).T |
| 184 | + new_pt = np.dot(t, new_pt) |
| 185 | + return new_pt[:2] |
| 186 | + |
| 187 | + |
| 188 | +def resize(img: CV_IMAGE, boxes: BBOX, ratio: float = 0.8) -> Tuple[CV_IMAGE, BBOX]: |
| 189 | + """ |
| 190 | + ratio: <= 1.0 |
| 191 | + """ |
| 192 | + assert ratio <= 1.0, f'resize ratio {ratio} must <= 1.0' |
| 193 | + |
| 194 | + h, w, _ = img.shape |
| 195 | + ow = int(w * ratio) |
| 196 | + oh = int(h * ratio) |
| 197 | + resize_img = cv2.resize(img, (ow, oh)) |
| 198 | + new_img = np.zeros_like(img) |
| 199 | + new_img[:oh, :ow] = resize_img |
| 200 | + |
| 201 | + if len(boxes) == 0: |
| 202 | + return new_img, boxes |
| 203 | + else: |
| 204 | + return new_img, boxes * ratio |
0 commit comments