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

于 2021-10-24 09:30:24 发布
阅读量7.8k 收藏 39
点赞数 4
分类专栏： python实操 文章标签： 1024程序员节 python
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
本文链接：https://blog.csdn.net/weixin_51015047/article/details/120929465
版权 
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 allfiles


def limit_value(a, b):
    if a < 1:
        a = 1
    if a > b:
        a = b - 1
    return a


# 读取xml文件，xmlfile参数表示xml的路径
def readXml(xmlfile):
    DomTree = xml.dom.minidom.parse(xmlfile)
    annotation = DomTree.documentElement
    sizelist = 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].data
        bndbox = 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_rot


def 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_flip


def flip_image(imgs_path, method, anno_new_dir_, img_new_dir):
    j=1
    for 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  # 得到翻转后的图像的高、宽、深度，用于书写xml
        file_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]  # xmin
            y1 = gt[1]  # ymin
            x2 = gt[2]  # xmax
            y2 = gt[3]  # ymax
            classname = str(gt[4])
            if method == 'H':
                x1 = w - 1 - x1  # xmax
                x2 = w - 1 - x2  # xmin
                x1 = limit_value(x1, w)
                x2 = limit_value(x2, w)
                gt_new.append([x2, y1, x1, y2, classname])
            elif method == 'V':
                y1 = h - 1 - y1  # ymax
                y2 = h - 1 - y2  # ymin
                y1 = 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  # 得到旋转后的图像的高、宽、深度，用于书写xml
            file_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]) + xc
                if int(x1) <= 0: x1 = 1.0
                if int(x1) > w - 1: x1 = w - 1
                y1 = yc - (gt[0] - xc) * math.sin(angle_rad[i]) - (yc - gt[1]) * math.cos(angle_rad[i])
                if int(y1) <= 0: y1 = 1.0
                if int(y1) > h - 1: y1 = h - 1
                # 计算右上角点的变换
                x2 = (gt[2] - xc) * math.cos(angle_rad[i]) - (yc - gt[1]) * math.sin(angle_rad[i]) + xc
                if int(x2) <= 0: x2 = 1.0
                if int(x2) > w - 1: x2 = w - 1
                y2 = yc - (gt[2] - xc) * math.sin(angle_rad[i]) - (yc - gt[1]) * math.cos(angle_rad[i])
                if int(y2) <= 0: y2 = 1.0
                if int(y2) > h - 1: y2 = h - 1
                # 计算左下角点的变换
                x3 = (gt[0] - xc) * math.cos(angle_rad[i]) - (yc - gt[3]) * math.sin(angle_rad[i]) + xc
                if int(x3) <= 0: x3 = 1.0
                if int(x3) > w - 1: x3 = w - 1
                y3 = yc - (gt[0] - xc) * math.sin(angle_rad[i]) - (yc - gt[3]) * math.cos(angle_rad[i])
                if int(y3) <= 0: y3 = 1.0
                if int(y3) > h - 1: y3 = h - 1
                # 计算右下角点的变换
                x4 = (gt[2] - xc) * math.cos(angle_rad[i]) - (yc - gt[3]) * math.sin(angle_rad[i]) + xc
                if int(x4) <= 0: x4 = 1.0
                if int(x4) > w - 1: x4 = w - 1
                y4 = yc - (gt[2] - xc) * math.sin(angle_rad[i]) - (yc - gt[3]) * math.cos(angle_rad[i])
                if int(y4) <= 0: y4 = 1.0
                if int(y4) > h - 1: y4 = h - 1
                xmin = 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])
                # 写出新的xml
                writeXml(anno_new_dir, new_img_name, W, H, D, gt_new)
            j = j + 1
            if j % 100 == 0: print('----%s----' % j)


# 写xml文件，参数中tmp表示路径，imgname是文件名（没有尾缀）ps有尾缀也无所谓
def writeXml(tmp, imgname, w, h, d, bboxes):
    doc = Document()
    # owner
    annotation = doc.createElement('annotation')
    doc.appendChild(annotation)
    # owner
    folder = 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)
    # owner
    owner = 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()
    return


if __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)
相关代码学习：
# encoding:utf-8

import cv2
image = cv2.imread("girl.jpg")

# Flipped Horizontally 水平翻转
h_flip = cv2.flip(image, 1)
cv2.imwrite("girl-h.jpg", h_flip)

# Flipped Vertically 垂直翻转
v_flip = cv2.flip(image, 0)
cv2.imwrite("girl-v.jpg", v_flip)

# Flipped Horizontally & Vertically 水平垂直翻转
hv_flip = cv2.flip(image, -1)
cv2.imwrite("girl-hv.jpg", hv_flip)

j+=1  #每次执行完一次代码，j 都会加1
缩进的重要性：（基本规范）
def flip_image(imgs_path, method, anno_new_dir, img_new_dir):
    # j=1
    for 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  # 得到翻转后的图像的高、宽、深度，用于书写xml
        file_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]  # xmin
            y1 = gt[1]  # ymin
            x2 = gt[2]  # xmax
            y2 = gt[3]  # ymax
            classname = str(gt[4])
            if method == 'H':
                x1 = w - 1 - x1  # xmax
                x2 = w - 1 - x2  # xmin
                x1 = limit_value(x1, w)
                x2 = limit_value(x2, w)
                gt_new.append([x2, y1, x1, y2, classname])
            elif method == 'V':
                y1 = h - 1 - y1  # ymax
                y2 = h - 1 - y2  # ymin
                y1 = 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)
    #一开始运行这份代码，他只生成了一个xml文件，而且是最后那张图片的xml文件，那么只可能是前面的都覆盖掉了，然后回归代码。找对应运行的那部分，很明显writeXml(anno_new_dir, new_img_name, W, H, D, gt_new) 这边首先是标黄了，说明是有问题的，他这边是缩进去的，并没有跟随for循环去遍历，所以就导致每生成一个就覆盖一个

最终版本：

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

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

这个完全是多余的

seven_不是赛文
关注
  • 4
    点赞
  • 踩
  • 39
    收藏
    觉得还不错? 一键收藏
  • 9
    评论
专栏目录
评论 9
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

当前余额3.43元 前往充值 >
需支付：10.00元
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付元
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明：

1.余额是钱包充值的虚拟货币，按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载，可以购买VIP、付费专栏及课程。

余额充值