from PIL import Image, ImageDraw
import numpy as np
import imgaug as ia
from imgaug import augmenters as iaa
import os


AUG_NUM = 200
PATH = '/home/jiangce/license_plate_unconstrain/images'
SAVE_PATH = '/home/jiangce/license_plate_unconstrain/data/images'

os.makedirs(SAVE_PATH, exist_ok=True)
os.makedirs(SAVE_PATH.replace('images', 'labels'), exist_ok = True)

def get_LP_loc(path):
    # 对标记数据进行处理，获取车牌四个点坐标的位置
    with open(path, 'r') as f:
        pos = f.readlines()[3].split(':')[1].split()
    x1 = int(pos[0])
    y1 = int(pos[1])
    x2 = int(pos[2])
    y2 = int(pos[3])
    x4 = int(pos[4])
    y4 = int(pos[5])
    x3 = int(pos[6])
    y3 = int(pos[7])
    return x1, y1, x2, y2, x3, y3, x4, y4


def data_augmentation(img_path):
    # 数据增强函数
    img_abs_path = os.path.join(PATH, img_path)
    label_path = img_abs_path.replace('images', 'labels').replace('.jpg', '.yaml')
    x1, y1, x2, y2, x3, y3, x4, y4 = get_LP_loc(label_path)
    im = Image.open(img_abs_path)
    image = np.array(im)
    keypoints = ia.KeypointsOnImage([ia.Keypoint(x = x1, y = y1),
                                     ia.Keypoint(x = x2, y = y2),
                                     ia.Keypoint(x = x3, y = y3),
                                     ia.Keypoint(x = x4, y = y4)], shape = image.shape)
    # 处理序列，来自imgaug的示例程序
    seq = iaa.Sequential([
    iaa.Fliplr(0.5), # horizontal flips
    iaa.Flipud(0.2), # vertically flip 20% of all images
    iaa.Crop(percent=(0, 0.1)), # random crops
    # Small gaussian blur with random sigma between 0 and 0.5.
    # But we only blur about 50% of all images.
    iaa.Sometimes(0.5,
        iaa.GaussianBlur(sigma=(0, 0.5))
    ),
    # Strengthen or weaken the contrast in each image.
    iaa.ContrastNormalization((0.75, 1.5)),
    # Add gaussian noise.
    # For 50% of all images, we sample the noise once per pixel.
    # For the other 50% of all images, we sample the noise per pixel AND
    # channel. This can change the color (not only brightness) of the
    # pixels.
    iaa.AdditiveGaussianNoise(loc=0, scale=(0.0, 0.05*255), per_channel=0.5),
    # Make some images brighter and some darker.
    # In 20% of all cases, we sample the multiplier once per channel,
    # which can end up changing the color of the images.
    iaa.Multiply((0.8, 1.2), per_channel=0.2),
    # Apply affine transformations to each image.
    # Scale/zoom them, translate/move them, rotate them and shear them.
    iaa.Affine(
        scale={"x": (0.8, 1.2), "y": (0.8, 1.2)},
        translate_percent={"x": (-0.2, 0.2), "y": (-0.2, 0.2)},
        rotate=(-25, 25),
        shear=(-8, 8)
    )
], random_order=True) # apply augmenters in random order
    seq_def = seq.to_deterministic()
    images_aug = seq_def.augment_images([image for _ in range(AUG_NUM)])
    keypoints_aug = seq_def.augment_keypoints([keypoints for _ in range(AUG_NUM)])
    print('\tFinish Argument')
    for i in range(AUG_NUM):
        save_flag = True
        image_aug = Image.fromarray(images_aug[i])
        # image_aug.save('/home/jiangce/index/images_new/%d.jpg' % i)
        # draw_pol(image_aug, keypoints_aug[i])
        hei, wid = image_aug.size
        # 判断增强后的照片，LP是否完整，若完整再储存
        for point in keypoints_aug[i].keypoints:
            if point.x < 0 or point.x > wid or point.y < 0 or point.y > hei:
                save_flag = False
                break
        if save_flag:
            img_save_path = os.path.join(SAVE_PATH , (img_path.split('.')[0]) + '_%d.jpg' % i)
            image_aug.save(img_save_path)
            label_save_path = img_save_path.replace('images', 'labels').replace('.jpg', '.yaml')
            # 存储新图片的标注信息
            with open(label_save_path, 'w') as file:
                file.write('image_file: ' + img_save_path + '\n')
                file.write('image_width: ' + str(wid) + '\n')
                file.write('image_height: ' + str(hei) + '\n')
                file.write('plate_corners_gt: ' +
                           str(int(keypoints_aug[i].keypoints[0].x)) + ' ' +
                           str(int(keypoints_aug[i].keypoints[0].y)) + ' ' +
                           str(int(keypoints_aug[i].keypoints[1].x)) + ' ' +
                           str(int(keypoints_aug[i].keypoints[1].y)) + ' ' +
                           str(int(keypoints_aug[i].keypoints[3].x)) + ' ' +
                           str(int(keypoints_aug[i].keypoints[3].y)) + ' ' +
                           str(int(keypoints_aug[i].keypoints[2].x)) + ' ' +
                           str(int(keypoints_aug[i].keypoints[2].y)) + '\n')
                file.write('plate_number_gt: ' + '\n')
                file.write('plate_inverted_gt: ')
        print("\tSave %d images" % i)


def draw_pol(image, keypoints_aug):
    # 在增强图片中绘制标注点，检测标注点的改变是否正确
    draw = ImageDraw.Draw(image)
    draw.polygon([int(keypoints_aug.keypoints[0].x),int(keypoints_aug.keypoints[0].y),int(keypoints_aug.keypoints[1].x),int(keypoints_aug.keypoints[1].y),
    int(keypoints_aug.keypoints[3].x),int(keypoints_aug.keypoints[3].y),int(keypoints_aug.keypoints[2].x),int(keypoints_aug.keypoints[2].y)], fill = (255,0,0))
    image.show()


if __name__ == '__main__':
    images_path = os.listdir(PATH)
    for image_path in images_path:
        data_augmentation(image_path)