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

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

一般VOC数据集样本格式

<annotation>
        <folder>train</folder>
        <filename>0013.jpg</filename>
        <size>
            <width>1215</width>
            <height>1098</height>
            <depth>1</depth>
        </size>
        <segmented>0</segmented>
    <object>
            <name>AE1</name>
            <pose>Unspecified</pose>
            <truncated>0</truncated>
            <difficult>0</difficult>
            <bndbox>
                <xmin>322</xmin>
                <ymin>92</ymin>
                <xmax>832</xmax>
                <ymax>577</ymax>
            </bndbox>
        </object>
</annotation>

Code

该类目前适配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')

生成结果

<?xml version="1.0" encoding="UTF-8"?>
<annotation>
	<object>
		<name>AE1</name>
		<bndbox>
			<xmin>0</xmin>
			<ymin>1</ymin>
			<xmax>2</xmax>
			<ymax>3</ymax>
		</bndbox>
	</object>
	<filename>test0.xml</filename>
	<size>
		<width>1024</width>
		<height>712</height>
		<depth>3</depth>
	</size>
</annotation>