python实现数据增强（翻转，旋转），并对应修改xml文件

import os
import cv2
import xml.dom.minidom
from xml.dom.minidom import Document
import math# 获取路径下所有文件的完整路径，用于读取文件用
def GetFileFromThisRootDir(dir, ext=None):allfiles = []needExtFilter = (ext != None)for root, dirs, files in os.walk(dir):for filespath in files:filepath = os.path.join(root, filespath)extension = os.path.splitext(filepath)[1][1:]if needExtFilter and extension in ext:allfiles.append(filepath)elif not needExtFilter:allfiles.append(filepath)return allfilesdef limit_value(a, b):if a < 1:a = 1if a > b:a = b - 1return a# 读取xml文件，xmlfile参数表示xml的路径
def readXml(xmlfile):DomTree = xml.dom.minidom.parse(xmlfile)annotation = DomTree.documentElementsizelist = annotation.getElementsByTagName('size')  # [<DOM Element: filename at 0x381f788>]heights = sizelist[0].getElementsByTagName('height')height = int(heights[0].childNodes[0].data)widths = sizelist[0].getElementsByTagName('width')width = int(widths[0].childNodes[0].data)depths = sizelist[0].getElementsByTagName('depth')depth = int(depths[0].childNodes[0].data)objectlist = annotation.getElementsByTagName('object')bboxes = []for objects in objectlist:namelist = objects.getElementsByTagName('name')class_label = namelist[0].childNodes[0].databndbox = objects.getElementsByTagName('bndbox')[0]x1_list = bndbox.getElementsByTagName('xmin')x1 = int(float(x1_list[0].childNodes[0].data))y1_list = bndbox.getElementsByTagName('ymin')y1 = int(float(y1_list[0].childNodes[0].data))x2_list = bndbox.getElementsByTagName('xmax')x2 = int(float(x2_list[0].childNodes[0].data))y2_list = bndbox.getElementsByTagName('ymax')y2 = int(float(y2_list[0].childNodes[0].data))# 这里我box的格式【xmin，ymin，xmax，ymax，classname】bbox = [x1, y1, x2, y2, class_label]bboxes.append(bbox)return bboxes, width, height, depth# 图像旋转用，里面的angle是角度制的
def im_rotate(im, angle, center=None, scale=1.0):h, w = im.shape[:2]if center is None:center = (w / 2, h / 2)M = cv2.getRotationMatrix2D(center, angle, scale)im_rot = cv2.warpAffine(im, M, (w, h))return im_rotdef im_flip(im, method='H'):  # 翻转图像if method == 'H':  # Flipped Horizontally 水平翻转im_flip = cv2.flip(im, 1)elif method == 'V':  # Flipped Vertically 垂直翻转im_flip = cv2.flip(im, 0)# elif method == 'HV':# Flipped Horizontally & Vertically 水平垂直翻转#    im_flip = cv2.flip(im, -1)return im_flipdef flip_image(imgs_path, method, anno_new_dir_, img_new_dir):j=1for img_path in imgs_path:# print(img_path)# 读取原图像j+=1# print(anno_new_dir)im = cv2.imread(img_path)flip_img = im_flip(im, method)  # 翻转(H, W, D) = flip_img.shape  # 得到翻转后的图像的高、宽、深度，用于书写xmlfile_name = os.path.basename(os.path.splitext(img_path)[0])  # 得到原图的名称anno_new_dir = anno_new_dir_ + '\\'+file_name+'H%s' % j# print(anno_new_dir)ext = os.path.splitext(img_path)[-1]  # 得到原图的后缀# 保存翻转后图像new_img_name = '%s_%s%s' % (method, file_name, ext)cv2.imwrite(os.path.join(img_new_dir, new_img_name), flip_img)  # 新的命名方式为H/V+原图名称# 读取anno标签数据，返回相应的信息anno = os.path.join(anno_path, '%s.xml' % file_name)[gts, w, h, d] = readXml(anno)gt_new = []for gt in gts:x1 = gt[0]  # xminy1 = gt[1]  # yminx2 = gt[2]  # xmaxy2 = gt[3]  # ymaxclassname = str(gt[4])if method == 'H':x1 = w - 1 - x1  # xmaxx2 = w - 1 - x2  # xminx1 = limit_value(x1, w)x2 = limit_value(x2, w)gt_new.append([x2, y1, x1, y2, classname])elif method == 'V':y1 = h - 1 - y1  # ymaxy2 = h - 1 - y2  # yminy1 = limit_value(y1, h)y2 = limit_value(y2, h)gt_new.append([x1, y2, x2, y1, classname])writeXml(anno_new_dir, new_img_name, W, H, D, gt_new)def rotate_image(angles, angle_rad, imgs_path, anno_new_dir, img_new_dir):j = 0  # 计数用angle_num = len(angles)for img_path in imgs_path:# 读取原图像im = cv2.imread(img_path)for i in range(angle_num):gt_new = []im_rot = im_rotate(im, angles[i])  # 旋转(H, W, D) = im_rot.shape  # 得到旋转后的图像的高、宽、深度，用于书写xmlfile_name = os.path.basename(os.path.splitext(img_path)[0])  # 得到原图的名称ext = os.path.splitext(img_path)[-1]  # 得到原图的后缀# 保存旋转后图像new_img_name = 'P%s_%s%s' % (angles[i], file_name, ext)cv2.imwrite(os.path.join(img_new_dir, new_img_name), im_rot)  # 新的命名方式为P+角度+原图名称# 读取anno标签数据，返回相应的信息anno = os.path.join(anno_path, '%s.xml' % file_name)[gts, w, h, d] = readXml(anno)# 计算旋转后gt框四点的坐标变换[xc, yc] = [float(w) / 2, float(h) / 2]for gt in gts:# 计算左上角点的变换x1 = (gt[0] - xc) * math.cos(angle_rad[i]) - (yc - gt[1]) * math.sin(angle_rad[i]) + xcif int(x1) <= 0: x1 = 1.0if int(x1) > w - 1: x1 = w - 1y1 = yc - (gt[0] - xc) * math.sin(angle_rad[i]) - (yc - gt[1]) * math.cos(angle_rad[i])if int(y1) <= 0: y1 = 1.0if int(y1) > h - 1: y1 = h - 1# 计算右上角点的变换x2 = (gt[2] - xc) * math.cos(angle_rad[i]) - (yc - gt[1]) * math.sin(angle_rad[i]) + xcif int(x2) <= 0: x2 = 1.0if int(x2) > w - 1: x2 = w - 1y2 = yc - (gt[2] - xc) * math.sin(angle_rad[i]) - (yc - gt[1]) * math.cos(angle_rad[i])if int(y2) <= 0: y2 = 1.0if int(y2) > h - 1: y2 = h - 1# 计算左下角点的变换x3 = (gt[0] - xc) * math.cos(angle_rad[i]) - (yc - gt[3]) * math.sin(angle_rad[i]) + xcif int(x3) <= 0: x3 = 1.0if int(x3) > w - 1: x3 = w - 1y3 = yc - (gt[0] - xc) * math.sin(angle_rad[i]) - (yc - gt[3]) * math.cos(angle_rad[i])if int(y3) <= 0: y3 = 1.0if int(y3) > h - 1: y3 = h - 1# 计算右下角点的变换x4 = (gt[2] - xc) * math.cos(angle_rad[i]) - (yc - gt[3]) * math.sin(angle_rad[i]) + xcif int(x4) <= 0: x4 = 1.0if int(x4) > w - 1: x4 = w - 1y4 = yc - (gt[2] - xc) * math.sin(angle_rad[i]) - (yc - gt[3]) * math.cos(angle_rad[i])if int(y4) <= 0: y4 = 1.0if int(y4) > h - 1: y4 = h - 1xmin = min(x1, x2, x3, x4)xmax = max(x1, x2, x3, x4)ymin = min(y1, y2, y3, y4)ymax = max(y1, y2, y3, y4)# 把因为旋转导致的特别小的 长线型的去掉# w_new = xmax-xmin+1# h_new = ymax-ymin+1# ratio1 = float(w_new)/h_new# ratio2 = float(h_new)/w_new# if(1.0/6.0<ratio1<6 and 1.0/6.0<ratio2<6 and w_new>9 and h_new>9):classname = str(gt[4])gt_new.append([xmin, ymin, xmax, ymax, classname])# 写出新的xmlwriteXml(anno_new_dir, new_img_name, W, H, D, gt_new)j = j + 1if j % 100 == 0: print('----%s----' % j)# 写xml文件，参数中tmp表示路径，imgname是文件名（没有尾缀）ps有尾缀也无所谓
def writeXml(tmp, imgname, w, h, d, bboxes):doc = Document()# ownerannotation = doc.createElement('annotation')doc.appendChild(annotation)# ownerfolder = doc.createElement('folder')annotation.appendChild(folder)folder_txt = doc.createTextNode("VOC2007")folder.appendChild(folder_txt)filename = doc.createElement('filename')annotation.appendChild(filename)filename_txt = doc.createTextNode(imgname)filename.appendChild(filename_txt)# ones#source = doc.createElement('source')annotation.appendChild(source)database = doc.createElement('database')source.appendChild(database)database_txt = doc.createTextNode("My Database")database.appendChild(database_txt)annotation_new = doc.createElement('annotation')source.appendChild(annotation_new)annotation_new_txt = doc.createTextNode("VOC2007")annotation_new.appendChild(annotation_new_txt)image = doc.createElement('image')source.appendChild(image)image_txt = doc.createTextNode("flickr")image.appendChild(image_txt)# ownerowner = doc.createElement('owner')annotation.appendChild(owner)flickrid = doc.createElement('flickrid')owner.appendChild(flickrid)flickrid_txt = doc.createTextNode("NULL")flickrid.appendChild(flickrid_txt)ow_name = doc.createElement('name')owner.appendChild(ow_name)ow_name_txt = doc.createTextNode("idannel")ow_name.appendChild(ow_name_txt)# onee## twos#size = doc.createElement('size')annotation.appendChild(size)width = doc.createElement('width')size.appendChild(width)width_txt = doc.createTextNode(str(w))width.appendChild(width_txt)height = doc.createElement('height')size.appendChild(height)height_txt = doc.createTextNode(str(h))height.appendChild(height_txt)depth = doc.createElement('depth')size.appendChild(depth)depth_txt = doc.createTextNode(str(d))depth.appendChild(depth_txt)# twoe#segmented = doc.createElement('segmented')annotation.appendChild(segmented)segmented_txt = doc.createTextNode("0")segmented.appendChild(segmented_txt)for bbox in bboxes:# threes#object_new = doc.createElement("object")annotation.appendChild(object_new)name = doc.createElement('name')object_new.appendChild(name)name_txt = doc.createTextNode(str(bbox[4]))name.appendChild(name_txt)pose = doc.createElement('pose')object_new.appendChild(pose)pose_txt = doc.createTextNode("Unspecified")pose.appendChild(pose_txt)truncated = doc.createElement('truncated')object_new.appendChild(truncated)truncated_txt = doc.createTextNode("0")truncated.appendChild(truncated_txt)difficult = doc.createElement('difficult')object_new.appendChild(difficult)difficult_txt = doc.createTextNode("0")difficult.appendChild(difficult_txt)# threes-1#bndbox = doc.createElement('bndbox')object_new.appendChild(bndbox)xmin = doc.createElement('xmin')bndbox.appendChild(xmin)xmin_txt = doc.createTextNode(str(bbox[0]))xmin.appendChild(xmin_txt)ymin = doc.createElement('ymin')bndbox.appendChild(ymin)ymin_txt = doc.createTextNode(str(bbox[1]))ymin.appendChild(ymin_txt)xmax = doc.createElement('xmax')bndbox.appendChild(xmax)xmax_txt = doc.createTextNode(str(bbox[2]))xmax.appendChild(xmax_txt)ymax = doc.createElement('ymax')bndbox.appendChild(ymax)ymax_txt = doc.createTextNode(str(bbox[3]))ymax.appendChild(ymax_txt)print(bbox[0], bbox[1], bbox[2], bbox[3])xmlname = os.path.splitext(imgname)[0]tempfile = tmp + ".xml"print(tempfile)with open(tempfile, 'wb') as f:f.write(doc.toprettyxml(indent='\t', encoding='utf-8'))f.close()returnif __name__ == '__main__':# 数据路径root = 'H:\Butterfly_detection\yolov5-1.0-baseline\VOCdevkit\VOC2007'img_dir = root + '/JPEGImages'anno_path = root + '/Annotations'imgs_path = GetFileFromThisRootDir(img_dir)  # 返回每一张原图的路径AUG = 'Flip'  # 数据扩增的方式 Rotate代表旋转，Flip表示翻转# 存储新的anno位置anno_new_dir = os.path.join(root, AUG, 'xml')if not os.path.isdir(anno_new_dir):os.makedirs(anno_new_dir)# 扩增后图片保存的位置img_new_dir = os.path.join(root, AUG, 'images')if not os.path.isdir(img_new_dir):os.makedirs(img_new_dir)if AUG == 'Rotate':# 旋转角的大小，正数表示逆时针旋转angles = [5, 90, 180, 270, 355]  # 角度im_rotate用到的是角度制angle_rad = [angle * math.pi / 180.0 for angle in angles]  # cos三角函数里要用到弧度制的# 开始旋转rotate_image(angles, angle_rad, imgs_path, anno_new_dir, img_new_dir)elif AUG == 'Flip':method = 'H'flip_image(imgs_path, method, anno_new_dir, img_new_dir)

最终版本：

上面那个生成出来的xml并不对应相应图片，因为我加多了一句代码：然后他的名称后面就在不断循环H1,H2…

 anno_new_dir = anno_new_dir_ + '\\'+file_name+'H%s' % j

这个完全是多余的

python实现数据增强（翻转，旋转），并对应修改xml文件相关推荐

【Tool】Augmentor和imgaug——python图像数据增强库
Augmentor和imgaug--python图像数据增强库 Tags: ComputerVision Python 介绍两个图像增强库:Augmentor和imgaug,Augmentor使用比较 ...
python图像增强_【Tool】Augmentor和imgaug——python图像数据增强库
Augmentor和imgaug--python图像数据增强库 Tags: ComputerVision Python 介绍两个图像增强库:Augmentor和imgaug,Augmentor使用比较 ...
python处理xml文件_Python解析并修改XML文件
XML 本质上只是一种数据格式,它的本意并不是管理数据,因此,在 XML 应用中,数据的管理仍然要借助数据库,尤其是当数据量很大.性能要求很高的时候. XML数据库具有以下优势: XML数据库能够对半 ...
python学习(二) ElementTree解析、读写、创建xml文件
python学习(二) 读写xml文件 1.xml格式将其存储为sample.xml Tag:使用<>包围的部分 Element:被Tag包围的部分,例如22中的22 Attrib ...
python修改xml标签的值_对python修改xml文件的节点值方法详解
这是我的xml文件结构 <?xml version='1.0' encoding='utf-8'?> JPEGImages train_2018-05-08_1000.jpg D:\all ...
数据增强-翻转、旋转、随机颜色、对比度增强、亮度增强、颜色增强
# imageDir 为原数据集的存放位置 # saveDir 为数据增强后数据的存放位置 ### #def flip(root_path,img_name): #翻转图像img = Image.op ...
python图像数据增强——imgaug （一）
(一)imgaug基础用法 (二)imgaug进阶示例 (三)imgaug图像分割数据增强在机器学习或者深度学习的问题中,因为受制于图像采集的方式,投入的精力和图像标注的难度等,用于训练的图像数量可 ...
python 离线数据增强（扩充数据集）
增强方式包括: # (一) 针对像素的数据增强 # 1. 改变亮度 # 2. 加噪声 # (二) 针对图像的数据增强 # 3. 裁剪(需改变bbox) # 4. 平移(需改变bbox) # 5. 镜像 ...
目标检测的Tricks | 【Trick8】数据增强——随机旋转、平移、缩放、错切、hsv增强
如有错误,恳请指出. 在之前使用opencv就介绍使用过一些常用的数据增强的实现方法,见:<数据增强 | 旋转.平移.缩放.错切.HSV增强>,当时介绍了旋转.平移.缩放.错切.HSV增强 ...

python实现数据增强（翻转，旋转），并对应修改xml文件

相关代码学习：

缩进的重要性：（基本规范）

最终版本：

python实现数据增强（翻转，旋转），并对应修改xml文件相关推荐

最新文章

热门文章