retinaface 检查数据

import cv2
import numpy as np
import random
from utils.box_utils import matrix_iof


def _crop(image, boxes, labels, img_size):
    #随意裁剪，裁剪到最后resize，会有放大效果，不会缩小，小目标检测无帮助
    height, width, _ = image.shape
    pad_image_flag = True

    if boxes[0][3] == 0 or boxes[0][2] - boxes[0][0] == 0 or boxes[0][3] - boxes[0][1] == 0:
        scale = random.uniform(0.3, 1.0)
        short_side = min(width, height)
        crop_width = int(scale * short_side)

        crop_left = 0
        if width != crop_width:
            crop_left = random.randrange(width - crop_width)
        crop_top = 0
        if height != crop_width:
            crop_top = random.randrange(height - crop_width)
        roi = np.array((crop_left, crop_top, crop_left + crop_width, crop_top + crop_width))

        image_t = image[roi[1]:roi[3], roi[0]:roi[2]]
        pad_image_flag = False
        return image_t, boxes, labels, pad_image_flag
    for _ in range(250):

        scale = random.uniform(0.3, 1.0)
        short_side = min(width, height)
        crop_width = int(scale * short_side)

        crop_left = 0
        if width != crop_width:
            crop_left = random.randrange(width - crop_width)
        crop_top = 0
        if height != crop_width:
            crop_top = random.randrange(height - crop_width)

        roi = np.array((crop_left, crop_top, crop_left + crop_width, crop_top + crop_width))
        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) / crop_width * img_size
        b_h_t = (boxes_t[:, 3] - boxes_t[:, 1] + 1) / crop_width * img_size
        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_t, boxes_t, labels_t, pad_image_flag
    return image, boxes, labels, 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):
    _, 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


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_size, rgb_means):
        self.img_size = img_size
        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, pad_image_flag = _crop(image, boxes, labels, self.img_size)
        # image_t, boxes_t, labels_t, landm_t, pad_image_flag = _crop(image, boxes, labels, landm, self.img_dim)

        for lable in boxes_t:
            cv2.rectangle(image_t, (int(lable[0]), int(lable[1])), (int(lable[2]), int(lable[3])),
                          (0, 0, 255), 1)
        cv2.imshow("asdf",image_t)
        cv2.waitKey()
        image_t = _distort(image_t)
        image_t = _pad_to_square(image_t,self.rgb_means, pad_image_flag)
        image_t, boxes_t = _mirror(image_t, boxes_t)
        height, width, _ = image_t.shape
        image_t = _resize_subtract_mean(image_t, self.img_size, self.rgb_means)
        boxes_t[:, 0::2] /= width
        boxes_t[:, 1::2] /= height



        labels_t = np.expand_dims(labels_t, 1)
        targets_t = np.hstack((boxes_t, labels_t))

        return image_t, targets_t