Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| # -*- coding:utf-8 -*- | |
| # This code is based on | |
| # https://github.com/ultralytics/yolov5/blob/master/utils/dataloaders.py | |
| import math | |
| import random | |
| import cv2 | |
| import numpy as np | |
| def augment_hsv(im, hgain=0.5, sgain=0.5, vgain=0.5): | |
| # HSV color-space augmentation | |
| if hgain or sgain or vgain: | |
| r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1 # random gains | |
| hue, sat, val = cv2.split(cv2.cvtColor(im, cv2.COLOR_BGR2HSV)) | |
| dtype = im.dtype # uint8 | |
| x = np.arange(0, 256, dtype=r.dtype) | |
| lut_hue = ((x * r[0]) % 180).astype(dtype) | |
| lut_sat = np.clip(x * r[1], 0, 255).astype(dtype) | |
| lut_val = np.clip(x * r[2], 0, 255).astype(dtype) | |
| im_hsv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val))) | |
| cv2.cvtColor(im_hsv, cv2.COLOR_HSV2BGR, dst=im) # no return needed | |
| def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleup=True, stride=32): | |
| # Resize and pad image while meeting stride-multiple constraints | |
| shape = im.shape[:2] # current shape [height, width] | |
| if isinstance(new_shape, int): | |
| new_shape = (new_shape, new_shape) | |
| # Scale ratio (new / old) | |
| r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) | |
| if not scaleup: # only scale down, do not scale up (for better val mAP) | |
| r = min(r, 1.0) | |
| # Compute padding | |
| new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) | |
| dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding | |
| if auto: # minimum rectangle | |
| dw, dh = np.mod(dw, stride), np.mod(dh, stride) # wh padding | |
| dw /= 2 # divide padding into 2 sides | |
| dh /= 2 | |
| if shape[::-1] != new_unpad: # resize | |
| im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR) | |
| top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) | |
| left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) | |
| im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border | |
| return im, r, (dw, dh) | |
| def mixup(im, labels, im2, labels2): | |
| # Applies MixUp augmentation https://arxiv.org/pdf/1710.09412.pdf | |
| r = np.random.beta(32.0, 32.0) # mixup ratio, alpha=beta=32.0 | |
| im = (im * r + im2 * (1 - r)).astype(np.uint8) | |
| labels = np.concatenate((labels, labels2), 0) | |
| return im, labels | |
| def box_candidates(box1, box2, wh_thr=2, ar_thr=20, area_thr=0.1, eps=1e-16): # box1(4,n), box2(4,n) | |
| # Compute candidate boxes: box1 before augment, box2 after augment, wh_thr (pixels), aspect_ratio_thr, area_ratio | |
| w1, h1 = box1[2] - box1[0], box1[3] - box1[1] | |
| w2, h2 = box2[2] - box2[0], box2[3] - box2[1] | |
| ar = np.maximum(w2 / (h2 + eps), h2 / (w2 + eps)) # aspect ratio | |
| return (w2 > wh_thr) & (h2 > wh_thr) & (w2 * h2 / (w1 * h1 + eps) > area_thr) & (ar < ar_thr) # candidates | |
| def random_affine(img, labels=(), degrees=10, translate=.1, scale=.1, shear=10, | |
| new_shape=(640,640)): | |
| n = len(labels) | |
| height,width = new_shape | |
| M,s = get_transform_matrix(img.shape[:2],(height,width),degrees,scale,shear,translate) | |
| if (M != np.eye(3)).any(): # image changed | |
| img = cv2.warpAffine(img, M[:2], dsize=(width, height), borderValue=(114, 114, 114)) | |
| # Transform label coordinates | |
| if n: | |
| new = np.zeros((n, 4)) | |
| xy = np.ones((n * 4, 3)) | |
| xy[:, :2] = labels[:, [1, 2, 3, 4, 1, 4, 3, 2]].reshape(n * 4, 2) # x1y1, x2y2, x1y2, x2y1 | |
| xy = xy @ M.T # transform | |
| xy = xy[:, :2].reshape(n, 8) # perspective rescale or affine | |
| # create new boxes | |
| x = xy[:, [0, 2, 4, 6]] | |
| y = xy[:, [1, 3, 5, 7]] | |
| new = np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T | |
| # clip | |
| new[:, [0, 2]] = new[:, [0, 2]].clip(0, width) | |
| new[:, [1, 3]] = new[:, [1, 3]].clip(0, height) | |
| # filter candidates | |
| i = box_candidates(box1=labels[:, 1:5].T * s, box2=new.T, area_thr=0.1) | |
| labels = labels[i] | |
| labels[:, 1:5] = new[i] | |
| return img, labels | |
| def get_transform_matrix(img_shape,new_shape,degrees,scale,shear,translate): | |
| new_height,new_width = new_shape | |
| # Center | |
| C = np.eye(3) | |
| C[0, 2] = -img_shape[1] / 2 # x translation (pixels) | |
| C[1, 2] = -img_shape[0] / 2 # y translation (pixels) | |
| # Rotation and Scale | |
| R = np.eye(3) | |
| a = random.uniform(-degrees, degrees) | |
| # a += random.choice([-180, -90, 0, 90]) # add 90deg rotations to small rotations | |
| s = random.uniform(1 - scale, 1 + scale) | |
| # s = 2 ** random.uniform(-scale, scale) | |
| R[:2] = cv2.getRotationMatrix2D(angle=a, center=(0, 0), scale=s) | |
| # Shear | |
| S = np.eye(3) | |
| S[0, 1] = math.tan(random.uniform(-shear, shear) * math.pi / 180) # x shear (deg) | |
| S[1, 0] = math.tan(random.uniform(-shear, shear) * math.pi / 180) # y shear (deg) | |
| # Translation | |
| T = np.eye(3) | |
| T[0, 2] = random.uniform(0.5 - translate, 0.5 + translate) * new_width # x translation (pixels) | |
| T[1, 2] = random.uniform(0.5 - translate, 0.5 + translate) * new_height # y transla ion (pixels) | |
| # Combined rotation matrix | |
| M = T @ S @ R @ C # order of operations (right to left) is IMPORTANT | |
| return M,s | |
| def mosaic_augmentation(img_size, imgs, hs, ws, labels, hyp): | |
| assert len(imgs)==4, "Mosaic augmentaion of current version only supports 4 images." | |
| labels4 = [] | |
| s = img_size | |
| yc, xc = (int(random.uniform(s//2, 3*s//2)) for _ in range(2)) # mosaic center x, y | |
| for i in range(len(imgs)): | |
| # Load image | |
| img, h, w = imgs[i],hs[i],ws[i] | |
| # place img in img4 | |
| if i == 0: # top left | |
| img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8) # base image with 4 tiles | |
| x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc # xmin, ymin, xmax, ymax (large image) | |
| x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h # xmin, ymin, xmax, ymax (small image) | |
| elif i == 1: # top right | |
| x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc | |
| x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h | |
| elif i == 2: # bottom left | |
| x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h) | |
| x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, w, min(y2a - y1a, h) | |
| elif i == 3: # bottom right | |
| x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h) | |
| x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h) | |
| img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b] # img4[ymin:ymax, xmin:xmax] | |
| padw = x1a - x1b | |
| padh = y1a - y1b | |
| # Labels | |
| labels_per_img= labels[i].copy() | |
| if labels_per_img.size: | |
| boxes = np.copy(labels_per_img[:,1:]) | |
| boxes[:, 0] = w * (labels_per_img[:, 1] - labels_per_img[:, 3] / 2) + padw # top left x | |
| boxes[:, 1] = h * (labels_per_img[:, 2] - labels_per_img[:, 4] / 2) + padh # top left y | |
| boxes[:, 2] = w * (labels_per_img[:, 1] + labels_per_img[:, 3] / 2) + padw # bottom right x | |
| boxes[:, 3] = h * (labels_per_img[:, 2] + labels_per_img[:, 4] / 2) + padh # bottom right y | |
| labels_per_img[:,1:] = boxes | |
| labels4.append(labels_per_img) | |
| # Concat/clip labels | |
| labels4 = np.concatenate(labels4, 0) | |
| for x in (labels4[:, 1:]): | |
| np.clip(x, 0, 2 * s, out=x) | |
| # Augment | |
| img4, labels4 = random_affine(img4, labels4, | |
| degrees=hyp['degrees'], | |
| translate=hyp['translate'], | |
| scale=hyp['scale'], | |
| shear=hyp['shear']) | |
| return img4, labels4 | |