保存检测框结果 生成 xml文件

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

def create_xml(image_name,w,h):
    '''
    <folder>images</folder>
	<filename>org_595_914061.jpg</filename>
	<source>
		<database>Unknown</database>
	</source>
	<size>
		<width>1280</width>
		<height>720</height>
		<depth>3</depth>
	</size>
	<segmented>0</segmented>
    '''
    # ET.parse('test.xml')
    # 创建树根annotation
    annotation = ET.Element('annotation')
    # 创建一级分支folder
    folder = ET.SubElement(annotation, 'folder')
    # 添加folder标签内容
    folder.text = ("VOC2012")

    # 创建一级分支filename
    filename = ET.SubElement(annotation, 'filename')
    filename.text = image_name
    # 创建一级分支path
    path = ET.SubElement(annotation, 'path')
    path.text = image_name  # 用于返回当前工作目录

    # 创建一级分支source
    source = ET.SubElement(annotation, 'source')
    # 创建source下的二级分支database
    database = ET.SubElement(source, 'database')
    database.text = 'Unknown'
    # 创建一级分支size
    size = ET.SubElement(annotation, 'size')
    # 创建size下的二级分支图像的宽、高及depth
    width = ET.SubElement(size, 'width')
    width.text = str(w)
    height = ET.SubElement(size, 'height')
    height.text = str(h)
    depth = ET.SubElement(size, 'depth')
    depth.text = '3'

    # 创建一级分支segmented
    segmented = ET.SubElement(annotation, 'segmented')
    segmented.text = '0'
    return annotation

# 定义一个创建一级分支object的函数
def create_object(root, xi, yi, xa, ya, lab):  # 参数依次，树根，xmin，ymin，xmax，ymax
    '''
	<object>
		<name>D10</name>
		<pose>Unspecified</pose>
		<truncated>0</truncated>
		<difficult>0</difficult>
		<bndbox>
			<xmin>33</xmin>
			<ymin>539</ymin>
			<xmax>647</xmax>
			<ymax>640</ymax>
		</bndbox>
	</object>
    '''
    # 创建一级分支object
    _object = ET.SubElement(root, 'object')
    # 创建二级分支
    name = ET.SubElement(_object, 'name')
    # print(obj_name)
    name.text = str(lab)
    pose = ET.SubElement(_object, 'pose')
    pose.text = 'Unspecified'
    truncated = ET.SubElement(_object, 'truncated')
    truncated.text = '0'
    difficult = ET.SubElement(_object, 'difficult')
    difficult.text = '0'
    # 创建bndbox
    bndbox = ET.SubElement(_object, 'bndbox')
    xmin = ET.SubElement(bndbox, 'xmin')
    xmin.text = '%s' % int(xi)
    ymin = ET.SubElement(bndbox, 'ymin')
    ymin.text = '%s' % int(yi)
    xmax = ET.SubElement(bndbox, 'xmax')
    xmax.text = '%s' % int(xa)
    ymax = ET.SubElement(bndbox, 'ymax')
    ymax.text = '%s' % int(ya)
    return root

def save_xml(image_name, w,h,np_boxes, labels, annotation_path_root, threshold=0.5):  #
    annotation = create_xml(image_name, w, h)  #3.jpg  创建一个基础的xml 还没与box

    expect_boxes = (np_boxes[:, 1] > threshold) & (np_boxes[:, 0] > -1)
    np_boxes = np_boxes[expect_boxes, :]  #过滤阈值的框
    for dt in np_boxes:  #有框就遍历写入到xml中 ，没有就没有
        clsid, bbox, score = int(dt[0]), dt[2:], dt[1]
        lab = labels[clsid]
        xmin, ymin, xmax, ymax = bbox
        annotation = create_object(annotation, xmin, ymin, xmax, ymax, lab)

    tree = ET.ElementTree(annotation)
    xml_path = annotation_path_root+image_name.split(".")[0] +".xml"
    tree.write(xml_path) # 1_2.xml
    # 保存时没有缩进，借用dom，添加缩进，重新保存
    tree = ET.parse(xml_path)  # 方式一
    root = tree.getroot()
    rawtext=ET.tostring(root)
    dom=minidom.parseString(rawtext)
    with open(xml_path, "w") as f:
        dom.writexml(f, indent='\t', addindent='\t', newl='\n', encoding='utf-8')