数据增强--对于目标检测（边框、图像）

import torch
from PIL import Image, ImageFont, ImageDraw
from functools import reduce
import scipy.io as scio
from PIL import Image
import cv2 as cv
import numpy as np
import random
import imutils
from imgaug import augmenters as iaa


def horisontal_flip(images, targets):
    images = torch.flip(images, [-1])
    targets[:, 2] = 1 - targets[:, 2]
    return images, targets


def compose(*funcs):
    """Compose arbitrarily many functions, evaluated left to right.

    Reference: https://mathieularose.com/function-composition-in-python/
    """
    # return lambda x: reduce(lambda v, f: f(v), funcs, x)
    if funcs:
        return reduce(lambda f, g: lambda *a, **kw: g(f(*a, **kw)), funcs)
    else:
        raise ValueError('Composition of empty sequence not supported.')


def letterbox_image(image, size):
    '''resize image with unchanged aspect ratio using padding'''
    iw, ih = image.size
    w, h = size
    scale = min(w / iw, h / ih)
    nw = int(iw * scale)
    nh = int(ih * scale)

    image = image.resize((nw, nh), Image.BICUBIC)
    new_image = Image.new('RGB', size, (128, 128, 128))
    new_image.paste(image, ((w - nw) // 2, (h - nh) // 2))
    return new_image


def rand(a=0, b=1):
    return np.random.rand() * (b - a) + a


def resize_3D_data(ima, size=(416, 416)):
    dat = np.zeros([416, 416, 15])
    for i in range(15):
        slice = np.squeeze(ima[:, :, i])
        re_slice = cv.resize(slice.astype('uint8'), size, interpolation=cv.INTER_AREA)
        dat[:, :, i] = re_slice
    return dat


def random_crop(cell, boxes):
    dx = random.randint(15, 20)
    dy = random.randint(15, 20)
    shape = cell.shape
    nx = shape[0]
    ny = shape[1]
    boxes[:, [0, 2]] = boxes[:, [0, 2]] - dx
    boxes[:, [1, 3]] = boxes[:, [1, 3]] - dy
    new_cell = np.zeros_like(cell)
    new_cell[0:int(nx - dx), 0:int(ny - dy)] = cell[dx:, dy:]
    new_cell = Image.fromarray(new_cell.astype('uint8')).convert('RGB')
    new_cell = cv.cvtColor(np.asarray(new_cell), cv.COLOR_RGB2BGR)
    boxes = np.where(boxes < 0, 0, boxes)
    return new_cell, boxes


def random_noise(cell, boxes):
    image = Image.fromarray(cell.astype('uint8')).convert('RGB')
    image = cv.cvtColor(np.asarray(image), cv.COLOR_RGB2BGR)
    seq = iaa.Sequential(
        [
            iaa.AdditiveGaussianNoise(scale=0.05 * 255),
            iaa.LinearContrast((0.75, 1.5)),
            iaa.GaussianBlur(sigma=(0, 4.0)),
            iaa.Dropout(p=(0, 0.2)),
            iaa.CoarseDropout(0.02, size_percent=0.5)
        ], random_order=True
    )
    images_aug = seq.augment_images([image])[0]
    new_cell = Image.fromarray(cv.cvtColor(images_aug, cv.COLOR_BGR2RGB))
    return new_cell, boxes


def gray_level_crop(cell, boxes):
    max_val = random.randint(80, 255) / 255
    new_cell = cell * max_val
    new_cell = Image.fromarray(new_cell.astype('uint8')).convert('RGB')
    new_cell = cv.cvtColor(np.asarray(new_cell), cv.COLOR_RGB2BGR)
    return new_cell, boxes


def rotate_box(box, M, shape):
    # print(box)
    y1, x1, y2, x2 = box
    p1 = np.array([x1, y1, 1]).reshape((3, 1))
    p2 = np.array([x1, y2, 1]).reshape((3, 1))
    p3 = np.array([x2, y2, 1]).reshape((3, 1))
    p4 = np.array([x2, y1, 1]).reshape((3, 1))
    p1 = np.matmul(M, p1)
    p2 = np.matmul(M, p2)
    p3 = np.matmul(M, p3)
    p4 = np.matmul(M, p4)
    x1 = np.min([p1[0, 0], p2[0, 0], p3[0, 0], p4[0, 0]])
    x2 = np.max([p1[0, 0], p2[0, 0], p3[0, 0], p4[0, 0]])
    y1 = np.min([p1[1, 0], p2[1, 0], p3[1, 0], p4[1, 0]])
    y2 = np.max([p1[1, 0], p2[1, 0], p3[1, 0], p4[1, 0]])
    if x1 < 0:
        x1 = 0
    if x1 > shape[1]:
        x1 = shape[1] - 1
    if x2 < 0:
        x2 = 0
    if x2 > shape[1]:
        x2 = shape[1] - 1

    if y1 < 0:
        y1 = 0
    if y1 > shape[0]:
        y1 = shape[0] - 1
    if y2 < 0:
        y2 = 0
    if y2 > shape[0]:
        y2 = shape[0] - 1

    box = [y1, x1, y2, x2]
    # print(box)
    # print('--------------')
    return box


def random_rotate(cell, boxes, angle=45):
    (h, w) = cell.shape
    (cX, cY) = (w // 2, h // 2)
    new_cell = imutils.rotate_bound(cell.astype('uint8'), angle)
    M = cv.getRotationMatrix2D((cX, cY), -angle, 1.0)
    cos = np.abs(M[0, 0])
    sin = np.abs(M[0, 1])

    # compute the new bounding dimensions of the image
    nW = int((h * sin) + (w * cos))
    nH = int((h * cos) + (w * sin))

    # adjust the rotation matrix to take into account translation
    M[0, 2] += (nW / 2) - cX
    M[1, 2] += (nH / 2) - cY
    new_boxes = []
    for i in range(len(boxes)):
        new_boxes.append(rotate_box(boxes[i], M, new_cell.shape))
    if len(new_boxes) > 0:
        new_boxes = np.array(new_boxes)
    new_cell = Image.fromarray(new_cell.astype('uint8')).convert('RGB')
    new_cell = cv.cvtColor(np.asarray(new_cell), cv.COLOR_RGB2BGR)
    return new_cell, new_boxes


def random_argument(data, cell, boxes):
    # crop region ...
    if rand() < .5:
        cell, boxes = random_crop(cell, boxes)
    # add noise ...
    if rand() < .5:
        cell, boxes = random_noise(cell, boxes)
    # add noise ...
    if rand() < .5:
        cell, boxes = gray_level_crop(cell, boxes)
    # rotate ...
    if rand() < .5:
       cell, boxes = random_rotate(cell, boxes, random.randint(0, 90))

    return cell, boxes


def pad_to_square(img, pad_value):
    h, w, c = img.shape
    dim_diff = np.abs(h - w)
    # (upper / left) padding and (lower / right) padding
    pad1, pad2 = dim_diff // 2, dim_diff - dim_diff // 2
    # Determine padding
    pad = ((0, 0), (pad1, pad2), (0, 0)) if w <= h else ((pad1, pad2), (0, 0), (0, 0))
    # Add padding
    img = np.pad(img, pad, "constant", constant_values=pad_value)

    return img, pad


if __name__ == '__main__':
    dat = scio.loadmat(
        '/home/xuxu/Data/CTC_Signal_Datasets/Processed/20200616_cell_signalmark_39frames/Green/653802_650142_0_675.mat')
    cell = dat['signal'].astype('uint8')
    boxe = dat['rects']
    print('boxe', boxe)
    # image = Image.fromarray(cell.astype('uint8')).convert('RGB')
    # image = cv.cvtColor(np.asarray(image), cv.COLOR_RGB2BGR)
    # for x1, y1, x2, y2 in boxes:
    #     cv.rectangle(image, (y1, x1), (y2, x2), (255, 255, 255), thickness=1)
    # cv.imwrite('1.jpg', image)
    new_cell, new_boxes = random_rotate(cell, boxe)
    # new_cell, new_boxes = gray_level_crop(cell, boxe)
    # print('new_boxes', new_boxes)
    # new_boxes = np.where(new_boxes < 0, 0, new_boxes)
    # print('new_boxes', new_boxes)
    # image = Image.fromarray(new_cell.astype('uint8')).convert('RGB')
    # image = cv.cvtColor(np.asarray(image), cv.COLOR_RGB2BGR)
    for x1, y1, x2, y2 in new_boxes:
        cv.rectangle(new_cell, (int(y1), int(x1)), (int(y2), int(x2)), (255, 255, 255), thickness=1)
    cv.imwrite('1.jpg', new_cell)

其中图像加噪使用的是imgaug库