目标检测利用imgaug包进行数据增强！！！

代码中的数据增强操作基本都罗列出来了，想使用哪些增强，就留哪些，不想用的注释掉就行。

'''
Author: CodingWZP
Email: codingwzp@gmail.com
Date: 2021-08-06 10:51:35
LastEditTime: 2021-08-09 10:53:43
Description: Image augmentation with label.
'''
import xml.etree.ElementTree as ET
import os
import imgaug as ia
import numpy as np
import shutil
from tqdm import tqdm
from PIL import Image
from imgaug import augmenters as iaa

ia.seed(1)


def read_xml_annotation(root, image_id):
    in_file = open(os.path.join(root, image_id))
    tree = ET.parse(in_file)
    root = tree.getroot()
    bndboxlist = []

    for object in root.findall('object'):  # 找到root节点下的所有country节点
        bndbox = object.find('bndbox')  # 子节点下节点rank的值

        xmin = int(bndbox.find('xmin').text)
        xmax = int(bndbox.find('xmax').text)
        ymin = int(bndbox.find('ymin').text)
        ymax = int(bndbox.find('ymax').text)
        # print(xmin,ymin,xmax,ymax)
        bndboxlist.append([xmin, ymin, xmax, ymax])
        # print(bndboxlist)

    bndbox = root.find('object').find('bndbox')
    return bndboxlist


def change_xml_list_annotation(root, image_id, new_target, saveroot, id):
    in_file = open(os.path.join(root, str(image_id) + '.xml'))  # 这里root分别由两个意思
    tree = ET.parse(in_file)
    # 修改增强后的xml文件中的filename
    elem = tree.find('filename')
    elem.text = (str(id) + '.jpg')
    xmlroot = tree.getroot()
    # 修改增强后的xml文件中的path
    elem = tree.find('path')
    if elem != None:
        elem.text = (saveroot + str(id) + '.jpg')

    index = 0
    for object in xmlroot.findall('object'):  # 找到root节点下的所有country节点
        bndbox = object.find('bndbox')  # 子节点下节点rank的值

        # xmin = int(bndbox.find('xmin').text)
        # xmax = int(bndbox.find('xmax').text)
        # ymin = int(bndbox.find('ymin').text)
        # ymax = int(bndbox.find('ymax').text)

        new_xmin = new_target[index][0]
        new_ymin = new_target[index][1]
        new_xmax = new_target[index][2]
        new_ymax = new_target[index][3]

        xmin = bndbox.find('xmin')
        xmin.text = str(new_xmin)
        ymin = bndbox.find('ymin')
        ymin.text = str(new_ymin)
        xmax = bndbox.find('xmax')
        xmax.text = str(new_xmax)
        ymax = bndbox.find('ymax')
        ymax.text = str(new_ymax)

        index = index + 1

    tree.write(os.path.join(saveroot, str(id + '.xml')))


def mkdir(path):
    # 去除首位空格
    path = path.strip()
    # 去除尾部 \ 符号
    path = path.rstrip("\\")
    # 判断路径是否存在
    # 存在     True
    # 不存在   False
    isExists = os.path.exists(path)
    # 判断结果
    if not isExists:
        # 如果不存在则创建目录
        # 创建目录操作函数
        os.makedirs(path)
        print(path + ' 创建成功')
        return True
    else:
        # 如果目录存在则不创建，并提示目录已存在
        print(path + ' 目录已存在')
        return False


if __name__ == "__main__":

    IMG_DIR = "F:/huawei/data/hangkong/nick/nick/image"
    XML_DIR = "F:/huawei/data/hangkong/nick/nick/annotations"

    AUG_XML_DIR = "F:/huawei/data/hangkong/augnickxml"  # 存储增强后的XML文件夹路径
    try:
        shutil.rmtree(AUG_XML_DIR)
    except FileNotFoundError as e:
        a = 1
    mkdir(AUG_XML_DIR)

    AUG_IMG_DIR = "F:/huawei/data/hangkong/augnick"  # 存储增强后的影像文件夹路径
    try:
        shutil.rmtree(AUG_IMG_DIR)
    except FileNotFoundError as e:
        a = 1
    mkdir(AUG_IMG_DIR)

    AUGLOOP = 44 # 每张影像增强的数量

    boxes_img_aug_list = []
    new_bndbox = []
    new_bndbox_list = []

    # 影像增强
    seq = iaa.Sequential([
        iaa.Invert(0.5),#反转图像像素
        iaa.Fliplr(0.5),  # 水平翻船
        #iaa.Grayscale(alpha=(0.0, 1.0)),#将图像更改为灰度，并通过改变强度将其与原始图像叠加，有效地删除0到100％的颜色：
        iaa.Flipud(0.5),#垂直翻转
        iaa.AddElementwise((-40, 40),per_channel=0.5),#向具有相邻像素的图像像素添加不同值。在图像中添加-40到40之间的随机值，每个像素采样一次：
        #iaa.Affine(0.5,rotate=(-45, 45)),#将图像旋转-45到45度：
        iaa.Dropout(p=(0, 0.2), per_channel=0.5),#每个图像的样本值为0<=p<=0.2，然后将图像中p％像素转换为黑色像素，但在50％的图像中每个通道独立执行此操作：
        iaa.Multiply((1.2, 1.5)),  # 将图像中相邻像素乘以不同的值，使每个像素更暗或更亮。
        iaa.MedianBlur(k=(1, 9), name=None, deterministic="deprecated", random_state=None),#中值模糊
        iaa.AverageBlur(k=(1, 9), name=None, deterministic="deprecated", random_state=None),#均值模糊
        iaa.GaussianBlur(sigma=(0, 3.0)),  # iaa.GaussianBlur(0.5),
        iaa.contrast.LinearContrast(alpha=(0.5, 1.5), per_channel=False, name=None, deterministic="deprecated",
                                    random_state=None),#线性对比度
        iaa.contrast.LogContrast(gain=(0.5, 1.5), per_channel=False, name=None, deterministic="deprecated",
                                 random_state=None),#对数对比度
        iaa.contrast.GammaContrast(gamma=(0.5, 1.5), per_channel=False, name=None, deterministic="deprecated",
                                   random_state=None),#伽马对比度
        iaa.AdditiveGaussianNoise(loc=0, scale=0, per_channel=False, name=None, deterministic="deprecated",
                                  random_state=None),#高斯噪声
        iaa.imgcorruptlike.ShotNoise(severity=2),#散粒噪声
        iaa.imgcorruptlike.ImpulseNoise(severity=2),#脉冲噪声
        iaa.imgcorruptlike.SpeckleNoise(severity=2),#斑点噪声
        iaa.imgcorruptlike.Brightness(severity=2),#调整图像亮度
        iaa.imgcorruptlike.Saturate(severity=3),#调整图像饱和度
        #iaa.contrast.AllChannelsHistogramEqualization(seed=None, name=None, deterministic="deprecated",
                                                      #random_state=None),#对图像所有通道做直方图均衡化。
        iaa.Sharpen(alpha=(0, 1.0), lightness=(0.75, 1.5)),#图像锐化：补偿图像的轮廓，增强图像的边缘及灰度跳变的部分，使图像变得清晰
    #     iaa.Affine(
    #         translate_px={"x": 15, "y": 15},#按照像素来进行平移
    #         scale=(0.8, 0.95),# # 图像缩放为80%到95%之间
    #     )  # translate by 40/60px on x/y axis, and scale to 50-70%, affects BBs
    #
         ])

    for name in tqdm(os.listdir(XML_DIR), desc='Processing'):

        bndbox = read_xml_annotation(XML_DIR, name)

        # 保存原xml文件
        shutil.copy(os.path.join(XML_DIR, name), AUG_XML_DIR)
        # 保存原图
        og_img = Image.open(IMG_DIR + '/' + name[:-4] + '.jpg')
        og_img.convert('RGB').save(AUG_IMG_DIR + name[:-4] + '.jpg', 'JPEG')
        og_xml = open(os.path.join(XML_DIR, name))
        tree = ET.parse(og_xml)
        # 修改增强后的xml文件中的filename
        elem = tree.find('filename')
        elem.text = (name[:-4] + '.jpg')
        tree.write(os.path.join(AUG_XML_DIR, name))

        for epoch in range(AUGLOOP):
            seq_det = seq.to_deterministic()  # 保持坐标和图像同步改变，而不是随机
            # 读取图片
            img = Image.open(os.path.join(IMG_DIR, name[:-4] + '.jpg'))
            # sp = img.size
            img = np.asarray(img)
            # bndbox 坐标增强
            for i in range(len(bndbox)):
                bbs = ia.BoundingBoxesOnImage([
                    ia.BoundingBox(x1=bndbox[i][0], y1=bndbox[i][1], x2=bndbox[i][2], y2=bndbox[i][3]),
                ], shape=img.shape)

                bbs_aug = seq_det.augment_bounding_boxes([bbs])[0]
                boxes_img_aug_list.append(bbs_aug)

                # new_bndbox_list:[[x1,y1,x2,y2],...[],[]]
                n_x1 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x1)))
                n_y1 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y1)))
                n_x2 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x2)))
                n_y2 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y2)))
                if n_x1 == 1 and n_x1 == n_x2:
                    n_x2 += 1
                if n_y1 == 1 and n_y2 == n_y1:
                    n_y2 += 1
                if n_x1 >= n_x2 or n_y1 >= n_y2:
                    print('error', name)
                new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2])
            # 存储变化后的图片
            image_aug = seq_det.augment_images([img])[0]
            path = os.path.join(AUG_IMG_DIR,
                                str(str(name[:-4]) + '_' + str(epoch)) + '.jpg')
            image_auged = bbs.draw_on_image(image_aug, size=0)
            Image.fromarray(image_auged).convert('RGB').save(path)

            # 存储变化后的XML
            change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR,
                                       str(name[:-4]) + '_' + str(epoch))
            # print(str(str(name[:-4]) + '_' + str(epoch)) + '.jpg')
            new_bndbox_list = []
    print('Finish!')