使用python生成VOC类型数据集样本

voc和coco是常用的数据集，它们有各自的dom结构，本文介绍使用python实现一个voc样本生成类，用于批量将非voc格式样本转化为voc格式

        train
        0013.jpg
        
            1215
            1098
            1
        
        0
    

该类目前适配PaddleDetection VOC格式，部分VOC格式字段有需要可自行添加

from xml.dom import minidom


class VOC_Sample_Generator:

    def __init__(self):
        self.dom = minidom.document()
        self.root_node = self.dom.createElement('annotation')
        self.object_node = self.dom.createElement('object')
        self.root_node.appendChild(self.object_node)

    def add_filename(self, filename):
        text = self.dom.createTextNode(filename)
        filename_node = self.dom.createElement('filename')
        filename_node.appendChild(text)
        self.root_node.appendChild(filename_node)

    def add_class(self, class_name):
        text = self.dom.createTextNode(class_name)
        name_node = self.dom.createElement('name')
        name_node.appendChild(text)
        self.object_node.appendChild(name_node)

    def add_size(self, width, height, depth):
        text = self.dom.createTextNode(str(width))
        width_node = self.dom.createElement('width')
        width_node.appendChild(text)

        text = self.dom.createTextNode(str(height))
        height_node = self.dom.createElement('height')
        height_node.appendChild(text)

        text = self.dom.createTextNode(str(depth))
        depth_node = self.dom.createElement('depth')
        depth_node.appendChild(text)

        size_node = self.dom.createElement('size')

        size_node.appendChild(width_node)
        size_node.appendChild(height_node)
        size_node.appendChild(depth_node)
        self.root_node.appendChild(size_node)


    def add_bndbox(self, xmin, ymin, xmax, ymax):
        text = self.dom.createTextNode(str(xmin))
        xmin_node = self.dom.createElement('xmin')
        xmin_node.appendChild(text)

        text = self.dom.createTextNode(str(ymin))
        ymin_node = self.dom.createElement('ymin')
        ymin_node.appendChild(text)

        text = self.dom.createTextNode(str(xmax))
        xmax_node = self.dom.createElement('xmax')
        xmax_node.appendChild(text)

        text = self.dom.createTextNode(str(ymax))
        ymax_node = self.dom.createElement('ymax')
        ymax_node.appendChild(text)

        bndbox_node = self.dom.createElement('bndbox')

        bndbox_node.appendChild(xmin_node)
        bndbox_node.appendChild(ymin_node)
        bndbox_node.appendChild(xmax_node)
        bndbox_node.appendChild(ymax_node)
        self.object_node.appendChild(bndbox_node)

    def build(self, path):
        self.dom.appendChild(self.root_node)
        with open(path, 'w') as f:
            self.dom.writexml(f, indent='', addindent='t', newl='n', encoding='UTF-8')
        f.close()

# 测试
if __name__ == "__main__":
    for i in range(2):
        voc = VOC_Sample_Generator()
        voc.add_filename('test' + str(i) + '.xml')
        voc.add_size(1024, 712, 3)
        voc.add_class('AE1')
        voc.add_bndbox(0, 1, 2, 3)
        voc.build('.\test' + str(i) + '.xml')

	
	test0.xml
	
		1024
		712
		3