生成VOC数据集的xml文件

脚本

采用写入xml的方式来生成voc的数据集标注文件，脚本及使用方法如下

import cv2
import os
import argparse
from lxml.etree import Element, SubElement, tostring
from xml.dom.minidom import parseString

def make_xml(xmin_tuple,ymin_tuple,xmax_tuple,ymax_tuple,image_name,image_shape):
    node_root = Element('annotation')
    # node_folder = SubElement(node_root, 'folder')
    # node_folder.text = 'VOC'

    node_filename = SubElement(node_root, 'filename')
    node_filename.text = image_name

    node_object_num = SubElement(node_root, 'object_num')
    node_object_num.text = str(len(xmin_tuple))

    node_size = SubElement(node_root, 'size')
    node_width = SubElement(node_size, 'width')
    node_width.text = str(image_shape[1])

    node_height = SubElement(node_size, 'height')
    node_height.text = str(image_shape[0])

    for i in range(len(xmin_tuple)):  
        node_object = SubElement(node_root, 'object')
        # node_name = SubElement(node_object, 'name')
        # node_name.text = 'ship'
        # node_difficult = SubElement(node_object, 'difficult')
        # node_difficult.text = '0'

        node_bndbox = SubElement(node_object, 'bndbox')
        node_xmin = SubElement(node_bndbox, 'xmin')
        node_xmin.text = str(xmin_tuple[i])
        node_ymin = SubElement(node_bndbox, 'ymin')
        node_ymin.text = str(ymin_tuple[i])
        node_xmax = SubElement(node_bndbox, 'xmax')
        node_xmax.text = str(xmax_tuple[i])
        node_ymax = SubElement(node_bndbox, 'ymax')
        node_ymax.text = str(ymax_tuple[i])

    xml = tostring(node_root)
    dom = parseString(xml)
    return dom

def saveXmlFiles(src,dst):
    img = cv2.imread(src,0) 

    if len(img.shape) == 3:
        gray_image = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    else:
        gray_image = img

    a,b,status,d = cv2.connectedComponentsWithStats(gray_image)
    # print(a,src,status)
    xmin_tuple = []
    ymin_tuple = []
    xmax_tuple = []
    ymax_tuple = []
    for i in range(1,len(status)):
        x=status[i][0]
        y=status[i][1]
        w=status[i][2]
        h=status[i][3]

        xmin_tuple.append(x)
        ymin_tuple.append(y)
        xmax_tuple.append(x + w)
        ymax_tuple.append(y + h)

    dom = make_xml(xmin_tuple, ymin_tuple, xmax_tuple, ymax_tuple, src.split('/')[-1],img.shape)
    with open(dst,'wb') as f:
        f.write(dom.toprettyxml(indent='\t',encoding = 'utf-8'))

if __name__ == "__main__":
    
    parser = argparse.ArgumentParser()
    parser.add_argument("--inputdir",default='trainlabel',help='要输入的数据目录')
    parser.add_argument("--outputdir",default='Annotations',help='xml存储目录')
    args = parser.parse_args()

    assert os.path.exists(args.inputdir),'The value of inputdir:{} is not found!'.format(args.inputdir)
    
    for dir in sorted(os.listdir(args.inputdir)):
        subdir = os.path.join(args.inputdir,dir)
        outputdir = os.path.join(args.outputdir,dir)
        
        if not os.path.exists(outputdir):
            os.makedirs(outputdir)

        for fname in os.listdir(subdir):
            imagefile = os.path.join(subdir,fname)
            outputfile =  os.path.join(outputdir,fname.replace('.bmp','.xml'))
            
            saveXmlFiles(imagefile,outputfile)

生成的文件如下所示

但我们需要的xml文件是没有第一行<?xml version="1.0" encoding="utf-8"?>标注的。这该如何解决。

取消生成的xml文件第一行

进入导入包from xml.dom.minidom import parseString中的minidom文件，搜索writexml函数，在1789行，注释掉下面这部分即可。

之后重新运行，生成的xml文件就符合我们的要求了。