【工具脚本】目标检测数据样本的扩增脚本

参考自：https://blog.csdn.net/Mr_health/article/details/82667538

原文中，writeXml(anno_new_dir, new_img_name, W, H, D, gt_new)的代码缩进有问题，无法保存新的xml。

针对HBB格式的标注样本进行扩增。

HBB格式下目标被标注为xmin，ymin，xmax，ymax。

下面脚本中支持两种数据样本扩增方式：

1）镜像操作（左右镜像、上下镜像变换）。

2）旋转操作，（旋转操作，要注意图像会被裁剪。）

脚本使用方法：

# 设置数据路径
    root = '/data_1/data/标注数据/标注数据20190826'       #数据样本的上级目录的路径
    img_dir = root + '/convert_img'                #存放图像的文件夹名称
    anno_path = root + '/convert_xml'           #存放xml的文件夹名称
    AUG = 'Rotate'      #选择数据扩增方式：Rotate代表旋转，Flip表示翻转

# -*- coding: UTF-8 -*-
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 = 0  # 计数用for img_path in imgs_path:# 读取原图像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])  # 得到原图的名称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)j = j + 1if j % 10 == 0:print ('process:%s...' % j)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 % 10 == 0:print ('process:%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], bbox[4])xmlname = os.path.splitext(imgname)[0]tempfile = tmp + "/%s.xml" % xmlnamewith open(tempfile, 'wb') as f:f.write(doc.toprettyxml(indent='\t', encoding='utf-8'))returnif __name__ == '__main__':# 数据路径root = '/data_1/data/标注数据/标注数据20190826'img_dir = root + '/convert_img'anno_path = root + '/convert_xml'imgs_path = GetFileFromThisRootDir(img_dir)  # 返回每一张原图的路径AUG = 'Rotate'  # 数据扩增的方式 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 = [10, 20, 30, 330, 340, 350]  # 角度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'    # 关于H轴镜像变换flip_image(imgs_path, method, anno_new_dir, img_new_dir)

【工具脚本】目标检测数据样本的扩增脚本相关推荐

深度学习目标检测数据VisDrone2019（to yolo / voc / coco）---MMDetection数据篇
1.VisDrone2019数据集介绍配备摄像头的无人机(或通用无人机)已被快速部署到广泛的应用领域,包括农业.航空摄影.快速交付和监视.因此,从这些平台上收集的视觉数据的自动理解要求越来越高,这使 ...
pascal行人voc_利用Pascal VOC目标检测数据深度学习进行目标检测
利用 Pascal VOC 目标检测数据深度学习进行目标检测穆玉理 [期刊名称] <通讯世界> [年 ( 卷 ), 期] 2018(000)005 [摘要] 随着信息社会的发展 , 尤其 ...
03- 目标检测数据集和标注工具介绍 (目标检测)
要点: 常用数据集和标注工具标注工具 PPOCRLabel github地址:paddleocrlabel 参考文档:目标检测简介 - 知乎一目标检测数据集 1. PASCAL VOC VOC数 ...
堪比Focal Loss！解决目标检测中样本不平衡的无采样方法
训练目标检测模型的一个难点是样本不均衡,特别是正负样本比例严重失衡.目前解决这类问题主要是两种方案(见综述Imbalance Problems in Object Detection: A Revie ...
从目标检测数据集中扣出所需类别进行分类
文章目录 1.获取VOC数据集中两轮车 2 .接着做COCO数据集的分类数据获取 3.YOLO 格式数据 4.openimage数据获取获取标签根据displayname 获取 labelname ...
解决one-stage目标检测正负样本不均衡的另类方法--Gradient Harmonized
正负样本不均衡问题一直是One-stage目标检测中被大家所诟病的地方,He Keming等人提出了Focal Loss来解决这个问题.而AAAI2019上的一篇论文<Gradient Harm ...
[数据集][VOC]高质量的目标检测数据集合集（持续更新）
[1][数据集名称]数据集VOC正版消防灭火器数据集VOC格式-5156张 [数据集信息]数据集格式:Pascal VOC格式(仅包含jpg图片和对应的xml) 图数量(jpg文件个数 xml文件个数 ...
目标检测之样本不平衡问题
样本不平衡问题感觉近期研究的论文相对较多,如:2019 AAAI GHM,2019 CVPR AP-loss, 还有2019 DR loss, 2019 IoU-balanced loss,two-s ...
最全目标检测相关资料整理 (目标检测+数据增强+卷价神经网络+类别不均衡...)
1 小目标检测: 综述: 综述论文Augmentation for small object detection 深度学习笔记(十)Augmentation for small object dete ...
coco 数据集_Tensorflow对COCO目标检测数据预处理
COCO数据集是微软发布的一个大型图像数据集, 专为对象检测.分割.人体关键点检测.语义分割和字幕生成而设计.这个数据集还提供了Matlab, Python 和 Lua 的 API 接口. 该 API ...

【工具脚本】目标检测数据样本的扩增脚本

针对HBB格式的标注样本进行扩增。

【工具脚本】目标检测数据样本的扩增脚本相关推荐

最新文章

热门文章