针对于目标检测任务实现的扩充脚本,能同时扩充图像和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
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 = np.random.normal(mean, sigma, img_noise.shape)
gaussian_out = img_noise + noise
gaussian_out = np.clip(gaussian_out, 0, 1)
gaussian_out = np.uint8(gaussian_out * 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)
xml_path = os.path.join(xml_dir, filename + '.xml')
tree = ET.parse(xml_path)
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'))
image_PILcontrast = ImageEnhance.Contrast(image_PIL).enhance(1.5)
img_noise = np.array(image_PILcontrast)
img_noise = gaussian_noise(img_noise, 0, 0.1)
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'))
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))