Spaces:
Running
Running
| import cv2 | |
| import numpy as np | |
| import random | |
| from utils.box_utils import matrix_iof | |
| def _crop(image, boxes, labels, landm, img_dim): | |
| height, width, _ = image.shape | |
| pad_image_flag = True | |
| for _ in range(250): | |
| """ | |
| if random.uniform(0, 1) <= 0.2: | |
| scale = 1.0 | |
| else: | |
| scale = random.uniform(0.3, 1.0) | |
| """ | |
| PRE_SCALES = [0.3, 0.45, 0.6, 0.8, 1.0] | |
| scale = random.choice(PRE_SCALES) | |
| short_side = min(width, height) | |
| w = int(scale * short_side) | |
| h = w | |
| if width == w: | |
| l = 0 | |
| else: | |
| l = random.randrange(width - w) | |
| if height == h: | |
| t = 0 | |
| else: | |
| t = random.randrange(height - h) | |
| roi = np.array((l, t, l + w, t + h)) | |
| value = matrix_iof(boxes, roi[np.newaxis]) | |
| flag = (value >= 1) | |
| if not flag.any(): | |
| continue | |
| centers = (boxes[:, :2] + boxes[:, 2:]) / 2 | |
| mask_a = np.logical_and(roi[:2] < centers, centers < roi[2:]).all(axis=1) | |
| boxes_t = boxes[mask_a].copy() | |
| labels_t = labels[mask_a].copy() | |
| landms_t = landm[mask_a].copy() | |
| landms_t = landms_t.reshape([-1, 5, 2]) | |
| if boxes_t.shape[0] == 0: | |
| continue | |
| image_t = image[roi[1]:roi[3], roi[0]:roi[2]] | |
| boxes_t[:, :2] = np.maximum(boxes_t[:, :2], roi[:2]) | |
| boxes_t[:, :2] -= roi[:2] | |
| boxes_t[:, 2:] = np.minimum(boxes_t[:, 2:], roi[2:]) | |
| boxes_t[:, 2:] -= roi[:2] | |
| # landm | |
| landms_t[:, :, :2] = landms_t[:, :, :2] - roi[:2] | |
| landms_t[:, :, :2] = np.maximum(landms_t[:, :, :2], np.array([0, 0])) | |
| landms_t[:, :, :2] = np.minimum(landms_t[:, :, :2], roi[2:] - roi[:2]) | |
| landms_t = landms_t.reshape([-1, 10]) | |
| # make sure that the cropped image contains at least one face > 16 pixel at training image scale | |
| b_w_t = (boxes_t[:, 2] - boxes_t[:, 0] + 1) / w * img_dim | |
| b_h_t = (boxes_t[:, 3] - boxes_t[:, 1] + 1) / h * img_dim | |
| mask_b = np.minimum(b_w_t, b_h_t) > 0.0 | |
| boxes_t = boxes_t[mask_b] | |
| labels_t = labels_t[mask_b] | |
| landms_t = landms_t[mask_b] | |
| if boxes_t.shape[0] == 0: | |
| continue | |
| pad_image_flag = False | |
| return image_t, boxes_t, labels_t, landms_t, pad_image_flag | |
| return image, boxes, labels, landm, pad_image_flag | |
| def _distort(image): | |
| def _convert(image, alpha=1, beta=0): | |
| tmp = image.astype(float) * alpha + beta | |
| tmp[tmp < 0] = 0 | |
| tmp[tmp > 255] = 255 | |
| image[:] = tmp | |
| image = image.copy() | |
| if random.randrange(2): | |
| #brightness distortion | |
| if random.randrange(2): | |
| _convert(image, beta=random.uniform(-32, 32)) | |
| #contrast distortion | |
| if random.randrange(2): | |
| _convert(image, alpha=random.uniform(0.5, 1.5)) | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) | |
| #saturation distortion | |
| if random.randrange(2): | |
| _convert(image[:, :, 1], alpha=random.uniform(0.5, 1.5)) | |
| #hue distortion | |
| if random.randrange(2): | |
| tmp = image[:, :, 0].astype(int) + random.randint(-18, 18) | |
| tmp %= 180 | |
| image[:, :, 0] = tmp | |
| image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR) | |
| else: | |
| #brightness distortion | |
| if random.randrange(2): | |
| _convert(image, beta=random.uniform(-32, 32)) | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) | |
| #saturation distortion | |
| if random.randrange(2): | |
| _convert(image[:, :, 1], alpha=random.uniform(0.5, 1.5)) | |
| #hue distortion | |
| if random.randrange(2): | |
| tmp = image[:, :, 0].astype(int) + random.randint(-18, 18) | |
| tmp %= 180 | |
| image[:, :, 0] = tmp | |
| image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR) | |
| #contrast distortion | |
| if random.randrange(2): | |
| _convert(image, alpha=random.uniform(0.5, 1.5)) | |
| return image | |
| def _expand(image, boxes, fill, p): | |
| if random.randrange(2): | |
| return image, boxes | |
| height, width, depth = image.shape | |
| scale = random.uniform(1, p) | |
| w = int(scale * width) | |
| h = int(scale * height) | |
| left = random.randint(0, w - width) | |
| top = random.randint(0, h - height) | |
| boxes_t = boxes.copy() | |
| boxes_t[:, :2] += (left, top) | |
| boxes_t[:, 2:] += (left, top) | |
| expand_image = np.empty( | |
| (h, w, depth), | |
| dtype=image.dtype) | |
| expand_image[:, :] = fill | |
| expand_image[top:top + height, left:left + width] = image | |
| image = expand_image | |
| return image, boxes_t | |
| def _mirror(image, boxes, landms): | |
| _, width, _ = image.shape | |
| if random.randrange(2): | |
| image = image[:, ::-1] | |
| boxes = boxes.copy() | |
| boxes[:, 0::2] = width - boxes[:, 2::-2] | |
| # landm | |
| landms = landms.copy() | |
| landms = landms.reshape([-1, 5, 2]) | |
| landms[:, :, 0] = width - landms[:, :, 0] | |
| tmp = landms[:, 1, :].copy() | |
| landms[:, 1, :] = landms[:, 0, :] | |
| landms[:, 0, :] = tmp | |
| tmp1 = landms[:, 4, :].copy() | |
| landms[:, 4, :] = landms[:, 3, :] | |
| landms[:, 3, :] = tmp1 | |
| landms = landms.reshape([-1, 10]) | |
| return image, boxes, landms | |
| def _pad_to_square(image, rgb_mean, pad_image_flag): | |
| if not pad_image_flag: | |
| return image | |
| height, width, _ = image.shape | |
| long_side = max(width, height) | |
| image_t = np.empty((long_side, long_side, 3), dtype=image.dtype) | |
| image_t[:, :] = rgb_mean | |
| image_t[0:0 + height, 0:0 + width] = image | |
| return image_t | |
| def _resize_subtract_mean(image, insize, rgb_mean): | |
| interp_methods = [cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_NEAREST, cv2.INTER_LANCZOS4] | |
| interp_method = interp_methods[random.randrange(5)] | |
| image = cv2.resize(image, (insize, insize), interpolation=interp_method) | |
| image = image.astype(np.float32) | |
| image -= rgb_mean | |
| return image.transpose(2, 0, 1) | |
| class preproc(object): | |
| def __init__(self, img_dim, rgb_means): | |
| self.img_dim = img_dim | |
| self.rgb_means = rgb_means | |
| def __call__(self, image, targets): | |
| assert targets.shape[0] > 0, "this image does not have gt" | |
| boxes = targets[:, :4].copy() | |
| labels = targets[:, -1].copy() | |
| landm = targets[:, 4:-1].copy() | |
| image_t, boxes_t, labels_t, landm_t, pad_image_flag = _crop(image, boxes, labels, landm, self.img_dim) | |
| image_t = _distort(image_t) | |
| image_t = _pad_to_square(image_t,self.rgb_means, pad_image_flag) | |
| image_t, boxes_t, landm_t = _mirror(image_t, boxes_t, landm_t) | |
| height, width, _ = image_t.shape | |
| image_t = _resize_subtract_mean(image_t, self.img_dim, self.rgb_means) | |
| boxes_t[:, 0::2] /= width | |
| boxes_t[:, 1::2] /= height | |
| landm_t[:, 0::2] /= width | |
| landm_t[:, 1::2] /= height | |
| labels_t = np.expand_dims(labels_t, 1) | |
| targets_t = np.hstack((boxes_t, landm_t, labels_t)) | |
| return image_t, targets_t | |