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
ia.seed(1)
def read_xml_annotation(root, image_id):
in_file = open(os.path.join(root, image_id), encoding="utf-8")
tree = ET.parse(in_file)
root = tree.getroot()
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
# (506.0000, 330.0000, 528.0000, 348.0000) -> (520.4747, 381.5080, 540.5596, 398.6603)
def change_xml_annotation(root, image_id, new_target):
new_xmin = new_target[0]
new_ymin = new_target[1]
new_xmax = new_target[2]
new_ymax = new_target[3]
in_file = open(os.path.join(root, str(image_id) + '.xml'), encoding="utf-8") # 这里root分别由两个意思
tree = ET.parse(in_file)
xmlroot = tree.getroot()
object = xmlroot.find('object')
bndbox = object.find('bndbox')
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)
tree.write(os.path.join(root, str(image_id) + "aug" + '.xml'))
def change_xml_list_annotation(root, image_id, new_target, saveroot, idd, new_w, new_h):
in_file = open(os.path.join(root, str(image_id) + '.xml'), encoding="utf-8") # 这里root分别由两个意思
tree = ET.parse(in_file)
elem = tree.find('filename')
elem.text = (idd + '.jpg')
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.find('size').find('width').text = str(new_w)
tree.find('size').find('height').text = str(new_h)
tree.write(os.path.join(saveroot, idd + '.xml'), encoding="utf-8")
in_file.close()
def mkdir(path):
# 去除首位空格
path = path.strip()
# 去除尾部 \ 符号
path = path.rstrip("\\")
# 判断路径是否存在
# 存在 True
# 不存在 False
isExists = os.path.exists(path)
# 判断结果
if not isExists:
# 如果不存在则创建目录
# 创建目录操作函数
os.makedirs(path)
print(path + ' 创建成功')
return True
else:
# 如果目录存在则不创建,并提示目录已存在
print(path + ' 目录已存在')
return False
if __name__ == "__main__":
IMG_DIR = "C:/Users/haohao/Desktop/11/all/images"
XML_DIR = "C:/Users/haohao/Desktop/11/all/annotations"
AUG_XML_DIR = "C:/Users/haohao/Desktop/11/output1/annotations" # 存储增强后的XML文件夹路径
AUG_IMG_DIR = "C:/Users/haohao/Desktop/11/output1/images" # 存储增强后的影像文件夹路径
try:
shutil.rmtree(AUG_XML_DIR)
except FileNotFoundError as e:
a = 1
os.makedirs(AUG_XML_DIR)
try:
shutil.rmtree(AUG_IMG_DIR)
except FileNotFoundError as e:
a = 1
os.makedirs(AUG_IMG_DIR)
AUGLOOP = 5 # 每张影像增强的数量
boxes_img_aug_list = []
new_bndbox = []
new_bndbox_list = []
# 影像增强
# 产生一个处理图片的Sequential函数原型:iaa.Sequential
seq = iaa.SomeOf(
(1, 2), # 1-2种变形
[ # 选几个执行
# iaa.CropAndPad(px=(0, 40)), # crop images from each side by 0 to 20px (randomly chosen)
# iaa.Cutout(nb_iterations=(1, 10), size=0.02, squared=False,fill_mode="constant", cval=0),
# 提取该通道,乘0.5和1.5之间的一个因子,然后转换回原始色彩空间。
iaa.MultiplyBrightness((1.0, 1.8)),
# 创建一个通过随机因素使图像亮度变差的增强器:
iaa.pillike.EnhanceBrightness(),
# 增亮
iaa.WithBrightnessChannels(
iaa.Add((-50, 50)), to_colorspace=[iaa.CSPACE_Lab, iaa.CSPACE_HSV]),
iaa.LinearContrast((0.7, 1.3)), ##增强或减弱图片的对比度
# iaa.Cartoon(blur_ksize=3, segmentation_size=1.0,saturation=2.0, edge_prevalence=1.0),
# 散焦模糊
iaa.imgcorruptlike.DefocusBlur(severity=2),
iaa.MotionBlur(k=6), # 像素移动模糊
# iaa.Rain(speed=(0.1, 0.3)), # 加雨线
# iaa.Affine(scale=(0.5, 1.5)), # 将图像缩放到原始大小的 50% 到 150%:
# iaa.Clouds(), # 加云
# iaa.Snowflakes(flake_size=(0.3, 0.5), speed=(0.007, 0.03)), # 加雪点
# iaa.Canny(alpha=(0.0, 0.3)),#边缘提取
# iaa.Dropout((0.01, 0.1), per_channel=True), # randomly remove up to 10% of the pixels
# iaa.Flipud(0.5), # 垂直
# iaa.Fliplr(0.5), # 水平 (垂直+水平=镜像)
# iaa.Crop(px=(1, 16), keep_size=False),
# mask 增强
iaa.CoarseSaltAndPepper(0.06, size_percent=(0.01, 0.1)),
# iaa.Multiply((0.7, 1.3)), # change brightness, doesn't affect BBs
iaa.GaussianBlur(sigma=(0.2, 0.5)),
# 参数sigma表示高斯模糊的强度,一般 sigma = 0 表示没有高斯模糊,sigma = 3.0 表示很强的模糊。采用 tripe 输入,表示 sigma 值在 (0, 0.03) 之间随机取。
# iaa.Rot90(1, keep_size=False),
# # 每个区域具有20% 相应的高度和宽度大小(对于非方形图像,这会导致要填充非方形区域)。
# iaa.Cutout(nb_iterations=(5, 15), size=0.08, squared=False),
# iaa.Sequential([iaa.Rot90(1, keep_size=False),#不保持原始图像尺寸
# iaa.CropToFixedSize(width=960, height=960)]),
# # 将所有图像的所有组件设置50 % 为零:
# iaa.TotalDropout(0.10),
# iaa.Affine(
# translate_percent=0,
# scale=1.0,#不缩放
# rotate=90,
# #mode="edge"
# ), # translate by 40/60px on x/y axis, and scale to 50-70%, affects BBs
])
for root, sub_folders, files in os.walk(XML_DIR):
for name in files:
bndbox = read_xml_annotation(XML_DIR, name)
shutil.copy(os.path.join(XML_DIR, name), AUG_XML_DIR)
shutil.copy(os.path.join(IMG_DIR, name[:-4] + '.jpg'), AUG_IMG_DIR)
for epoch in range(AUGLOOP):
error = 0
seq_det = seq.to_deterministic() # 保持坐标和图像同步改变,而不是随机
# 读取图片
img = Image.open(os.path.join(IMG_DIR, name[:-4] + '.jpg'))
# 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)))
n_x1 = int(bbs_aug.bounding_boxes[0].x1)
n_y1 = int(bbs_aug.bounding_boxes[0].y1)
n_x2 = int(bbs_aug.bounding_boxes[0].x2)
n_y2 = int(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
# print(n_x1 >= n_x2,n_y1 >= n_y2)
if n_x1 >= n_x2 or n_y1 >= n_y2 or abs(n_x1 - n_x2) <= 20 or abs(n_y1 - n_y2) <= 20:
print('框不合规范舍去(error):', name)
error = 1
break
# new_bndbox_list.append([n_x2, n_y2, n_x1, n_y1])
new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2])
if error == 1:
continue
# 存储变化后的图片
image_aug = seq_det.augment_images([img])[0]
path = os.path.join(AUG_IMG_DIR, str(name[:-4] + "_a%02d" % epoch) + '.jpg')
print(path)
image_auged = bbs.draw_on_image(image_aug)
Image.fromarray(image_auged).save(path)
print(image_auged.shape) # (高,宽,通道)
new_width = image_auged.shape[1]
new_height = image_auged.shape[0]
# 存储变化后的XML
change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR,
str(name[:-4] + "_a%02d" % epoch), new_width, new_height)
# print(str(name[:-4] + "_a%s"%epoch ) + '.jpg')
new_bndbox_list = []
# break
深度学习,数据增强批量处理代码。
最新推荐文章于 2023-06-23 20:22:45 发布