├── LICENSE ├── README.md └── augmentation.py /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 王泽鹏 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # image_augmentation 2 | Image augmentation and change bounding box at the same time. 3 | 4 | # 安装依赖库和imgaug库 5 | 在训练yolo的时候，对已有数据集已经打好了标签，想要进行数据的增强（数据扩增），可以通过imgaug实现对图片和标签中的boundingbox同时变换。 6 | [imgaug使用文档](https://imgaug.readthedocs.io/en/latest/index.html) 7 | 首先，安装依赖库。 8 | ```python 9 | pip install six numpy scipy matplotlib scikit-image opencv-python imageio` 10 | ``` 11 | 安装imgaug 12 | 13 | ```python 14 | pip install imgaug 15 | ``` 16 | 17 | # Bounding Boxes实现 18 | ## 读取原影像bounding boxes坐标 19 | 读取xml文件并使用ElementTree对xml文件进行解析，找到每个object的坐标值。 20 | 21 | ```python 22 | def change_xml_list_annotation(root, image_id, new_target, saveroot, id): 23 | in_file = open(os.path.join(root, str(image_id) + '.xml')) # 这里root分别由两个意思 24 | tree = ET.parse(in_file) 25 | #修改增强后的xml文件中的filename 26 | elem = tree.find('filename') 27 | elem.text = (str(id) + '.jpg') 28 | xmlroot = tree.getroot() 29 | #修改增强后的xml文件中的path 30 | elem = tree.find('path') 31 | if elem != None: 32 | elem.text = (saveroot + str(id) + '.jpg') 33 | 34 | index = 0 35 | for object in xmlroot.findall('object'): # 找到root节点下的所有country节点 36 | bndbox = object.find('bndbox') # 子节点下节点rank的值 37 | 38 | # xmin = int(bndbox.find('xmin').text) 39 | # xmax = int(bndbox.find('xmax').text) 40 | # ymin = int(bndbox.find('ymin').text) 41 | # ymax = int(bndbox.find('ymax').text) 42 | 43 | new_xmin = new_target[index][0] 44 | new_ymin = new_target[index][1] 45 | new_xmax = new_target[index][2] 46 | new_ymax = new_target[index][3] 47 | 48 | xmin = bndbox.find('xmin') 49 | xmin.text = str(new_xmin) 50 | ymin = bndbox.find('ymin') 51 | ymin.text = str(new_ymin) 52 | xmax = bndbox.find('xmax') 53 | xmax.text = str(new_xmax) 54 | ymax = bndbox.find('ymax') 55 | ymax.text = str(new_ymax) 56 | 57 | index = index + 1 58 | 59 | tree.write(os.path.join(saveroot, str(id + '.xml'))) 60 | ``` 61 | 62 | ## 生成变换序列 63 | 产生一个处理图片的Sequential。 64 | 65 | ```python 66 | # 影像增强 67 | seq = iaa.Sequential([ 68 | iaa.Invert(0.5), 69 | iaa.Fliplr(0.5), # 镜像 70 | iaa.Multiply((1.2, 1.5)), # change brightness, doesn't affect BBs 71 | iaa.GaussianBlur(sigma=(0, 3.0)), # iaa.GaussianBlur(0.5), 72 | iaa.Affine( 73 | translate_px={"x": 15, "y": 15}, 74 | scale=(0.8, 0.95), 75 | ) # translate by 40/60px on x/y axis, and scale to 50-70%, affects BBs 76 | ]) 77 | ``` 78 | 79 | ## bounding box 变化后坐标计算 80 | 先读取该影像对应xml文件，获取所有目标的bounding boxes，然后依次计算每个box变化后的坐标。 81 | 82 | ```python 83 | seq_det = seq.to_deterministic() # 保持坐标和图像同步改变，而不是随机 84 | # 读取图片 85 | img = Image.open(os.path.join(IMG_DIR, name[:-4] + '.jpg')) 86 | # sp = img.size 87 | img = np.asarray(img) 88 | # bndbox 坐标增强 89 | for i in range(len(bndbox)): 90 | bbs = ia.BoundingBoxesOnImage([ 91 | ia.BoundingBox(x1=bndbox[i][0], y1=bndbox[i][1], x2=bndbox[i][2], y2=bndbox[i][3]), 92 | ], shape=img.shape) 93 | 94 | bbs_aug = seq_det.augment_bounding_boxes([bbs])[0] 95 | boxes_img_aug_list.append(bbs_aug) 96 | 97 | # new_bndbox_list:[[x1,y1,x2,y2],...[],[]] 98 | n_x1 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x1))) 99 | n_y1 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y1))) 100 | n_x2 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x2))) 101 | n_y2 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y2))) 102 | if n_x1 == 1 and n_x1 == n_x2: 103 | n_x2 += 1 104 | if n_y1 == 1 and n_y2 == n_y1: 105 | n_y2 += 1 106 | if n_x1 >= n_x2 or n_y1 >= n_y2: 107 | print('error', name) 108 | new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2]) 109 | # 存储变化后的图片 110 | image_aug = seq_det.augment_images([img])[0] 111 | path = os.path.join(AUG_IMG_DIR, 112 | str(str(name[:-4]) + '_' + str(epoch)) + '.jpg') 113 | image_auged = bbs.draw_on_image(image_aug, thickness=0) 114 | Image.fromarray(image_auged).save(path) 115 | 116 | # 存储变化后的XML 117 | change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR, 118 | str(name[:-4]) + '_' + str(epoch)) 119 | # print(str(str(name[:-4]) + '_' + str(epoch)) + '.jpg') 120 | new_bndbox_list = [] 121 | ``` 122 | 123 | # 使用示例 124 | ## 数据准备 125 | 输入数据为两个文件夹一个是需要增强的影像数据（JPEGImages），一个是对应的xml文件（Annotations）。注意：影像文件名需和xml文件名相对应！ 126 | ![在这里插入图片描述](https://img-blog.csdnimg.cn/20210325164244643.png)![在这里插入图片描述](https://img-blog.csdnimg.cn/20210325164202912.png) 127 | ## 设置文件路径 128 | 129 | ```python 130 | IMG_DIR = "./VOCdevkit/VOC2007/JPEGImages" 131 | XML_DIR = "./VOCdevkit/VOC2007/Annotations" 132 | 133 | AUG_XML_DIR = "./AUG/Annotations" # 存储增强后的XML文件夹路径 134 | try: 135 | shutil.rmtree(AUG_XML_DIR) 136 | except FileNotFoundError as e: 137 | a = 1 138 | mkdir(AUG_XML_DIR) 139 | 140 | AUG_IMG_DIR = "./AUG/JPEGImages" # 存储增强后的影像文件夹路径 141 | try: 142 | shutil.rmtree(AUG_IMG_DIR) 143 | except FileNotFoundError as e: 144 | a = 1 145 | mkdir(AUG_IMG_DIR) 146 | ``` 147 | ## 设置增强次数 148 | 149 | ```python 150 | AUGLOOP = 10 # 每张影像增强的数量 151 | ``` 152 | ## 设置增强参数 153 | 通过修改Sequential函数参数进行设置，具体设置参考[imgaug使用文档](https://imgaug.readthedocs.io/en/latest/index.html) 154 | 155 | ```python 156 | seq = iaa.Sequential([ 157 | iaa.Invert(0.5), 158 | iaa.Fliplr(0.5), # 镜像 159 | iaa.Multiply((1.2, 1.5)), # change brightness, doesn't affect BBs 160 | iaa.GaussianBlur(sigma=(0, 3.0)), # iaa.GaussianBlur(0.5), 161 | iaa.Affine( 162 | translate_px={"x": 15, "y": 15}, 163 | scale=(0.8, 0.95), 164 | ) # translate by 40/60px on x/y axis, and scale to 50-70%, affects BBs 165 | ]) 166 | 167 | ``` 168 | ## 修改xml文件中filename和path 169 | 170 | ```python 171 | tree = ET.parse(in_file) 172 | #修改xml中的filename 173 | elem = tree.find('filename') 174 | elem.text = (str(id) + '.jpg') 175 | xmlroot = tree.getroot() 176 | #修改xml中的path 177 | elem = tree.find('path') 178 | elem.text = ('C:\working\yolo3\VOCdevkit\VOC2007\JPEGImages\\' + str(id) + '.jpg') 179 | xmlroot = tree.getroot() 180 | ``` 181 | 182 | 183 | ## 输出 184 | 运行augmentation.py ，运行结束后即可得到增强的影像和对应的xml文件夹。 185 | ![在这里插入图片描述](https://img-blog.csdnimg.cn/20210325164837254.png) 186 | ![在这里插入图片描述](https://img-blog.csdnimg.cn/20210325164900829.png) 187 | 188 | 189 | # 完整代码 190 | 191 | ```python 192 | import xml.etree.ElementTree as ET 193 | import pickle 194 | import os 195 | from os import getcwd 196 | import numpy as np 197 | from PIL import Image 198 | import shutil 199 | import matplotlib.pyplot as plt 200 | 201 | import imgaug as ia 202 | from imgaug import augmenters as iaa 203 | 204 | 205 | ia.seed(1) 206 | 207 | 208 | def read_xml_annotation(root, image_id): 209 | in_file = open(os.path.join(root, image_id)) 210 | tree = ET.parse(in_file) 211 | root = tree.getroot() 212 | bndboxlist = [] 213 | 214 | for object in root.findall('object'): # 找到root节点下的所有country节点 215 | bndbox = object.find('bndbox') # 子节点下节点rank的值 216 | 217 | xmin = int(bndbox.find('xmin').text) 218 | xmax = int(bndbox.find('xmax').text) 219 | ymin = int(bndbox.find('ymin').text) 220 | ymax = int(bndbox.find('ymax').text) 221 | # print(xmin,ymin,xmax,ymax) 222 | bndboxlist.append([xmin, ymin, xmax, ymax]) 223 | # print(bndboxlist) 224 | 225 | bndbox = root.find('object').find('bndbox') 226 | return bndboxlist 227 | 228 | 229 | 230 | def change_xml_list_annotation(root, image_id, new_target, saveroot, id): 231 | in_file = open(os.path.join(root, str(image_id) + '.xml')) # 这里root分别由两个意思 232 | tree = ET.parse(in_file) 233 | #修改xml中的filename 234 | elem = tree.find('filename') 235 | elem.text = (str(id) + '.jpg') 236 | xmlroot = tree.getroot() 237 | #修改xml中的path 238 | elem = tree.find('path') 239 | elem.text = ('C:\working\yolo3\VOCdevkit\VOC2007\JPEGImages\\' + str(id) + '.jpg') 240 | xmlroot = tree.getroot() 241 | 242 | index = 0 243 | for object in xmlroot.findall('object'): # 找到root节点下的所有country节点 244 | bndbox = object.find('bndbox') # 子节点下节点rank的值 245 | 246 | # xmin = int(bndbox.find('xmin').text) 247 | # xmax = int(bndbox.find('xmax').text) 248 | # ymin = int(bndbox.find('ymin').text) 249 | # ymax = int(bndbox.find('ymax').text) 250 | 251 | new_xmin = new_target[index][0] 252 | new_ymin = new_target[index][1] 253 | new_xmax = new_target[index][2] 254 | new_ymax = new_target[index][3] 255 | 256 | xmin = bndbox.find('xmin') 257 | xmin.text = str(new_xmin) 258 | ymin = bndbox.find('ymin') 259 | ymin.text = str(new_ymin) 260 | xmax = bndbox.find('xmax') 261 | xmax.text = str(new_xmax) 262 | ymax = bndbox.find('ymax') 263 | ymax.text = str(new_ymax) 264 | 265 | index = index + 1 266 | 267 | tree.write(os.path.join(saveroot, str(id + '.xml'))) 268 | 269 | 270 | def mkdir(path): 271 | # 去除首位空格 272 | path = path.strip() 273 | # 去除尾部 \ 符号 274 | path = path.rstrip("\\") 275 | # 判断路径是否存在 276 | # 存在 True 277 | # 不存在 False 278 | isExists = os.path.exists(path) 279 | # 判断结果 280 | if not isExists: 281 | # 如果不存在则创建目录 282 | # 创建目录操作函数 283 | os.makedirs(path) 284 | print(path + ' 创建成功') 285 | return True 286 | else: 287 | # 如果目录存在则不创建，并提示目录已存在 288 | print(path + ' 目录已存在') 289 | return False 290 | 291 | 292 | if __name__ == "__main__": 293 | 294 | IMG_DIR = "./VOCdevkit/VOC2007/JPEGImages" 295 | XML_DIR = "./VOCdevkit/VOC2007/Annotations" 296 | 297 | AUG_XML_DIR = "./AUG/Annotations" # 存储增强后的XML文件夹路径 298 | try: 299 | shutil.rmtree(AUG_XML_DIR) 300 | except FileNotFoundError as e: 301 | a = 1 302 | mkdir(AUG_XML_DIR) 303 | 304 | AUG_IMG_DIR = "./AUG/JPEGImages" # 存储增强后的影像文件夹路径 305 | try: 306 | shutil.rmtree(AUG_IMG_DIR) 307 | except FileNotFoundError as e: 308 | a = 1 309 | mkdir(AUG_IMG_DIR) 310 | 311 | AUGLOOP = 10 # 每张影像增强的数量 312 | 313 | boxes_img_aug_list = [] 314 | new_bndbox = [] 315 | new_bndbox_list = [] 316 | 317 | # 影像增强 318 | seq = iaa.Sequential([ 319 | iaa.Invert(0.5), 320 | iaa.Fliplr(0.5), # 镜像 321 | iaa.Multiply((1.2, 1.5)), # change brightness, doesn't affect BBs 322 | iaa.GaussianBlur(sigma=(0, 3.0)), # iaa.GaussianBlur(0.5), 323 | iaa.Affine( 324 | translate_px={"x": 15, "y": 15}, 325 | scale=(0.8, 0.95), 326 | ) # translate by 40/60px on x/y axis, and scale to 50-70%, affects BBs 327 | ]) 328 | 329 | for root, sub_folders, files in os.walk(XML_DIR): 330 | 331 | for name in files: 332 | 333 | bndbox = read_xml_annotation(XML_DIR, name) 334 | shutil.copy(os.path.join(XML_DIR, name), AUG_XML_DIR) 335 | shutil.copy(os.path.join(IMG_DIR, name[:-4] + '.jpg'), AUG_IMG_DIR) 336 | 337 | for epoch in range(AUGLOOP): 338 | seq_det = seq.to_deterministic() # 保持坐标和图像同步改变，而不是随机 339 | # 读取图片 340 | img = Image.open(os.path.join(IMG_DIR, name[:-4] + '.jpg')) 341 | # sp = img.size 342 | img = np.asarray(img) 343 | # bndbox 坐标增强 344 | for i in range(len(bndbox)): 345 | bbs = ia.BoundingBoxesOnImage([ 346 | ia.BoundingBox(x1=bndbox[i][0], y1=bndbox[i][1], x2=bndbox[i][2], y2=bndbox[i][3]), 347 | ], shape=img.shape) 348 | 349 | bbs_aug = seq_det.augment_bounding_boxes([bbs])[0] 350 | boxes_img_aug_list.append(bbs_aug) 351 | 352 | # new_bndbox_list:[[x1,y1,x2,y2],...[],[]] 353 | n_x1 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x1))) 354 | n_y1 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y1))) 355 | n_x2 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x2))) 356 | n_y2 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y2))) 357 | if n_x1 == 1 and n_x1 == n_x2: 358 | n_x2 += 1 359 | if n_y1 == 1 and n_y2 == n_y1: 360 | n_y2 += 1 361 | if n_x1 >= n_x2 or n_y1 >= n_y2: 362 | print('error', name) 363 | new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2]) 364 | # 存储变化后的图片 365 | image_aug = seq_det.augment_images([img])[0] 366 | path = os.path.join(AUG_IMG_DIR, 367 | str(str(name[:-4]) + '_' + str(epoch)) + '.jpg') 368 | image_auged = bbs.draw_on_image(image_aug, thickness=0) 369 | Image.fromarray(image_auged).save(path) 370 | 371 | # 存储变化后的XML 372 | change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR, 373 | str(name[:-4]) + '_' + str(epoch)) 374 | # print(str(str(name[:-4]) + '_' + str(epoch)) + '.jpg') 375 | new_bndbox_list = [] 376 | ``` 377 | -------------------------------------------------------------------------------- /augmentation.py: -------------------------------------------------------------------------------- 1 | import xml.etree.ElementTree as ET 2 | import pickle 3 | import os 4 | from os import getcwd 5 | import numpy as np 6 | from PIL import Image 7 | import shutil 8 | import matplotlib.pyplot as plt 9 | 10 | import imgaug as ia 11 | from imgaug import augmenters as iaa 12 | 13 | 14 | ia.seed(1) 15 | 16 | 17 | def read_xml_annotation(root, image_id): 18 | in_file = open(os.path.join(root, image_id)) 19 | tree = ET.parse(in_file) 20 | root = tree.getroot() 21 | bndboxlist = [] 22 | 23 | for object in root.findall('object'): # 找到root节点下的所有country节点 24 | bndbox = object.find('bndbox') # 子节点下节点rank的值 25 | 26 | xmin = int(bndbox.find('xmin').text) 27 | xmax = int(bndbox.find('xmax').text) 28 | ymin = int(bndbox.find('ymin').text) 29 | ymax = int(bndbox.find('ymax').text) 30 | # print(xmin,ymin,xmax,ymax) 31 | bndboxlist.append([xmin, ymin, xmax, ymax]) 32 | # print(bndboxlist) 33 | 34 | bndbox = root.find('object').find('bndbox') 35 | return bndboxlist 36 | 37 | 38 | def change_xml_list_annotation(root, image_id, new_target, saveroot, id): 39 | in_file = open(os.path.join(root, str(image_id) + '.xml')) # 这里root分别由两个意思 40 | tree = ET.parse(in_file) 41 | #修改xml中的folder 42 | elem = tree.find('folder') 43 | elem.text = ('JPEGImages') 44 | #修改xml中的filename 45 | elem = tree.find('filename') 46 | elem.text = (str(id) + '.jpg') 47 | #修改xml中的path 48 | elem = tree.find('path') 49 | elem.text = ('C:\working\yolo3\VOCdevkit\VOC2007\JPEGImages\\' + str(id) + '.jpg') 50 | xmlroot = tree.getroot() 51 | 52 | index = 0 53 | for object in xmlroot.findall('object'): # 找到root节点下的所有country节点 54 | bndbox = object.find('bndbox') # 子节点下节点rank的值 55 | 56 | # xmin = int(bndbox.find('xmin').text) 57 | # xmax = int(bndbox.find('xmax').text) 58 | # ymin = int(bndbox.find('ymin').text) 59 | # ymax = int(bndbox.find('ymax').text) 60 | 61 | new_xmin = new_target[index][0] 62 | new_ymin = new_target[index][1] 63 | new_xmax = new_target[index][2] 64 | new_ymax = new_target[index][3] 65 | 66 | xmin = bndbox.find('xmin') 67 | xmin.text = str(new_xmin) 68 | ymin = bndbox.find('ymin') 69 | ymin.text = str(new_ymin) 70 | xmax = bndbox.find('xmax') 71 | xmax.text = str(new_xmax) 72 | ymax = bndbox.find('ymax') 73 | ymax.text = str(new_ymax) 74 | 75 | index = index + 1 76 | 77 | tree.write(os.path.join(saveroot, str(id + '.xml'))) 78 | 79 | 80 | def mkdir(path): 81 | # 去除首位空格 82 | path = path.strip() 83 | # 去除尾部 \ 符号 84 | path = path.rstrip("\\") 85 | # 判断路径是否存在 86 | # 存在 True 87 | # 不存在 False 88 | isExists = os.path.exists(path) 89 | # 判断结果 90 | if not isExists: 91 | # 如果不存在则创建目录 92 | # 创建目录操作函数 93 | os.makedirs(path) 94 | print(path + ' 创建成功') 95 | return True 96 | else: 97 | # 如果目录存在则不创建，并提示目录已存在 98 | print(path + ' 目录已存在') 99 | return False 100 | 101 | 102 | if __name__ == "__main__": 103 | 104 | IMG_DIR = "./VOCdevkit/VOC2007/JPEGImages" 105 | 106 | XML_DIR = "./VOCdevkit/VOC2007/Annotations" 107 | AUG_XML_DIR = "./AUG/Annotations" # 存储增强后的XML文件夹路径 108 | try: 109 | shutil.rmtree(AUG_XML_DIR) 110 | except FileNotFoundError as e: 111 | a = 1 112 | mkdir(AUG_XML_DIR) 113 | 114 | AUG_IMG_DIR = "./AUG/JPEGImages" # 存储增强后的影像文件夹路径 115 | try: 116 | shutil.rmtree(AUG_IMG_DIR) 117 | except FileNotFoundError as e: 118 | a = 1 119 | mkdir(AUG_IMG_DIR) 120 | 121 | AUGLOOP = 10 # 每张影像增强的数量 122 | 123 | boxes_img_aug_list = [] 124 | new_bndbox = [] 125 | new_bndbox_list = [] 126 | 127 | # 影像增强 128 | seq = iaa.Sequential([ 129 | iaa.Invert(0.5), 130 | iaa.Fliplr(0.5), # 镜像 131 | iaa.Multiply((1.2, 1.5)), # change brightness, doesn't affect BBs 132 | iaa.GaussianBlur(sigma=(0, 3.0)), # iaa.GaussianBlur(0.5), 133 | iaa.Affine( 134 | translate_px={"x": 15, "y": 15}, 135 | scale={"x": (0.7, 1), "y": (1, 1)}, 136 | ) # translate by 40/60px on x/y axis, and scale to 50-70%, affects BBs 137 | ]) 138 | 139 | for root, sub_folders, files in os.walk(XML_DIR): 140 | 141 | for name in files: 142 | 143 | bndbox = read_xml_annotation(XML_DIR, name) 144 | shutil.copy(os.path.join(XML_DIR, name), AUG_XML_DIR) 145 | shutil.copy(os.path.join(IMG_DIR, name[:-4] + '.jpg'), AUG_IMG_DIR) 146 | 147 | for epoch in range(AUGLOOP): 148 | seq_det = seq.to_deterministic() # 保持坐标和图像同步改变，而不是随机 149 | # 读取图片 150 | img = Image.open(os.path.join(IMG_DIR, name[:-4] + '.jpg')) 151 | # sp = img.size 152 | img = np.asarray(img) 153 | # bndbox 坐标增强 154 | for i in range(len(bndbox)): 155 | bbs = ia.BoundingBoxesOnImage([ 156 | ia.BoundingBox(x1=bndbox[i][0], y1=bndbox[i][1], x2=bndbox[i][2], y2=bndbox[i][3]), 157 | ], shape=img.shape) 158 | 159 | bbs_aug = seq_det.augment_bounding_boxes([bbs])[0] 160 | boxes_img_aug_list.append(bbs_aug) 161 | 162 | # new_bndbox_list:[[x1,y1,x2,y2],...[],[]] 163 | n_x1 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x1))) 164 | n_y1 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y1))) 165 | n_x2 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x2))) 166 | n_y2 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y2))) 167 | if n_x1 == 1 and n_x1 == n_x2: 168 | n_x2 += 1 169 | if n_y1 == 1 and n_y2 == n_y1: 170 | n_y2 += 1 171 | if n_x1 >= n_x2 or n_y1 >= n_y2: 172 | print('error', name) 173 | new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2]) 174 | # 存储变化后的图片 175 | image_aug = seq_det.augment_images([img])[0] 176 | path = os.path.join(AUG_IMG_DIR, 177 | str(str(name[:-4]) + '_' + str(epoch)) + '.jpg') 178 | image_auged = bbs.draw_on_image(image_aug, thickness=0) 179 | Image.fromarray(image_auged).save(path) 180 | 181 | # 存储变化后的XML 182 | change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR, 183 | str(name[:-4]) + '_' + str(epoch)) 184 | # print(str(str(name[:-4]) + '_' + str(epoch)) + '.jpg') 185 | new_bndbox_list = [] 186 | --------------------------------------------------------------------------------