# -*- coding=utf-8 -*-
##############################################################
# description:
#     data augmentation for obeject detection
##############################################################

# 包括:
#     1. 裁剪(需改变bbox)
#     2. 平移(需改变bbox)
#     3. 改变亮度
#     4. 加噪声
#     5. 旋转角度(需要改变bbox)
#     6. 镜像(需要改变bbox)
#     7. cutout
# 注意:
#     random.seed(),相同的seed,产生的随机数是一样的!!

import time
import random
import cv2
import os
import math
import numpy as np
from skimage.util import random_noise
from skimage import exposure
from preprocessing.xml_h import *


def show_pic(img, bboxes=None, class_name=None):
    '''
    输入:
        img:图像array
        bboxes:图像的所有boudning box list, 格式为[[x_min, y_min, x_max, y_max]....]
        names:每个box对应的名称
    输出：
        画好标注信息的图像
    '''

    cv2.imwrite('./1.jpg', img)
    img = cv2.imread('./1.jpg')
    font = cv2.FONT_HERSHEY_COMPLEX_SMALL
    # print(len(bboxes))
    for i in range(len(bboxes)):
        bbox = bboxes[i]
        classname = class_name[i]
        if classname == 'pos':
            r = 0
            g = 255
            b = 0
        elif classname == 'neg':
            r = 0
            g = 0
            b = 255
        x_min = bbox[0]
        y_min = bbox[1]
        x_max = bbox[2]
        y_max = bbox[3]
        cv2.circle(img, (int(bbox[4]), int(bbox[5])), 2, (r, g, b), -1)  # circle（图像array， 圆心坐标， 半径， 颜色， 是否实心'-1'表示实心）
        cv2.rectangle(img, (int(x_min), int(y_min)), (int(x_max), int(y_max)), (r, g, b), 1)
        cv2.putText(img, classname, (int(x_min), int(y_min - 2)), fontFace=font, fontScale=0.8, thickness=1,
                    color=(r, g, b))
        cv2.rectangle(img, (int(x_min), int(y_min)), (int(x_max), int(y_max)), (r, g, b), 1)
    return img


# 图像均为cv2读取
class DataAugmentForObjectDetection():
    def __init__(self, rotation_rate=0.2, max_rotation_angle=5,
                 crop_rate=0.5, shift_rate=0.6, change_light_rate=0.5,
                 add_noise_rate=0.5, flip_rate=0.5,
                 cutout_rate=0.5, cut_out_length=50, cut_out_holes=1, cut_out_threshold=0.5):
        self.rotation_rate = rotation_rate
        self.max_rotation_angle = max_rotation_angle
        self.crop_rate = crop_rate
        self.shift_rate = shift_rate
        self.change_light_rate = change_light_rate
        self.add_noise_rate = add_noise_rate
        self.flip_rate = flip_rate
        self.cutout_rate = cutout_rate

        self.cut_out_length = cut_out_length
        self.cut_out_holes = cut_out_holes
        self.cut_out_threshold = cut_out_threshold

    # 加噪声
    def _addNoise(self, img):
        '''
        输入:
            img:图像array
        输出:
            加噪声后的图像array,由于输出的像素是在[0,1]之间,所以得乘以255
        '''
        # random.seed(int(time.time()))
        # return random_noise(img, mode='gaussian', seed=int(time.time()), clip=True)*255
        return random_noise(img, mode='gaussian', clip=True) * 255

    # 调整亮度
    def _changeLight(self, img):
        # random.seed(int(time.time()))
        shape = img.shape
        for m in range(shape[0]):
            for n in range(shape[1]):
                if np.any(img[m][n] < 0):
                    img[m][n][:] = 0
                    print('error occur in picture: %s' % i)
        flag = random.uniform(0.5, 1.5)  # flag>1为调暗,小于1为调亮
        return exposure.adjust_gamma(img, flag)

    # cutout
    def _cutout(self, img, bboxes, length=100, n_holes=1, threshold=0.5):
        '''
        原版本：https://github.com/uoguelph-mlrg/Cutout/blob/master/util/cutout.py
        Randomly mask out one or more patches from an image.
        Args:
            img : a 3D numpy array,(h,w,c)
            bboxes : 框的坐标
            n_holes (int): Number of patches to cut out of each image.
            length (int): The length (in pixels) of each square patch.
        '''

        def cal_iou(boxA, boxB):
            '''
            boxA, boxB为两个框，返回iou
            boxB为bouding box
            '''

            # determine the (x, y)-coordinates of the intersection rectangle
            xA = max(boxA[0], boxB[0])
            yA = max(boxA[1], boxB[1])
            xB = min(boxA[2], boxB[2])
            yB = min(boxA[3], boxB[3])

            if xB <= xA or yB <= yA:
                return 0.0

            # compute the area of intersection rectangle
            interArea = (xB - xA + 1) * (yB - yA + 1)

            # compute the area of both the prediction and ground-truth
            # rectangles
            boxAArea = (boxA[2] - boxA[0] + 1) * (boxA[3] - boxA[1] + 1)
            boxBArea = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1)

            # compute the intersection over union by taking the intersection
            # area and dividing it by the sum of prediction + ground-truth
            # areas - the interesection area
            # iou = interArea / float(boxAArea + boxBArea - interArea)
            iou = interArea / float(boxBArea)

            # return the intersection over union value
            return iou

        # 得到h和w
        if img.ndim == 3:
            h, w, c = img.shape
        else:
            _, h, w, c = img.shape

        mask = np.ones((h, w, c), np.float32)

        for n in range(n_holes):

            chongdie = True  # 看切割的区域是否与box重叠太多

            while chongdie:
                y = np.random.randint(h)
                x = np.random.randint(w)

                y1 = np.clip(y - length // 2, 0,
                             h)  # numpy.clip(a, a_min, a_max, out=None), clip这个函数将将数组中的元素限制在a_min, a_max之间，大于a_max的就使得它等于 a_max，小于a_min,的就使得它等于a_min
                y2 = np.clip(y + length // 2, 0, h)
                x1 = np.clip(x - length // 2, 0, w)
                x2 = np.clip(x + length // 2, 0, w)

                chongdie = False
                for box in bboxes:
                    if cal_iou([x1, y1, x2, y2], box) > threshold:
                        chongdie = True
                        break

            mask[y1: y2, x1: x2, :] = 0.

        # mask = np.expand_dims(mask, axis=0)
        img = img * mask

        return img

    # 旋转
    def _rotate_img_bbox(self, img, bboxes, angle=5, scale=1.):
        '''
        参考:https://blog.csdn.net/u014540717/article/details/53301195crop_rate
        输入:
            img:图像array,(h,w,c)
            bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
            angle:旋转角度
            scale:默认1
        输出:
            rot_img:旋转后的图像array
            rot_bboxes:旋转后的boundingbox坐标list
        '''
        # ---------------------- 旋转图像 ----------------------
        w = img.shape[1]
        h = img.shape[0]
        # 角度变弧度
        rangle = np.deg2rad(angle)  # angle in radians
        # now calculate new image width and height
        nw = (abs(np.sin(rangle) * h) + abs(np.cos(rangle) * w)) * scale
        nh = (abs(np.cos(rangle) * h) + abs(np.sin(rangle) * w)) * scale
        # ask OpenCV for the rotation matrix

        rot_mat = cv2.getRotationMatrix2D((nw * 0.5, nh * 0.5), angle, scale)
        # calculate the move from the old center to the new center combined
        # with the rotation
        rot_move = np.dot(rot_mat, np.array([(nw - w) * 0.5, (nh - h) * 0.5, 0]))
        # the move only affects the translation, so update the translation
        # part of the transform
        rot_mat[0, 2] += rot_move[0]
        rot_mat[1, 2] += rot_move[1]
        # 仿射变换
        rot_img = cv2.warpAffine(img, rot_mat, (int(math.ceil(nw)), int(math.ceil(nh))), flags=cv2.INTER_LANCZOS4)

        # ---------------------- 矫正bbox坐标 ----------------------
        # rot_mat是最终的旋转矩阵
        # 获取原始bbox的四个中点，然后将这四个点转换到旋转后的坐标系下
        rot_bboxes = list()
        for bbox in bboxes:
            xmin = bbox[0]
            ymin = bbox[1]
            xmax = bbox[2]
            ymax = bbox[3]
            c_x = bbox[4]
            c_y = bbox[5]
            point1 = np.dot(rot_mat, np.array([(xmin + xmax) / 2, ymin, 1]))
            point2 = np.dot(rot_mat, np.array([xmax, (ymin + ymax) / 2, 1]))
            point3 = np.dot(rot_mat, np.array([(xmin + xmax) / 2, ymax, 1]))
            point4 = np.dot(rot_mat, np.array([xmin, (ymin + ymax) / 2, 1]))
            point5 = np.dot(rot_mat, np.array([c_x, c_y, 1]))
            # 合并np.array
            concat = np.vstack((point1, point2, point3, point4))
            # 改变array类型
            concat = concat.astype(np.int32)
            # 得到旋转后的坐标
            rx, ry, rw, rh = cv2.boundingRect(concat)
            rx_min = rx
            ry_min = ry
            rx_max = rx + rw
            ry_max = ry + rh
            c_x = int(point5[0])
            c_y = int(point5[1])
            # 加入list中
            if bbox[6] == 'pos':
                rot_bboxes.append([rx_min, ry_min, rx_max, ry_max, c_x, c_y, bbox[6]])
            elif bbox[6] == 'neg':
                rot_bboxes.append([rx_min, ry_min, rx_max, ry_max, 0, 0, bbox[6]])
        return rot_img, rot_bboxes

    # 裁剪
    def _crop_img_bboxes(self, img, bboxes):
        '''
        裁剪后的图片要包含所有的框
        输入:
            img:图像array
            bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
        输出:
            crop_img:裁剪后的图像array
            crop_bboxes:裁剪后的bounding box的坐标list
        '''
        # ---------------------- 裁剪图像 ----------------------
        w = img.shape[1]
        h = img.shape[0]
        x_min = w  # 裁剪后的包含所有目标框的最小的框
        x_max = 0
        y_min = h
        y_max = 0
        for bbox in bboxes:
            x_min = min(x_min, bbox[0])
            y_min = min(y_min, bbox[1])
            x_max = max(x_max, bbox[2])
            y_max = max(y_max, bbox[3])

        d_to_left = x_min  # 包含所有目标框的最小框到左边的距离
        d_to_right = w - x_max  # 包含所有目标框的最小框到右边的距离
        d_to_top = y_min  # 包含所有目标框的最小框到顶端的距离
        d_to_bottom = h - y_max  # 包含所有目标框的最小框到底部的距离

        # 随机扩展这个最小框
        crop_x_min = int(x_min - random.uniform(0, d_to_left))
        crop_y_min = int(y_min - random.uniform(0, d_to_top))
        crop_x_max = int(x_max + random.uniform(0, d_to_right))
        crop_y_max = int(y_max + random.uniform(0, d_to_bottom))

        # 随机扩展这个最小框 , 防止别裁的太小
        # crop_x_min = int(x_min - random.uniform(d_to_left//2, d_to_left))
        # crop_y_min = int(y_min - random.uniform(d_to_top//2, d_to_top))
        # crop_x_max = int(x_max + random.uniform(d_to_right//2, d_to_right))
        # crop_y_max = int(y_max + random.uniform(d_to_bottom//2, d_to_bottom))

        # 确保不要越界
        crop_x_min = max(0, crop_x_min)
        crop_y_min = max(0, crop_y_min)
        crop_x_max = min(w, crop_x_max)
        crop_y_max = min(h, crop_y_max)

        crop_img = img[crop_y_min:crop_y_max, crop_x_min:crop_x_max]

        # ---------------------- 裁剪boundingbox ----------------------
        # 裁剪后的boundingbox坐标计算
        crop_bboxes = list()
        for bbox in bboxes:
            if bbox[6] == 'pos':
                crop_bboxes.append([bbox[0] - crop_x_min, bbox[1] - crop_y_min, bbox[2] - crop_x_min, bbox[3] - crop_y_min,
                               bbox[4] - crop_x_min, bbox[5] - crop_y_min, bbox[6]])
            elif bbox[6] == 'neg':
                crop_bboxes.append([bbox[0] - crop_x_min, bbox[1] - crop_y_min, bbox[2] - crop_x_min, bbox[3] - crop_y_min, 0, 0, bbox[6]])
        return crop_img, crop_bboxes

    # 平移
    def _shift_pic_bboxes(self, img, bboxes):
        '''
        参考:https://blog.csdn.net/sty945/article/details/79387054
        平移后的图片要包含所有的框
        输入:
            img:图像array
            bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
        输出:
            shift_img:平移后的图像array
            shift_bboxes:平移后的bounding box的坐标list
        '''
        # ---------------------- 平移图像 ----------------------
        w = img.shape[1]
        h = img.shape[0]
        x_min = w  # 裁剪后的包含所有目标框的最小的框
        x_max = 0
        y_min = h
        y_max = 0
        for bbox in bboxes:
            x_min = min(x_min, bbox[0])
            y_min = min(y_min, bbox[1])
            x_max = max(x_max, bbox[2])
            y_max = max(y_max, bbox[3])
            c_x = bbox[4]
            c_y = bbox[5]
        d_to_left = x_min  # 包含所有目标框的最大左移动距离
        d_to_right = w - x_max  # 包含所有目标框的最大右移动距离
        d_to_top = y_min  # 包含所有目标框的最大上移动距离
        d_to_bottom = h - y_max  # 包含所有目标框的最大下移动距离

        x = random.uniform(-(d_to_left - 1) / 3, (d_to_right - 1) / 3)
        y = random.uniform(-(d_to_top - 1) / 3, (d_to_bottom - 1) / 3)

        M = np.float32([[1, 0, x], [0, 1, y]])  # x为向左或右移动的像素值,正为向右负为向左; y为向上或者向下移动的像素值,正为向下负为向上
        shift_img = cv2.warpAffine(img, M, (img.shape[1], img.shape[0]))

        # ---------------------- 平移boundingbox ----------------------
        shift_bboxes = list()
        for bbox in bboxes:
            if bbox[6] == 'pos':
                shift_bboxes.append([int(bbox[0] + x), int(bbox[1] + y), int(bbox[2] + x), int(bbox[3] + y), int(bbox[4] + x), int(bbox[5] + y), bbox[6]])
            elif bbox[6] == 'neg':
                shift_bboxes.append([int(bbox[0] + x), int(bbox[1] + y), int(bbox[2] + x), int(bbox[3] + y), 0, 0, bbox[6]])
        return shift_img, shift_bboxes

    # 镜像
    def _filp_pic_bboxes(self, img, bboxes):
        '''
            参考:https://blog.csdn.net/jningwei/article/details/78753607
            平移后的图片要包含所有的框
            输入:
                img:图像array
                bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
            输出:
                flip_img:平移后的图像array
                flip_bboxes:平移后的bounding box的坐标list
        '''
        # ---------------------- 翻转图像 ----------------------
        import copy
        flip_img = copy.deepcopy(img)
        if random.random() < 0.5:  # 0.5的概率水平翻转，0.5的概率垂直翻转
            horizon = True
        else:
            horizon = False
        h, w, _ = img.shape
        if horizon:  # 水平翻转
            flip_img = cv2.flip(flip_img, 1)
        else:
            flip_img = cv2.flip(flip_img, 0)

        # ---------------------- 调整boundingbox ----------------------
        flip_bboxes = list()
        for box in bboxes:
            x_min = box[0]
            y_min = box[1]
            x_max = box[2]
            y_max = box[3]
            c_x = box[4]
            c_y = box[5]
            if box[6] == 'pos':
                if horizon:
                    flip_bboxes.append([w - x_max, y_min, w - x_min, y_max, w - c_x, c_y, box[6]])
                else:
                    flip_bboxes.append([x_min, h - y_max, x_max, h - y_min, c_x, h - c_y, box[6]])
            elif box[6] == 'neg':
                if horizon:
                    flip_bboxes.append([w - x_max, y_min, w - x_min, y_max, 0, 0, box[6]])
                else:
                    flip_bboxes.append([x_min, h - y_max, x_max, h - y_min, 0, 0, box[6]])

        return flip_img, flip_bboxes
    
    # 方法汇总
    def methods(self, num, img, bboxes):
        method = {
            0: self._crop_img_bboxes(img, bboxes),
            1: self._shift_pic_bboxes(img, bboxes),
            2: self._changeLight(img),
            3: self._addNoise(img),
            4: self._cutout(img, bboxes, length=self.cut_out_length, n_holes=self.cut_out_holes,
                            threshold=self.cut_out_threshold),
            5: self._filp_pic_bboxes(img, bboxes)
        }
        if num == 2 or num == 3 or num == 4:
            img = method[num]
        else:
            img, bboxes = method[num]
        return img, bboxes

    def dataAugment(self, img, bboxes):
        '''
        图像增强
        输入:
            img:图像array
            bboxes:该图像的所有框坐标
        输出:
            img:增强后的图像
            bboxes:增强后图片对应的box
        '''
        change_num = 0  # 改变的次数
        print('------')
        while change_num < 2:  # 默认至少有一种数据增强生效
            if random.random() < 0.4:  # 裁剪

                print('裁剪')
                change_num += 1
                img, bboxes = self._crop_img_bboxes(img, bboxes)

            if random.random() > 0.8:  # 旋转
                print('旋转')
                change_num += 1
                # angle = random.uniform(-self.max_rotation_angle, self.max_rotation_angle)
                angle = random.sample([90, 180, 270], 1)[0]
                scale = random.uniform(0.7, 0.8)
                img, bboxes = self._rotate_img_bbox(img, bboxes, angle, scale)

            if 0.1 < random.random() < 0.5:  # 平移
                print('平移')
                change_num += 1
                img, bboxes = self._shift_pic_bboxes(img, bboxes)

            if random.random() < 0.5:  # 改变亮度
                print('亮度')
                change_num += 1
                img = self._changeLight(img)

            if 0.45 < random.random() < 0.7:  # 加噪声
                print('加噪声')
                change_num += 1
                img = self._addNoise(img)

            if 0.5 < random.random() < 0.9:  # cutout
                print('cutout')
                change_num += 1
                img = self._cutout(img, bboxes, length=self.cut_out_length, n_holes=self.cut_out_holes,
                                   threshold=self.cut_out_threshold)

            if 0.3 < random.random() < 0.6:  # 翻转
                print('翻转')
                change_num += 1
                img, bboxes = self._filp_pic_bboxes(img, bboxes)

            print('\n')
        # print('------')
        return img, bboxes


if __name__ == '__main__':

    need_aug_num = 16 #一张图片将生成（15+360/旋转角度)张图片

    dataAug = DataAugmentForObjectDetection()

    path = 'data.txt'  #原始图片、标注信息路径
    files = open(path).readlines()
    i = 0
    for file in files:
        j = 0
        file = file.strip().split()
        pic_path = file[0]
        point = [int(file[i]) for i in range(1, len(file))]
        coords = [point[i:i + 6] for i in range(0, len(point) - 1, 6)]
        num = 0
        for coord in coords:
            num += 1
            if num > point[-1]:
                coord.append('neg')
            else:
                coord.append('pos')

        classname = [coord[-1] for coord in coords]
        coords = [coord[:7] for coord in coords]
        img = cv2.imread(pic_path)
        while j < need_aug_num:
            # show_pic(img, coords, classname)    # 原图
            if j == 0:
                auged_img, auged_bboxes = dataAug._crop_img_bboxes(img, coords)
                i += 1
            elif j == 1:
                auged_img, auged_bboxes = dataAug._shift_pic_bboxes(img, coords)
                i += 1
            elif j == 2:
                auged_img = dataAug._changeLight(img)
                auged_bboxes = coords
                i += 1
            elif j == 3:
                auged_img, auged_bboxes = dataAug._filp_pic_bboxes(img, coords)
                i += 1
            elif j == 4:
                auged_img = dataAug._addNoise(img)
                auged_bboxes = coords
                i += 1
            elif j == (need_aug_num - 1):
                for angle in range(0, 360, 6): #每旋转6度生成一张新图片
                    scale = random.uniform(0.7, 0.9)
                    auged_img, auged_bboxes = dataAug._rotate_img_bbox(img, coords, angle, scale)
                    a = random.randint(0, 7) #随机为每张旋转图片增加一种额外的增强方式
                    if a < 6:
                        auged_img, auged_bboxes = dataAug.methods(a, auged_img, auged_bboxes)
                    i += 1
                    m = 5 - len(str(i))
                    txt_name = str(0) * m + str(i) + '.txt'
                    txt_path = os.path.join('/PycharmProjects/lx/data/txt', txt_name) #标注信息保存路径
                    es_anchor = os.path.join('/PycharmProjects/lx/data/pre_anchor', 'anchor.txt')
                    f = open(txt_path, 'w')
                    for k in range(len(auged_bboxes)):
                        
                        f.writelines(str(auged_bboxes[k]) + '\n')
                        anchor = open(es_anchor, 'a')
                        a1 = auged_bboxes[k][2] - auged_bboxes[k][0]
                        a2 = auged_bboxes[k][3] - auged_bboxes[k][1]
                        anchor.writelines(str(a1) + ' ' + str(a2) + '\n')
                    f.close()
                    img_path = '/PycharmProjects/lx/data/img'  #增强后的图片保存路径
                    img_name = str(0) * m + str(i) + '.png'
                    img_name = os.path.join(img_path, img_name)
                    cv2.imwrite(img_name, auged_img)

                    pic_path = '/PycharmProjects/lx/data/pic' #带有标注信息的增强后的图片保存路径
                    aug_name = os.path.join(pic_path, 'pic' + str(i) + '.png')
                    # raw_name = os.path.join(pic_path, 'pic' + str(i) + '.png')
                    aug_img = show_pic(auged_img, auged_bboxes, classname)  # 强化后的图
                    cv2.imwrite(aug_name, aug_img)

                break
                # pic_path = '/home/lixiang/PycharmProjects/lx/data/pic'
            else:
                auged_img, auged_bboxes = dataAug.dataAugment(img, coords)
                i += 1
            n = 5 - len(str(i))
            txt_name = str(0) * n + str(i) + '.txt'
            txt_path = os.path.join('/PycharmProjects/lx/data/txt', txt_name) #增强后的图片保存路径
            es_anchor = os.path.join('/PycharmProjects/lx/data/pre_anchor', 'anchor.txt')
            f = open(txt_path, 'w')
            for k in range(len(auged_bboxes)):
                f.writelines(str(auged_bboxes[k]) + '\n')
                anchor = open(es_anchor, 'a')
                a1 = auged_bboxes[k][2] - auged_bboxes[k][0]
                a2 = auged_bboxes[k][3] - auged_bboxes[k][1]
                anchor.writelines(str(a1) + ' ' + str(a2) + '\n')
            f.close()
            img_path = '/PycharmProjects/lx/data/img'  #增强后的图片保存路径
            img_name = str(0) * n + str(i) + '.png'
            img_name = os.path.join(img_path, img_name)
            cv2.imwrite(img_name, auged_img)
            pic_path = '/PycharmProjects/lx/data/pic'  #带有标注信息的增强后的图片保存路径
            aug_name = os.path.join(pic_path, 'pic' + str(i) + '.png')
            raw_name = os.path.join(pic_path, 'pic' + str(i) + '.png')
            aug_img = show_pic(auged_img, auged_bboxes, classname)  # 强化后的图
            cv2.imwrite(aug_name, aug_img)
            # raw_image = show_pic(img, coords, classname)
            # cv2.imwrite(raw_name, raw_image)
            j += 1
            #print(i)
