import xml.etree.ElementTree as ET
import os
from os import getcwd
import numpy as np
from PIL import Image
import cv2 as cv
import shutil
import matplotlib.pyplot as plt
import imgaug as ia
from imgaug import augmenters as iaa
from pathlib import Path

ia.seed(1)


def find_x_y(root):
    for object in root.findall("object"):
        bndbox = object.find("bndbox")
        #print(f"type:{type(bndbox)}")
        xmin = bndbox.find("xmin")
        ymin = bndbox.find("ymin")
        xmax = bndbox.find("xmax")
        ymax = bndbox.find("ymax")
    return (xmin, ymin, xmax, ymax)


def find_target(new_target, index):
    new_xmin = new_target[index][0]
    new_ymin = new_target[index][1]
    new_xmax = new_target[index][2]
    new_ymax = new_target[index][3]
    return new_xmin, new_ymin, new_xmax, new_ymax


def write_x_y(*coordinate):
    coordinate[0].text = str(coordinate[4])
    coordinate[1].text = str(coordinate[5])
    coordinate[2].text = str(coordinate[6])
    coordinate[3].text = str(coordinate[7])


def read_xml_annotation(root, image_id):
    with open(str(Path(root + os.sep) / image_id if isinstance(image_id, str) else str(image_id)),
              encoding="UTF-8") as in_file:
        tree = ET.parse(in_file)
        root = tree.getroot()
        xmin, ymin, xmax, ymax = find_x_y(root)
        bndboxlist = []
        bndboxlist.append([int(xmin.text), int(ymin.text),
                           int(xmax.text), int(ymax.text)])

    return bndboxlist


def change_xml_list_annotation(root, image_id, new_target, save_root, id):
    with open(str(Path(root + os.sep) / image_id if isinstance(image_id, str) else str(image_id)) + ".xml",
              encoding="UTF-8") as in_file:
        tree = ET.parse(in_file)
        elem = tree.find("filename")
        elem.text = (str("%06d" % int(id)) + ".jpg")
        xmlroot = tree.getroot()
        index = 0
        for i in xmlroot.findall("object"):
            xmin, ymin, xmax, ymax = find_x_y(xmlroot)
            new_xmin, new_ymin, new_xmax, new_ymax = find_target(new_target, index)
            coordinate = [xmin, ymin, xmax, ymax, new_xmin, new_ymin, new_xmax, new_ymax]
            write_x_y(*coordinate)
            index += 1
            #print(f"index={index}")
        tree.write(str(Path(save_root + os.sep) / str("%06d" % int(id))) + ".xml")


def mkdir(dir_path):
    dir_path = dir_path.strip()
    dir_path = dir_path.rstrip("\\")
    os.makedirs(dir_path, exist_ok=True)
def imgaug():
    seq = iaa.Sequential([
        iaa.Fliplr(0.5),
        iaa.Flipud(0.5),
        iaa.ContrastNormalization((.75, 1.5), per_channel=True),
        iaa.AdditiveGaussianNoise(loc=0, scale=(0.0, 0.1 * 255), per_channel=0.5),
        iaa.Multiply((0.8, 1.2), per_channel=0.2)
    ])
    return seq

if __name__ == "__main__":
    IMG_DIR = r"D:/imgaug_test/jpegimages"
    XML_DIR = r"D:/imgaug_test/annotation"
    AUG_XML_DIR = r"D:/imgaug_test/annotation_augmentation"
    AUG_IMG_DIR = r"D:/imgaug_test/jpegimages_augmentation"
    try:
        # 递归删除目录以及子目录和目录下的所有文件
        shutil.rmtree(AUG_XML_DIR)
    except FileNotFoundError as e:
        print(f"aug_xml_dir_error:{e}")
    mkdir(AUG_XML_DIR)

    try:
        shutil.rmtree(AUG_IMG_DIR)
    except FileNotFoundError as e:
        print(f"aug_img_dir_error:{e}")
    mkdir(AUG_IMG_DIR)

    AUG_LOOP = 20

    boxes_img_aug_list = []
    new_bndbox = []
    new_bndbox_list = []
	seq = imgaug()

    for root, sub_folders, files in os.walk(XML_DIR):
        for name in files:
            bndbox = read_xml_annotation(XML_DIR, name)
            # 原图和xml复制到增强目录下
            shutil.copy(str(Path(XML_DIR + os.sep) / name), AUG_XML_DIR)
            shutil.copy(str(Path(IMG_DIR + os.sep) / name[:-4]) + ".jpg", AUG_IMG_DIR)
            for epoch in range(AUG_LOOP):
                # 固定变换
                seq_det = seq.to_deterministic()
                img = cv.imread(str(Path(IMG_DIR + os.sep) / name[:-4]) + ".jpg")
                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)
                    # 变换后的bounding box
                    bbs_aug = seq_det.augment_bounding_boxes([bbs])[0]
                    boxes_img_aug_list.append(bbs_aug)
                    # 控制框不出原图并且是正整数
                    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(f"error", name)
                    new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2])
                    # 变换后的图像
                    image_aug = seq_det.augment_images([img])[0]
                    path = str(Path(AUG_IMG_DIR + os.sep) / \
                               str("%06d" % (len(files) + int(name[:-4]) + epoch * 250))) + ".jpg"
                    image_aug = bbs.draw_on_image(image_aug, thickness=0)
                    cv.imwrite(path, image_aug)
                change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR,
                                           len(files) + int(name[:-4]) + epoch * 250)
                print(str("%06d" % (len(files) + int(name[:-4]) + epoch * 250)) + '.jpg')
                new_bndbox_list = []

参考imgaug数据增强