將自己的數據轉爲VOC數據格式

將自己的數據轉爲VOC數據格式用來目標檢測

from lxml.etree import Element, SubElement, tostring
from xml.dom.minidom import parseString
from xml.dom.minidom import Document

def make_xml(infos):
    for image_id, image_info in infos.items():
        image_name = image_info['image_name']
        boxes = image_info['boxes']
        height, width = image_info['height'], image_info['width']
        channel = 3

        node_root = Element('annotation')
        node_folder = SubElement(node_root, 'folder')
        node_folder.text = 'JPEGImages'
        node_filename = SubElement(node_root, 'filename')
        node_filename.text = image_name
        node_size = SubElement(node_root, 'size')
        node_width = SubElement(node_size, 'width')
        node_width.text = '%s' % width
        node_height = SubElement(node_size, 'height')
        node_height.text = '%s' % height
        node_depth = SubElement(node_size, 'depth')
        node_depth.text = '%s' % channel

        if len(boxes) > 0:
            for box in boxes:
                xmin, ymin, w, h, cls = box
                xmax, ymax = xmin + w - 1, ymin + h - 1
                cls = label_list[cls]
                xmin, ymin, xmax, ymax = int(xmin), int(ymin), int(xmax), int(ymax)
                node_object = SubElement(node_root, 'object')
                node_name = SubElement(node_object, 'name')
                node_name.text = cls

                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 = '%s' % xmin
                node_ymin = SubElement(node_bndbox, 'ymin')
                node_ymin.text = '%s' % ymin
                node_xmax = SubElement(node_bndbox, 'xmax')
                node_xmax.text = '%s' % xmax
                node_ymax = SubElement(node_bndbox, 'ymax')
                node_ymax.text = '%s' % ymax
            xml = tostring(node_root, pretty_print=True)
            dom = parseString(xml)

            save_xml = os.path.join(xml_path, image_name.replace('jpg', 'xml'))
            with open(save_xml, 'wb') as f:
                f.write(xml)