# coding=UTF-8
# This Python file uses the following encoding: utf-8

import cv2
import numpy as np
import xml.etree.cElementTree as ET
from random import sample

#default args:
default_args = {'noise_prob': 0.1,
                'gasuss_mean': 0,
                'gasuss_var': 0.001,
                'rand_hug': 30,
                'rand_saturation':30,
                'rand_light': 30,
                'rot_angle': 15,
                'bordervalue': (127, 127, 127),
                'zoom_out_value': 0.7,
                'output_shape': (416, 416),
                'take_value' : 5
               }


#添加黑色noise
def sp_noise(image, box_loc=None, **kwargs):
    h, w = image.shape[0:2]
    noise = np.random.rand(h,w)
    out_img = image.copy()
    out_img[noise < kwargs['noise_prob']] = 0
    if box_loc is None:
        return out_img
    else:
        return out_img, box_loc

#高斯noise
def gasuss_noise(image, box_loc=None, **kwargs):
    out_img = (image / 255.) - 0.5
    noise = np.random.normal(kwargs['gasuss_mean'], kwargs['gasuss_var']** 0.5, image.shape)
    out_img = out_img + noise + 0.5
    out_img[out_img < 0] = 0
    out_img[out_img > 1] = 1
    out_img = (out_img * 255).astype(np.uint8)
    if box_loc is None:
        return out_img
    else:
        return out_img, box_loc

#調整彩度(彩度通道加上隨機-N~N之值)
def mod_hue(image, box_loc=None, **kwargs):
    out_img = cv2.cvtColor(image, cv2.COLOR_BGR2HSV).astype(np.float32)
    out_img[:,:,0] += np.random.randint(-kwargs['rand_hug'], kwargs['rand_hug'])
    out_img = cv2.cvtColor(np.clip(out_img, 0, 180).astype(np.uint8), cv2.COLOR_HSV2BGR)
    if box_loc is None:
        return out_img
    else:
        return out_img, box_loc

#調整飽和度(飽和度通道加上隨機-N~N之值)
def mod_saturation(image, box_loc=None, **kwargs):
    out_img = cv2.cvtColor(image, cv2.COLOR_BGR2HSV).astype(np.float32)
    out_img[:,:,1] += np.random.randint(-kwargs['rand_saturation'], kwargs['rand_saturation'])
    out_img = cv2.cvtColor(np.clip(out_img, 0, 255).astype(np.uint8), cv2.COLOR_HSV2BGR)
    if box_loc is None:
        return out_img
    else:
        return out_img, box_loc

#調整亮度(亮度通道加上隨機-N~N之值)
def mod_light(image, box_loc=None, **kwargs):
    out_img = cv2.cvtColor(image, cv2.COLOR_BGR2HSV).astype(np.float32)
    out_img[:,:,2] += np.random.randint(-kwargs['rand_light'], kwargs['rand_light'])
    out_img = cv2.cvtColor(np.clip(out_img, 0, 255).astype(np.uint8), cv2.COLOR_HSV2BGR)
    if box_loc is None:
        return out_img
    else:
        return out_img, box_loc

#水平翻轉
def horizontal_flip(image, box_loc=None, **kwargs):
    '''
    Args:
        box_loc: bounding box location(x_min, y_min, x_max, y_max)
    '''
    if box_loc is None:
        return cv2.flip(image, 1)
    else:
        w = image.shape[1]
        for i in box_loc:
            if i[2] == 0:
                break
            else:
                x_min, x_max = i[0], i[2]
                i[0] = w - x_max
                i[2] = w - x_min
        return cv2.flip(image, 1), box_loc

#垂直翻轉
def vertical_flip(image, box_loc=None, **kwargs):
    '''
    Args:
        box_loc: bounding box location(num box,(x_min, y_min, x_max, y_max, label))
    '''
    if box_loc is None:
        return cv2.flip(image, 0)
    else:
        h = image.shape[0]
        for i in box_loc:
            if i[3] == 0:
                break
            else:
                y_min, y_max = i[1], i[3]
                i[1] = h - y_max
                i[3] = h - y_min
        return cv2.flip(image, 0), box_loc

#旋轉-n~n度
def rot_image(image, box_loc=None, **kwargs):
    '''
    Args:
        box_loc: bounding box location(num box,(x_min, y_min, x_max, y_max, label))
        rot: 要選轉的範圍
        bordervalue: 空白處補的值
    '''
    h, w, _ = image.shape
    center = ( w // 2, h // 2)
    angle = np.random.randint(-kwargs['rot_angle'], kwargs['rot_angle'])
    M = cv2.getRotationMatrix2D(center, angle, 1)
    out_img = cv2.warpAffine(image, M, (w, h), borderValue = kwargs['bordervalue'])
    if box_loc is None:
        return out_img
    else:
        loc = box_loc[:,0:4].copy()
        loc = np.append(loc, loc[:, 0:1], axis=-1)
        loc = np.append(loc, loc[:, 3:4], axis=-1)
        loc = np.append(loc, loc[:, 2:3], axis=-1)
        loc = np.append(loc, loc[:, 1:2], axis=-1)
        loc = loc.reshape(-1, 4, 2)
        loc = loc - np.array(center)
        rot_loc = loc.dot(np.transpose(M[:,0:2]))
        rot_loc = rot_loc + np.array(center)
        rot_box = np.hstack([np.min(rot_loc, axis=-2), np.max(rot_loc, axis=-2), box_loc[:, 4:5]])
        rot_box = np.floor(rot_box)
        rot_box[...,0:4] = np.clip(rot_box[...,0:4], [0,0,0,0], [w-1, h-1, w-1, h-1])

        return out_img, rot_box

#等比例縮放影像
def resize_img(image, box_loc=None, **kwargs):
    h, w, _ = image.shape
    max_edge = max(kwargs['output_shape'][0], kwargs['output_shape'][1])
    scale = min( max_edge / h, max_edge / w)
    h = int(h * scale)
    w = int(w * scale)

    if box_loc is None:
        return cv2.resize(image, (w, h))
    else:
        box_loc[:,0] = box_loc[:,0] * scale
        box_loc[:,1] = box_loc[:,1] * scale
        box_loc[:,2] = box_loc[:,2] * scale
        box_loc[:,3] = box_loc[:,3] * scale
        return cv2.resize(image, (w, h)), box_loc.astype(np.int32)

#將樸片補至指定大小
def padding_img(image, box_loc=None, **kwargs):
    h, w, _ = image.shape

    dx = int((kwargs['output_shape'][1] - w) / 2)
    dy = int((kwargs['output_shape'][0] - h) / 2)

    out_img = np.ones((kwargs['output_shape'][0], kwargs['output_shape'][1], 3), np.uint8) * kwargs['bordervalue'][0]
    out_img[dy: dy + h, dx: dx + w] = cv2.resize(image, (w, h))

    if box_loc is None:
        return out_img
    else:
        box_loc[:,0] = box_loc[:,0] + dx
        box_loc[:,1] = box_loc[:,1] + dy
        box_loc[:,2] = box_loc[:,2] + dx
        box_loc[:,3] = box_loc[:,3] + dy
        return out_img, box_loc.astype(np.int32)

#隨機縮小 value~1倍
def random_zoom_out(image, box_loc=None, **kwargs):

    h, w, _ = image.shape
    scale = np.random.uniform(kwargs['zoom_out_value'], 1)
    h = int(h * scale)
    w = int(w * scale)
    dx = int((image.shape[1] - w) / 2)
    dy = int((image.shape[0] - h) / 2)
    out_img = np.ones(image.shape, np.uint8) * kwargs['bordervalue'][0]
    out_img[dy: dy + h, dx: dx + w] = cv2.resize(image, (w, h))

    if box_loc is None:
        return out_img
    else:
        box_loc[:,0] = box_loc[:,0] * scale + dx
        box_loc[:,1] = box_loc[:,1] * scale + dy
        box_loc[:,2] = box_loc[:,2] * scale + dx
        box_loc[:,3] = box_loc[:,3] * scale + dy
        return out_img, box_loc.astype(np.int32)


#load csv data
def load_csv(xml_path, max_boxes=4):
    tree = ET.parse(xml_path)
    root = tree.getroot()

    #location list
    loc_list = np.zeros((0, 5))
    box_count = 0

    for obj in root.iter('object'):
        if box_count >= max_boxes:
            break

        '''
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue

        cls_id = classes.index(cls)
        '''
        loc = obj.find('bndbox')
        x_min = int(loc.find('xmin').text)
        y_min = int(loc.find('ymin').text)
        x_max = int(loc.find('xmax').text)
        y_max = int(loc.find('ymax').text)

        loc_list = np.vstack([loc_list, np.array([x_min, y_min, x_max, y_max, 0])])

        box_count += 1
    
    return loc_list.astype(np.float32)

#draw rectangle
def draw_rect(image, box_loc):
    for i in box_loc:
        cv2.rectangle(image, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), (0, 255, 0), 4)

def print_args(**kwargs):
    for key, value in kwargs.items():
        print('key name: {}\nvalue:{}\n'.format(key, value))

#隨機選擇0~N個 image augmentation方法
def rand_aug_image(image, box_loc=None, **kwargs):

    if box_loc is None:
        out_img = resize_img(image, **kwargs)
    else:
        out_img, box_loc = resize_img(image, box_loc, **kwargs)

    #total augmentation function
    func_list = [sp_noise, gasuss_noise, mod_hue, mod_saturation, mod_light,
                 horizontal_flip, vertical_flip, rot_image, random_zoom_out]
    
    #rand take function
    take_func = sample(func_list, np.random.randint(kwargs['take_value']))

    for func in take_func:
        if box_loc is None:
            out_img = func(out_img, **kwargs)
        else:
            out_img, box_loc = func(out_img, box_loc, **kwargs)
    if box_loc is None:
        out_img = padding_img(out_img, **kwargs)
        return out_img
    else:
        out_img, box_loc = padding_img(out_img, box_loc, **kwargs)
        return out_img, box_loc


if __name__ == "__main__":
    img = cv2.imread('./00002.jpg')
    bbox = load_csv('./00002.xml')

    #黑點noise
    #aug_img = sp_noise(img, **default_args)
    #aug_img, bbox = sp_noise(img, bbox, **default_args)

    #gasuss_noise
    #aug_img = gasuss_noise(img, **default_args)
    #aug_img, bbox = gasuss_noise(img, bbox, **default_args)

    #調整Hue
    #aug_img = mod_hue(img, **default_args)
    #aug_img, bbox = mod_hue(img, bbox, **default_args)

    #調整saturation
    #aug_img = mod_saturation(img, **default_args)
    #aug_img, bbox = mod_saturation(img, bbox, **default_args)

    #調整light
    #aug_img = mod_light(img, **default_args)
    #aug_img, bbox = mod_light(img, bbox, **default_args)

    #水平翻轉
    #aug_img = horizontal_flip(img, **default_args)
    #aug_img, bbox = horizontal_flip(img, bbox, **default_args)

    #垂直翻轉
    #aug_img = vertical_flip(img, **default_args)
    #aug_img, bbox = vertical_flip(img, bbox, **default_args)

    #旋轉角度
    #aug_img = rot_image(img, **default_args)
    #aug_img, bbox = rot_image(img, bbox, **default_args)

    #等比例resize至指定大小
    #aug_img = resize_img(img, **default_args)
    #aug_img, bbox = resize_img(img, bbox, **default_args)

    #補形狀至指定大小
    #aug_img = padding_img(aug_img, **default_args)
    #aug_img, bbox = padding_img(aug_img, bbox, **default_args)

    #隨機縮小 N~1倍
    #aug_img = random_zoom_out(img, **default_args)
    #aug_img, bbox = random_zoom_out(img, bbox, **default_args)

    #隨機選擇augmentation方法
    aug_img = rand_aug_image(img, **default_args)
    #aug_img, bbox = rand_aug_image(img, bbox, **default_args)
    print(bbox)

    draw_rect(aug_img, bbox)
    cv2.imshow('img', img)
    cv2.imshow('aug img', aug_img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
[((27, 12, 27, 31), 'numpy.random.rand', 'np.random.rand', ({(27, 27, 27, 28): 'h', (27, 29, 27, 30): 'w'}, {}), '(h, w)', True, 'import numpy as np\n'), ((38, 12, 38, 92), 'numpy.random.normal', 'np.random.normal', ({(38, 29, 38, 50): "kwargs['gasuss_mean']", (38, 52, 38, 78): "kwargs['gasuss_var'] ** 0.5", (38, 80, 38, 91): 'image.shape'}, {}), "(kwargs['gasuss_mean'], kwargs['gasuss_var'] ** 0.5, image.\n    shape)", True, 'import numpy as np\n'), ((51, 22, 51, 80), 'numpy.random.randint', 'np.random.randint', ({(51, 40, 51, 59): "(-kwargs['rand_hug'])", (51, 61, 51, 79): "kwargs['rand_hug']"}, {}), "(-kwargs['rand_hug'], kwargs['rand_hug'])", True, 'import numpy as np\n'), ((61, 22, 61, 94), 'numpy.random.randint', 'np.random.randint', ({(61, 40, 61, 66): "(-kwargs['rand_saturation'])", (61, 68, 61, 93): "kwargs['rand_saturation']"}, {}), "(-kwargs['rand_saturation'], kwargs['rand_saturation'])", True, 'import numpy as np\n'), ((71, 22, 71, 84), 'numpy.random.randint', 'np.random.randint', ({(71, 40, 71, 61): "(-kwargs['rand_light'])", (71, 63, 71, 83): "kwargs['rand_light']"}, {}), "(-kwargs['rand_light'], kwargs['rand_light'])", True, 'import numpy as np\n'), ((126, 12, 126, 72), 'numpy.random.randint', 'np.random.randint', ({(126, 30, 126, 50): "-kwargs['rot_angle']", (126, 52, 126, 71): "kwargs['rot_angle']"}, {}), "(-kwargs['rot_angle'], kwargs['rot_angle'])", True, 'import numpy as np\n'), ((127, 8, 127, 49), 'cv2.getRotationMatrix2D', 'cv2.getRotationMatrix2D', ({(127, 32, 127, 38): 'center', (127, 40, 127, 45): 'angle', (127, 47, 127, 48): '1'}, {}), '(center, angle, 1)', False, 'import cv2\n'), ((128, 14, 128, 83), 'cv2.warpAffine', 'cv2.warpAffine', (), '', False, 'import cv2\n'), ((172, 38, 172, 63), 'cv2.resize', 'cv2.resize', ({(172, 49, 172, 54): 'image', (172, 56, 172, 62): '(w, h)'}, {}), '(image, (w, h))', False, 'import cv2\n'), ((187, 12, 187, 58), 'numpy.random.uniform', 'np.random.uniform', ({(187, 30, 187, 54): "kwargs['zoom_out_value']", (187, 56, 187, 57): '1'}, {}), "(kwargs['zoom_out_value'], 1)", True, 'import numpy as np\n'), ((193, 38, 193, 63), 'cv2.resize', 'cv2.resize', ({(193, 49, 193, 54): 'image', (193, 56, 193, 62): '(w, h)'}, {}), '(image, (w, h))', False, 'import cv2\n'), ((207, 11, 207, 29), 'xml.etree.cElementTree.parse', 'ET.parse', ({(207, 20, 207, 28): 'xml_path'}, {}), '(xml_path)', True, 'import xml.etree.cElementTree as ET\n'), ((211, 15, 211, 31), 'numpy.zeros', 'np.zeros', ({(211, 24, 211, 30): '(0, 5)'}, {}), '((0, 5))', True, 'import numpy as np\n'), ((276, 10, 276, 35), 'cv2.imread', 'cv2.imread', ({(276, 21, 276, 34): '"""./00002.jpg"""'}, {}), "('./00002.jpg')", False, 'import cv2\n'), ((329, 4, 329, 26), 'cv2.imshow', 'cv2.imshow', ({(329, 15, 329, 20): '"""img"""', (329, 22, 329, 25): 'img'}, {}), "('img', img)", False, 'import cv2\n'), ((330, 4, 330, 34), 'cv2.imshow', 'cv2.imshow', ({(330, 15, 330, 24): '"""aug img"""', (330, 26, 330, 33): 'aug_img'}, {}), "('aug img', aug_img)", False, 'import cv2\n'), ((331, 4, 331, 18), 'cv2.waitKey', 'cv2.waitKey', ({(331, 16, 331, 17): '(0)'}, {}), '(0)', False, 'import cv2\n'), ((332, 4, 332, 27), 'cv2.destroyAllWindows', 'cv2.destroyAllWindows', ({}, {}), '()', False, 'import cv2\n'), ((85, 15, 85, 33), 'cv2.flip', 'cv2.flip', ({(85, 24, 85, 29): 'image', (85, 31, 85, 32): '(1)'}, {}), '(image, 1)', False, 'import cv2\n'), ((104, 15, 104, 33), 'cv2.flip', 'cv2.flip', ({(104, 24, 104, 29): 'image', (104, 31, 104, 32): '(0)'}, {}), '(image, 0)', False, 'import cv2\n'), ((133, 14, 133, 50), 'numpy.append', 'np.append', (), '', True, 'import numpy as np\n'), ((134, 14, 134, 50), 'numpy.append', 'np.append', (), '', True, 'import numpy as np\n'), ((135, 14, 135, 50), 'numpy.append', 'np.append', (), '', True, 'import numpy as np\n'), ((136, 14, 136, 50), 'numpy.append', 'np.append', (), '', True, 'import numpy as np\n'), ((142, 18, 142, 35), 'numpy.floor', 'np.floor', ({(142, 27, 142, 34): 'rot_box'}, {}), '(rot_box)', True, 'import numpy as np\n'), ((143, 27, 143, 85), 'numpy.clip', 'np.clip', ({(143, 35, 143, 51): 'rot_box[(...), 0:4]', (143, 53, 143, 62): '[0, 0, 0, 0]', (143, 64, 143, 84): '[w - 1, h - 1, w - 1, h - 1]'}, {}), '(rot_box[(...), 0:4], [0, 0, 0, 0], [w - 1, h - 1, w - 1, h - 1])', True, 'import numpy as np\n'), ((156, 15, 156, 40), 'cv2.resize', 'cv2.resize', ({(156, 26, 156, 31): 'image', (156, 33, 156, 39): '(w, h)'}, {}), '(image, (w, h))', False, 'import cv2\n'), ((171, 14, 171, 90), 'numpy.ones', 'np.ones', ({(171, 22, 171, 79): "(kwargs['output_shape'][0], kwargs['output_shape'][1], 3)", (171, 81, 171, 89): 'np.uint8'}, {}), "((kwargs['output_shape'][0], kwargs['output_shape'][1], 3), np.uint8)", True, 'import numpy as np\n'), ((192, 14, 192, 44), 'numpy.ones', 'np.ones', ({(192, 22, 192, 33): 'image.shape', (192, 35, 192, 43): 'np.uint8'}, {}), '(image.shape, np.uint8)', True, 'import numpy as np\n'), ((260, 34, 260, 73), 'numpy.random.randint', 'np.random.randint', ({(260, 52, 260, 72): "kwargs['take_value']"}, {}), "(kwargs['take_value'])", True, 'import numpy as np\n'), ((50, 14, 50, 52), 'cv2.cvtColor', 'cv2.cvtColor', ({(50, 27, 50, 32): 'image', (50, 34, 50, 51): 'cv2.COLOR_BGR2HSV'}, {}), '(image, cv2.COLOR_BGR2HSV)', False, 'import cv2\n'), ((60, 14, 60, 52), 'cv2.cvtColor', 'cv2.cvtColor', ({(60, 27, 60, 32): 'image', (60, 34, 60, 51): 'cv2.COLOR_BGR2HSV'}, {}), '(image, cv2.COLOR_BGR2HSV)', False, 'import cv2\n'), ((70, 14, 70, 52), 'cv2.cvtColor', 'cv2.cvtColor', ({(70, 27, 70, 32): 'image', (70, 34, 70, 51): 'cv2.COLOR_BGR2HSV'}, {}), '(image, cv2.COLOR_BGR2HSV)', False, 'import cv2\n'), ((95, 15, 95, 33), 'cv2.flip', 'cv2.flip', ({(95, 24, 95, 29): 'image', (95, 31, 95, 32): '(1)'}, {}), '(image, 1)', False, 'import cv2\n'), ((114, 15, 114, 33), 'cv2.flip', 'cv2.flip', ({(114, 24, 114, 29): 'image', (114, 31, 114, 32): '(0)'}, {}), '(image, 0)', False, 'import cv2\n'), ((138, 20, 138, 36), 'numpy.array', 'np.array', ({(138, 29, 138, 35): 'center'}, {}), '(center)', True, 'import numpy as np\n'), ((139, 26, 139, 48), 'numpy.transpose', 'np.transpose', ({(139, 39, 139, 47): 'M[:, 0:2]'}, {}), '(M[:, 0:2])', True, 'import numpy as np\n'), ((140, 28, 140, 44), 'numpy.array', 'np.array', ({(140, 37, 140, 43): 'center'}, {}), '(center)', True, 'import numpy as np\n'), ((162, 15, 162, 40), 'cv2.resize', 'cv2.resize', ({(162, 26, 162, 31): 'image', (162, 33, 162, 39): '(w, h)'}, {}), '(image, (w, h))', False, 'import cv2\n'), ((52, 27, 52, 51), 'numpy.clip', 'np.clip', ({(52, 35, 52, 42): 'out_img', (52, 44, 52, 45): '0', (52, 47, 52, 50): '180'}, {}), '(out_img, 0, 180)', True, 'import numpy as np\n'), ((62, 27, 62, 51), 'numpy.clip', 'np.clip', ({(62, 35, 62, 42): 'out_img', (62, 44, 62, 45): '0', (62, 47, 62, 50): '255'}, {}), '(out_img, 0, 255)', True, 'import numpy as np\n'), ((72, 27, 72, 51), 'numpy.clip', 'np.clip', ({(72, 35, 72, 42): 'out_img', (72, 44, 72, 45): '0', (72, 47, 72, 50): '255'}, {}), '(out_img, 0, 255)', True, 'import numpy as np\n'), ((141, 29, 141, 53), 'numpy.min', 'np.min', (), '', True, 'import numpy as np\n'), ((141, 55, 141, 79), 'numpy.max', 'np.max', (), '', True, 'import numpy as np\n'), ((232, 40, 232, 81), 'numpy.array', 'np.array', ({(232, 49, 232, 80): '[x_min, y_min, x_max, y_max, 0]'}, {}), '([x_min, y_min, x_max, y_max, 0])', True, 'import numpy as np\n')]