python实现xml标注数据正矩形框的目标检测数据增强并保存新的xml

# -*- coding:utf-8 -*-

import xml.etree.ElementTree as ET
import pickle
import os
from os import getcwd
import numpy as np
from PIL import Image
import shutil
import matplotlib.pyplot as plt

import imgaug as ia
from imgaug import augmenters as iaa
import  itertools

def read_xml_annotation(root, image_id):
    in_file = open(os.path.join(root, image_id),encoding='utf-8', errors='ignore')
    tree = ET.parse(in_file)
    root = tree.getroot()
    in_file.close()
    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'),encoding='utf-8', errors='ignore')  # 这里root分别由两个意思
    tree = ET.parse(in_file)
    elem = tree.find('filename')
    elem.text = id
    xmlroot = tree.getroot()
    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, id[:-4] + '.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 + ' Success')
        return True
    else:
        # 如果目录存在则不创建，并提示目录已存在
        print(path + ' Dir exist')
        return False

if __name__ == "__main__":

    IMG_DIR = "C:\\Users\\Desktop\\temp\\test"
    XML_DIR = "C:\\Users\\Desktop\\temp\\test1"

    AUG_XML_DIR = "./Augxml/"  # 存储增强后的xml文件夹路径
    try:
        shutil.rmtree(AUG_XML_DIR)
    except FileNotFoundError as e:
        a = 1
    mkdir(AUG_XML_DIR)

    AUG_IMG_DIR = "./AugImg/"  # 存储增强后?bmp文件夹路径
    try:
        shutil.rmtree(AUG_IMG_DIR)
    except FileNotFoundError as e:
        a = 1
    mkdir(AUG_IMG_DIR)

    AUGLOOP = 100      # 每张影像增强的数量,可自行设置

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

    # 影像增强
    sometimes = lambda aug: iaa.Sometimes(0.5, aug) #建立lambda表达式
    seq = iaa.Sequential([
        #iaa.Flipud(0.5),  # 50%的概率上下翻转
        #iaa.Fliplr(0.5),  # 50%的概率镜像翻转
        iaa.Sharpen(alpha=(0, 1.0), lightness=(0.75, 1.5)), # 锐化处理
        iaa.Emboss(alpha=(0, 1.0), strength=(0, 2.0)), #浮雕效果
        sometimes(iaa.OneOf([
                    iaa.EdgeDetect(alpha=(0, 0.7)),
                    iaa.DirectedEdgeDetect(
                        alpha=(0, 0.7), direction=(0.0, 1.0)
                    ),
                ])),#边缘检测，将检测到的赋值0或者255然后叠在原图上
        iaa.ContrastNormalization((0.5, 2.0), per_channel=0.5),# 将整个图像的对比度变为原来的一半或者二倍
        iaa.Multiply((0.5, 1.5), per_channel=0.5),
        iaa.Add((-10, 10), per_channel=0.5), # 每个像素随机加减-10到10之间的数
        #iaa.AverageBlur(k=(2, 5)), # 核大小2~7之间，k=((5, 7), (1, 3))时，核高度5~7，宽度1~3
        #iaa.MedianBlur(k=(3, 5)),
        #iaa.Multiply((0.8, 1.5)),  # 像素乘上0.8到1.5之间的数字.
        #iaa.Superpixels(p_replace=(0, 1.0),n_segments=(20, 200)), #超像素增强
        #iaa.GaussianBlur(sigma=(0.5, 0.8)),  # iaa.GaussianBlur(0.5),参数 sigma 表示高斯模糊的强度，一般 sigma = 0 表示没有高斯模糊，sigma = 3.0 表示很强的模糊。采用 tripe 输入，表示 sigma 值在 (0, 0.03) 之间随机取。
        iaa.Affine(translate_px={"x": 0, "y": 0},scale=(1.0, 1.0),rotate=(-90, 90)),     # 旋转角度
        #iaa.CropAndPad(px=(-10, 0), percent=None, pad_mode='constant', pad_cval=0, keep_size=False),
        #iaa.PiecewiseAffine(scale=(0, 0.05), nb_rows=4, nb_cols=4, cval=0)
        #iaa.ElasticTransformation(alpha=(0, 50), sigma=(4.0, 6.0))
    ],
    random_order=True # 随机的顺序把这些操作用在图像上
    )
    suffix = '.jpg'
    # mylist=("".join(x) for  x in  itertools.product("0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ",repeat=16))
    for root, sub_folders, files in os.walk(IMG_DIR):

        for name in files:
            suffix = name[-5:]
            bndbox = read_xml_annotation(XML_DIR, name[:-4] + '.xml')
            shutil.copy(os.path.join(XML_DIR, name[:-4] + '.xml'), AUG_XML_DIR)
            shutil.copy(os.path.join(IMG_DIR, name), AUG_IMG_DIR)

            for epoch in range(AUGLOOP):
                seq_det = seq.to_deterministic()  # 保持坐标和图像同步改变，而不是随机
                # 读取图片
                img = Image.open(os.path.join(IMG_DIR, name))
                # 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])
                # 存储变化后的图片
                addname = str(epoch)
                image_aug = seq_det.augment_images([img])[0]
                path = os.path.join(AUG_IMG_DIR, name[:-4] + "_" + addname + suffix)
                print(path)
                image_auged = bbs.draw_on_image(image_aug, thickness=0)
                Image.fromarray(image_auged).save(path)

                # 存储变化后的XML
                change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR, name[:-4] + "_" + addname + suffix)
                print(name[:-4] + "_" + addname + suffix)
                new_bndbox_list = []