图像处理（CV）数据增强

# coding=utf-8
import os
import shutil
import numpy as np
import imgaug as ia
import xml.etree.ElementTree as ET
from tqdm import tqdm
import cv2
from imgaug import augmenters as iaa
import warnings
from concurrent.futures import ProcessPoolExecutor

warnings.filterwarnings("ignore")

ia.seed(1)


def read_xml_annotation(root, image_id):
    """
    xml文件格式不同可以进行修改
    :param root:
    :param image_id:
    :return:
    """
    xml_file = os.path.join(root, image_id + '.xml')
    tree = ET.parse(xml_file)
    root = tree.getroot()
    bndboxlist = []

    for object in root.findall('object'):
        bndbox = object.find('bndbox')

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

        bndboxlist.append([xmin, ymin, xmax, ymax])

    return bndboxlist


def change_xml_annotation(root, image_id, new_target):
    new_xmin, new_ymin, new_xmax, new_ymax = new_target

    with open(os.path.join(root, str(image_id) + '.xml'), 'r') as in_file:
        tree = ET.parse(in_file)
    xmlroot = tree.getroot()
    object = xmlroot.find('object')
    bndbox = object.find('bndbox')
    for elem, value in zip([bndbox.find('xmin'), bndbox.find('ymin'), bndbox.find('xmax'), bndbox.find('ymax')],
                           [new_xmin, new_ymin, new_xmax, new_ymax]):
        elem.text = str(value)

    tree.write(os.path.join(root, str("%06d" % image_id) + '.xml'))


def change_xml_list_annotation(root, image_id, new_target, saveroot, _id):
    with open(os.path.join(root, str(image_id)), 'r') as in_file:
        tree = ET.parse(in_file)
    elem = tree.find('filename')
    elem.text = _id + '.jpg'
    xmlroot = tree.getroot()

    for i, object in enumerate(xmlroot.findall('object')):
        bndbox = object.find('bndbox')
        for elem, value in zip([bndbox.find('xmin'), bndbox.find('ymin'), bndbox.find('xmax'), bndbox.find('ymax')],
                               new_target[i]):
            elem.text = str(value)

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


def mkdir(path):
    path = path.strip().rstrip("\\")
    if not os.path.exists(path):
        os.makedirs(path)
        print(path + ' 创建成功')
    else:
        print(path + ' 目录已存在')


def process_image(args):
    name, image_id, xml_file, img_path, bndbox, AUGLOOP, seq, AUG_IMG_DIR, AUG_XML_DIR, XML_DIR = args
    shutil.copy(img_path, AUG_IMG_DIR)
    shutil.copy(os.path.join(XML_DIR, xml_file), AUG_XML_DIR)
    img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
    bbs = [ia.BoundingBox(x1=x[0], y1=x[1], x2=x[2], y2=x[3]) for x in bndbox]
    bbs_on_img = ia.BoundingBoxesOnImage(bbs, shape=img.shape)

    for epoch in range(AUGLOOP):
        # 为每个扩增创建一个带有新随机种子的新序列
        seq_det = seq.to_deterministic()  # 如果你想要非确定性的增加，你也可以使用seq.to_stochastic（）

        img_aug, bbs_aug = seq_det.augment(image=img, bounding_boxes=bbs_on_img)
        img_aug_path = os.path.join(AUG_IMG_DIR, f"{image_id}_{epoch}.jpg")
        cv2.imencode('.jpg', img_aug)[1].tofile(img_aug_path)

        new_bndbox_list = [[int(max(1, min(img.shape[1], b.x1))), int(max(1, min(img.shape[0], b.y1))),
                            int(max(1, min(img.shape[1], b.x2))), int(max(1, min(img.shape[0], b.y2)))] for b in
                           bbs_aug.bounding_boxes]
        change_xml_list_annotation(XML_DIR, xml_file, new_bndbox_list, AUG_XML_DIR, f"{image_id}_{epoch}")

    return f"{image_id} processed"


def Enhanced_image(XML_DIR, IMG_DIR, AUG_IMG_DIR, AUG_XML_DIR, AUGLOOP):
    """
    :param XML_DIR:
    :param IMG_DIR:
    :param AUG_IMG_DIR:
    :param AUG_XML_DIR:
    :param AUGLOOP:
    :return:
    """
    # 创建一个顺序执行的图像增强序列
    seq = iaa.Sequential([
        # 以50%的概率上下翻转图像
        iaa.Flipud(0.5),
        # 以50%的概率左右翻转图像
        iaa.Fliplr(0.5),
        # 将图像的亮度乘以一个在1.2到1.5之间的随机值
        iaa.Multiply((1.2, 1.5)),
        # 对图像应用高斯模糊，sigma值在0到3.0之间随机选择
        iaa.GaussianBlur(sigma=(0, 3.0)),
        # 对图像应用仿射变换
        iaa.Affine(
            # 在x轴和y轴上分别随机平移15像素
            translate_px={"x": 15, "y": 15},
            # 随机缩放图像，缩放因子在0.8到0.95之间
            scale=(0.8, 0.95),
            # 随机旋转图像，旋转角度在-30度到30度之间
            rotate=(-30, 30)
        )
    ])

    image_files = [f for f in os.listdir(IMG_DIR) if f.endswith('.jpg')]
    xml_files = [f for f in os.listdir(XML_DIR) if f.endswith('.xml')]

    args_list = []
    for name in image_files:
        image_id = name.split('.')[0]
        xml_file = image_id + '.xml'
        if xml_file in xml_files:
            img_path = os.path.join(IMG_DIR, name)
            bndbox = read_xml_annotation(XML_DIR, image_id)
            args_list.append(
                (name, image_id, xml_file, img_path, bndbox, AUGLOOP, seq, AUG_IMG_DIR, AUG_XML_DIR, XML_DIR))
    # 创建进程池提升效率
    with ProcessPoolExecutor() as executor:
        list(tqdm(executor.map(process_image, args_list), total=len(args_list),
                  desc=f'增强倍数:{AUGLOOP}\t原始图数量:{len(args_list)}\t增强后应有:{len(args_list) * (AUGLOOP + 1)}\t已完成进度',
                  colour="GREEN"))


if __name__ == "__main__":
    """
        用于进行图像和标签的数据增强，增强选项可在Enhanced_image函数中自行修改
        
        IMG_DIR : 图片路径
        XML_DIR : XML标签路径
        
        AUG_IMG_DIR : 增强后的图像保存地址
        AUG_XML_DIR : 增强后的XML标签路径
        
        AUGLOOP ：随机增强的数量
        
        最终的结果将原始数据图片及标签进行图像增强和xml增强后保存到AUG文件地址
                
    """

    IMG_DIR = r"E:\DATA set\img_path"
    XML_DIR = r"E:\DATA set\xml_path"
    AUG_IMG_DIR = r"E:\DATA set\result"
    AUG_XML_DIR = r"E:\DATA set\result"

    try:
        shutil.rmtree(AUG_IMG_DIR)
    except FileNotFoundError:
        pass
    mkdir(AUG_IMG_DIR)

    try:
        shutil.rmtree(AUG_XML_DIR)
    except FileNotFoundError:
        pass
    mkdir(AUG_XML_DIR)

    AUGLOOP = 5  # 选择增强数量

    Enhanced_image(XML_DIR, IMG_DIR, AUG_IMG_DIR, AUG_XML_DIR, AUGLOOP)