针对于目标检测任务实现的扩充脚本，能同时扩充图像和xml文件

'''
还不错的链接：http://t.csdn.cn/gw679
'''

# 导入头文件
import os
import cv2
import numpy as np
import random
import xml.etree.ElementTree as ET

from PIL import Image, ImageFilter, ImageEnhance

# r"G:\pycharmprojects\ScriptCodes_lwd\linepaste\save\save_7\images",   # 原图地址
#               r"G:\pycharmprojects\ScriptCodes_lwd\linepaste\save\save_7\save\save_3\images",   # 图像保存地址
#               r"G:\pycharmprojects\ScriptCodes_lwd\linepaste\save\save_7\xmls",      # 原始xml地址
#               r"G:\pycharmprojects\ScriptCodes_lwd\linepaste\save\save_7\save\save_3\xmls")

image_dir = r"G:\pycharmprojects\ScriptCodes_lwd\linepaste\save\save_7\images"  # 输入数据文件夹
xml_dir = r'G:\pycharmprojects\ScriptCodes_lwd\linepaste\save\save_7\xmls'
save_dir = r"G:\pycharmprojects\ScriptCodes_lwd\linepaste\save\save_7\save\save_3"  # 输出数据文件夹

filter_list = [ImageFilter.BLUR, ImageFilter.CONTOUR, ImageFilter.EDGE_ENHANCE, ImageFilter.EDGE_ENHANCE_MORE,
               ImageFilter.EMBOSS, ImageFilter.FIND_EDGES, ImageFilter.SMOOTH, ImageFilter.SMOOTH_MORE,
               ImageFilter.SHARPEN]

# 调整颜色平衡、对比度、亮度、锐度
enhence_list = [ImageEnhance.Color, ImageEnhance.Contrast, ImageEnhance.Brightness, ImageEnhance.Sharpness]


# 添加高斯噪声
def gaussian_noise(image, mean, sigma):
    img_noise = image.copy()
    # 将图片灰度标准化
    img_noise = img_noise / 255
    # 产生高斯 noise
    noise = np.random.normal(mean, sigma, img_noise.shape)
    # 将噪声和图片叠加
    gaussian_out = img_noise + noise
    # 将超过 1 的置 1，低于 0 的置 0
    gaussian_out = np.clip(gaussian_out, 0, 1)
    # 将图片灰度范围的恢复为 0-255
    gaussian_out = np.uint8(gaussian_out * 255)
    # 将噪声范围搞为 0-255
    # noise = np.uint8(noise*255)
    return gaussian_out  # 这里也会返回噪声，注意返回值


# 添加随机噪声
def random_noise(image, noise_num):
    a = random.randint(3, 6)
    img_noise = image.copy()
    h, w, _ = img_noise.shape

    for i in range(noise_num):
        x = random.randint(0, w - a)  # 随机生成指定范围的整数
        y = random.randint(0, h - a)
        img_noise[y:y + a, x:x + a, :] = random.choice((0, 255))
    return img_noise


if __name__ == '__main__':
    if not os.path.exists(image_dir) or not os.path.exists(xml_dir):
        print('image_dir or xml_dir not exists.')
    else:
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
        son_dir = os.path.join(save_dir, 'save_' + str(len(os.listdir(save_dir)) + 1))
        son_images_dir = os.path.join(son_dir, 'images')
        son_labels_dir = os.path.join(son_dir, 'labels_xml')
        os.makedirs(son_images_dir)
        os.makedirs(son_labels_dir)

        for file in os.listdir(image_dir):
            filename, ext = os.path.splitext(file)
            # image = cv2.imread(os.path.join(image_dir, file))

            xml_path = os.path.join(xml_dir, filename + '.xml')
            tree = ET.parse(xml_path)

            # 锐化1.5+随机噪声
            image_PIL = Image.open(os.path.join(image_dir, file))
            image_PILsharpe = ImageEnhance.Sharpness(image_PIL).enhance(1.5)
            img_noise = np.array(image_PILsharpe)
            img_noise = random_noise(img_noise, 1000)  # 随机噪声
            save_name = 'noise_img_' + str(len(os.listdir(son_images_dir)) + 1)
            cv2.imwrite(os.path.join(son_images_dir, save_name + ext), cv2.cvtColor(img_noise, cv2.COLOR_BGR2RGB))
            tree.write(os.path.join(son_labels_dir, save_name + '.xml'))

            # 对比度增强1.5 + 高斯噪声
            image_PILcontrast = ImageEnhance.Contrast(image_PIL).enhance(1.5)
            img_noise = np.array(image_PILcontrast)
            img_noise = gaussian_noise(img_noise, 0, 0.1)  # 高斯噪声 0.12
            save_name = 'noise_img_' + str(len(os.listdir(son_images_dir)) + 1)
            cv2.imwrite(os.path.join(son_images_dir, save_name + ext), cv2.cvtColor(img_noise, cv2.COLOR_BGR2RGB))
            tree.write(os.path.join(son_labels_dir, save_name + '.xml'))

            # 滤镜边界增强+调亮度1.2
            image = image_PIL.filter(ImageFilter.EDGE_ENHANCE_MORE)
            image = ImageEnhance.Brightness(image).enhance(1.2)
            save_name = 'noise_img_' + str(len(os.listdir(son_images_dir)) + 1)
            image.save(os.path.join(son_images_dir, save_name + ext))
            tree.write(os.path.join(son_labels_dir, save_name + '.xml'))

            print('file:{} done.'.format(file))