实用脚本分享02：目标检测数据集的数据增强代码分享

最新推荐文章于 2023-03-26 16:50:58 发布

DLNovice

最新推荐文章于 2023-03-26 16:50:58 发布

阅读量4.1k

点赞数 3

分类专栏： Practical Scripts 文章标签：目标检测人工智能计算机视觉 python big data

本文链接：https://blog.csdn.net/qq_55535816/article/details/122504297

版权

Practical Scripts 专栏收录该内容

4 篇文章 0 订阅

订阅专栏

以下脚本可直接从百度网盘下载：
链接：https://pan.baidu.com/s/1ofQgoePombKPh86o7u4wsw
提取码：1234

注意：

以下脚本的使用需要依赖另两个脚本： voc_xml.py （生成标签文件）、utils.py

请大家在使用以下脚本时，在同级目录下创建这两个文件

1、 voc_xml.py ：

from xml.dom.minidom import Document
import xml.etree.ElementTree as ET
import os


def get_xml_tree(xmlfile):
    '''
    获取xml tree
    Args:
        xmlfile: xml文件路径
    return:
        tree:xml tree
    '''
    tree = ET.parse(xmlfile)
    return tree


class CreateXML():
    def __init__(self, img_name, img_w, img_h, img_channels):
        '''
        Args:
            img_name:图片名
            img_w,img_h,img_channels:图片宽、高、通道数
        '''
        self.img_name = img_name
        self.doc = Document()

        self.annotation = self.doc.createElement('annotation')
        self.doc.appendChild(self.annotation)

        folder = self.doc.createElement('folder')
        folder.appendChild(self.doc.createTextNode("AIA AUTO"))
        self.annotation.appendChild(folder)

        filename = self.doc.createElement('filename')
        filename.appendChild(self.doc.createTextNode(img_name))
        self.annotation.appendChild(filename)

        source = self.doc.createElement('source')
        database = self.doc.createElement('database')
        database.appendChild(self.doc.createTextNode('The AUTO Database'))
        anno = self.doc.createElement("annotation")
        anno.appendChild(self.doc.createTextNode("AUTO by zxl"))
        image = self.doc.createElement("JPEGImages")
        image.appendChild(self.doc.createTextNode("flickr"))
        source.appendChild(database)
        source.appendChild(anno)
        source.appendChild(image)
        self.annotation.appendChild(source)

        sizeimage = self.doc.createElement('size')
        imagewidth = self.doc.createElement('width')
        imagewidth.appendChild(self.doc.createTextNode(str(img_w)))
        imageheight = self.doc.createElement('height')
        imageheight.appendChild(self.doc.createTextNode(str(img_h)))
        imagedepth = self.doc.createElement("depth")
        imagedepth.appendChild(self.doc.createTextNode(str(img_channels)))
        sizeimage.appendChild(imagewidth)
        sizeimage.appendChild(imageheight)
        sizeimage.appendChild(imagedepth)
        self.annotation.appendChild(sizeimage)

    def add_object_node(self, obj_name, xmin_v, ymin_v, xmax_v, ymax_v, truncated_v=0, difficult_v=0):
        '''
        添加目标框节点
        obj_name:目标名
        xmin_v,ymin_v,xmax_v,ymax_v:目标框左上右上坐标
        truncated_v:截断程度
        difficult:困难程度
        '''
        obj = self.doc.createElement("object")
        objname = self.doc.createElement("name")
        objname.appendChild(self.doc.createTextNode(obj_name))
        pose = self.doc.createElement("pose")
        pose.appendChild(self.doc.createTextNode("front"))
        truncated = self.doc.createElement("truncated")
        truncated.appendChild(self.doc.createTextNode(str(truncated_v)))
        difficult = self.doc.createElement('difficult')
        difficult.appendChild(self.doc.createTextNode(str(difficult_v)))
        obj.appendChild(objname)
        obj.appendChild(pose)
        obj.appendChild(truncated)
        obj.appendChild(difficult)

        bndbox = self.doc.createElement("bndbox")
        xmin = self.doc.createElement("xmin")
        ymin = self.doc.createElement("ymin")
        xmax = self.doc.createElement("xmax")
        ymax = self.doc.createElement("ymax")
        xmin.appendChild(self.doc.createTextNode(str(xmin_v)))
        ymin.appendChild(self.doc.createTextNode(str(ymin_v)))
        xmax.appendChild(self.doc.createTextNode(str(xmax_v)))
        ymax.appendChild(self.doc.createTextNode(str(ymax_v)))
        bndbox.appendChild(xmin)
        bndbox.appendChild(ymin)
        bndbox.appendChild(xmax)
        bndbox.appendChild(ymax)
        obj.appendChild(bndbox)
        self.annotation.appendChild(obj)

    def save_xml(self, save_path, xml_save_name):
        '''
        save_path:保存路径
        xml_save_name:xml文件保存名字       
        '''
        xml_file = open(os.path.join(save_path, xml_save_name), 'w')
        xml_file.write(self.doc.toprettyxml(indent=' ' * 4))

    def get_doc(self):
        '''
        return:
            doc:xml文件的Document()
        '''
        return self.doc

2、utils.py

import os
import random


def confine(value, v_min, v_max):
    '''
    值的边界限制
    Args:
        value:输入值
        v_min,v_max:最大最小边界
    return:
        value:限制值
    '''
    value = v_min if value < v_min else value
    value = v_max if value > v_max else value
    return value


def fileCountIn(dir):
    '''
    计算文件夹下文件个数
    Args:
        dir:文件目录
    return:
        文件个数
    '''
    return sum([len(files) for root, dirs, files in os.walk(dir)])


def randomChoiceIn(dir):
    '''
    目录下随机选择一个文件
    Args:
        dir:目录
    return:
        filename:随机选择的文件名
    '''
    for root, dirs, files in os.walk(dir):
        index = random.randint(0, len(files) - 1)
        filename = files[index]
    return filename


def calc_rect_area(rect):
    '''计算矩形框面积
    Args:
        rect:矩形框 [xmin,ymin,xmax,ymax]
    return:
        dst:矩形框面积
    '''
    return (rect[2] - rect[0] + 0.001) * (rect[3] - rect[1] + 0.001)


def calc_iou(rect1, rect2):
    '''计算两个矩形框的交并比
    Args:
        rect1,rect2:两个矩形框
    return:
        iou:交并比
    '''
    bd_i = (max(rect1[0], rect2[0]), max(rect1[1], rect2[1]), \
            min(rect1[2], rect2[2]), min(rect1[3], rect2[3]))
    iw = bd_i[2] - bd_i[0] + 0.001
    ih = bd_i[3] - bd_i[1] + 0.001
    iou = 0
    if (iw > 0 and ih > 0):
        ua = calc_rect_area(rect1) + calc_rect_area(rect2) - iw * ih
        iou = iw * ih / ua
    return iou