-
Notifications
You must be signed in to change notification settings - Fork 9
/
Copy pathvoc_augument.py
268 lines (228 loc) · 11.1 KB
/
voc_augument.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
'''
VOC格式数据集离线扩增
参考链接：
[1] https://zhuanlan.zhihu.com/p/85292901
[2]
'''
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
from tqdm import tqdm
import imgaug as ia
from imgaug import augmenters as iaa
import argparse
sometimes = lambda aug: iaa.Sometimes(0.5, aug) #建立lambda表达式，
ia.seed(1)
def read_xml_annotation(root, image_id):
'''从xml标注文件所在路径读取标注框
root: xml文件路径
image_id: xml文件名称（带.xml后缀）
'''
in_file = open(os.path.join(root, image_id))
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)
return bndboxlist
def change_xml_list_annotation(root, image_id, new_target, saveroot, id):
'''
root: 原始voc数据集中xml文件所在路径
image_id：xml文件名，带后缀
new_target: 新的bndbox:[[x1,y1,x2,y2],...[],[]]
saveroot: 扩增数据集后xml文件的保存路径
id: 新的xml标注文件名名称（不含后缀）
'''
in_file = open(os.path.join(root, str(image_id) + '.xml')) # 这里root分别由两个意思
tree = ET.parse(in_file)
elem = tree.find('filename')
elem.text = (str("%d" % int(id)) + '.jpg')
xmlroot = tree.getroot()
index = 0
for object in xmlroot.findall('object'): # 找到root节点下的所有country节点
bndbox = object.find('bndbox') # 子节点下节点rank的值
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, str("%d" % int(id)) + '.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 + ' 创建成功')
return True
else:
# 如果目录存在则不创建，并提示目录已存在
print(path + ' 目录已存在')
return False
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--voc-root', type=str, required=True,
help='VOC格式数据集根目录，该目录下必须包含JPEGImages和Annotations这两个文件夹')
parser.add_argument('--aug-dir',type=str, default='VOCAugumented',
help='数据增强后保存的路径，默认在voc-root路径下创建一个文件夹Augument进行保存')
parser.add_argument('--aug_num',type=int, default=3,
help='每张图片进行扩增的次数')
parser.add_argument('--ext', type=str, default='.jpg', help='图像后缀，默认为.jpg')
opt = parser.parse_args()
ext = opt.ext
IMG_DIR = os.path.join(opt.voc_root, "JPEGImages")
XML_DIR = os.path.join(opt.voc_root, "Annotations")
AUGUMENT = os.path.join(os.path.dirname(opt.voc_root), opt.aug_dir)
if not os.path.exists(AUGUMENT):
os.mkdir(AUGUMENT)
# 存储增强后的XML文件夹路径
AUG_XML_DIR = os.path.join(AUGUMENT, "Annotations")
try:
shutil.rmtree(AUG_XML_DIR)
except FileNotFoundError as e:
a = 1
mkdir(AUG_XML_DIR)
# 存储增强后的影像文件夹路径
AUG_IMG_DIR = os.path.join(AUGUMENT, "JPEGImages")
try:
shutil.rmtree(AUG_IMG_DIR)
except FileNotFoundError as e:
a = 1
mkdir(AUG_IMG_DIR)
AUGLOOP = opt.aug_num # 每张影像增强的数量
boxes_img_aug_list = []
new_bndbox = []
new_bndbox_list = []
# 影像增强
'''
seq = iaa.Sequential([
iaa.Flipud(0.2), # vertically flip 20% of all images
iaa.Fliplr(0.5), # 镜像
iaa.Multiply((1.2, 1.5)), # change brightness, doesn't affect BBs
iaa.GaussianBlur(sigma=(0, 3.0)), # iaa.GaussianBlur(0.5),
iaa.Affine(
translate_px={"x": 15, "y": 15},
scale=(0.8, 0.95),
rotate=(-30, 30)
) # translate by 40/60px on x/y axis, and scale to 50-70%, affects BBs
])
'''
seq = iaa.Sequential([
iaa.Sometimes(p=0.5,
#高斯模糊
then_list=[iaa.GaussianBlur(sigma=(0, 0.5))],
#锐化
else_list=[iaa.ContrastNormalization((0.15, 0.75), per_channel=True)]
), #以p的概率执行then_list的增强方法，以1-p的概率执行else_list的增强方法，其中then_list,else_list默认为None
iaa.SomeOf(4,[
# 以下一共10个，随机选7个进行处理，也可以将7改为其他数值，继续对数据集进行扩充
# 边缘检测，将检测到的赋值0或者255然后叠在原图上
# sometimes(iaa.OneOf([
# iaa.EdgeDetect(alpha=(0, 0.7)),
# iaa.DirectedEdgeDetect(
# alpha=(0, 0.7), direction=(0.0, 1.0)
# ),
# ])),
# 将RGB变成灰度图然后乘alpha加在原图上
# iaa.Grayscale(alpha=(0.0, 1.0)),
# 扭曲图像的局部区域
sometimes(iaa.PiecewiseAffine(scale=(0.001, 0.005))),
# 每个像素随机加减-10到10之间的数
iaa.Add((-10, 10), per_channel=0.5),
# 中值模糊
iaa.MedianBlur(k=1, name=None, deterministic=False, random_state=None),
#锐化
iaa.Sharpen(alpha=0, lightness=1, name=None, deterministic=False, random_state=None),
# 从最邻近像素中取均值来扰动。
iaa.AverageBlur(k=1, name=None, deterministic=False, random_state=None),
# 0-0.05的数值，分别乘以图片的宽和高为剪裁的像素个数，保持原尺寸
# iaa.Crop(percent=(0.01, 0.01)),
# iaa.Affine(
# # 对图片进行仿射变化，scale缩放x,y取值，translate_percent左右上下移动
# # rotate为旋转角度，shear为剪切取值范围0-360
# scale={"x": (0.99, 1), "y": (0.99, 1)},
# translate_percent={"x": (-0.01, 0.01), "y": (-0.01, 0.01)},
# rotate=(-1, 1),
# shear=(-1, 1)),
# 20%的图片像素值乘以0.8-1.2中间的数值,用以增加图片明亮度或改变颜色
iaa.Multiply((0.8, 1.2), per_channel=0.2),
# 随机去掉一些像素点, 即把这些像素点变成0。
iaa.Dropout(p=0, per_channel=False, name=None, deterministic=False, random_state=None),
# 浮雕效果
# iaa.Emboss(alpha=0, strength=2, name=None, deterministic=False, random_state=None),
# loc 噪声均值，scale噪声方差，50%的概率，对图片进行添加白噪声并应用于每个通道
iaa.AdditiveGaussianNoise(loc=0, scale=(0.0, 0.1 * 255), per_channel=0.3)],
),
], random_order=True) # 打乱定义图像增强的顺序
number = 1 # 在原来图像数量基础上叠加
for root, sub_folders, files in os.walk(XML_DIR):
for name in tqdm(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] + ext), AUG_IMG_DIR)
for epoch in range(AUGLOOP): # 在每个扩增循环里处理
seq_det = seq.to_deterministic() # 保持坐标和图像同步改变，而不是随机
# 读取图片
img = Image.open(os.path.join(IMG_DIR, name[:-4] + ext))
# 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] # 对bbox扩增
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])
# 存储变化后的图片
image_aug = seq_det.augment_images([img])[0] # 对图像进行扩增
path = os.path.join(AUG_IMG_DIR,
str("%d" % (len(files) + number)) + ext)
# image_auged = bbs.draw_on_image(image_aug, size=0)
Image.fromarray(image_aug).save(path)
# 存储变化后的XML
change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR,
len(files) + number)
# print(str("%d" % (len(files) + number)) + ext)
number += 1
new_bndbox_list = []