制作Pascal Voc 格式标柱文件

# -*- coding: utf-8 -*-
import os
import cv2
import shutil
from xml.etree import ElementTree
from xml.etree.ElementTree import Element, SubElement


def pretty_xml(element, indent, newline, level=0):
    """
    通过加入缩进、换行，美化xml
    :param element: Elment类
    :param indent: 用于缩进
    :param newline: 用于换行
    :param level:
    :return:
    """
    if element.text is None or element.text.isspace():  # 如果element的text没有内容
        element.text = newline + indent * (level + 1)
    else:
        element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * (level + 1)

    # 将element转成list
    temp = list(element)
    for ele in temp:
        # 如果不是list的最后一个元素，说明下一个行是同级别元素的起始，缩进应一致
        if temp.index(ele) < (len(temp) - 1):
            ele.tail = newline + indent * (level + 1)
        else:  # 如果是list的最后一个元素， 说明下一行是母元素的结束，缩进应该少一个
            ele.tail = newline + indent * level
        # 对子元素进行递归操作
        pretty_xml(ele, indent, newline, level=level + 1)


def convert_voc_xml(img_path, boxes):
    """
    convert image boxes annotation into VOC-formatted xml file.
    :param img_path: (str) image path
    :param boxes: (list) list of box each one is [x1, y1, x2, y2, text]
    :return: saved xml path
    """
    img = cv2.imread(img_path)
    h, w, c = img.shape
    img_folder, img_name = img_path.split('/')[-2:]

    # 创建一个root element<annotation>
    annotation = Element("annotation")
    # 直接通过SubElement类为<annotation>添加多个子元素<folder><filename><path><source><size><segmented><object>
    folder = SubElement(annotation, "folder")
    folder.text = img_folder
    filename = SubElement(annotation, "filename")
    filename.text = img_name
    path = SubElement(annotation, "path")
    path.text = img_path
    source = SubElement(annotation, "source")
    database = SubElement(source, "database")
    size = SubElement(annotation, "size")
    width = SubElement(size, "width")
    width.text = str(w)
    height = SubElement(size, "height")
    height.text = str(h)
    depth = SubElement(size, "depth")
    depth.text = str(c)
    segmented = SubElement(annotation, "segmented")
    segmented.text = "0"

    for x1, y1, x2, y2, text in boxes:
        object = SubElement(annotation, 'object')

        name = SubElement(object, 'name')
        name.text = text
        pose = SubElement(object, 'pose')
        pose.text = "Unspecified"
        truncated = SubElement(object, 'truncated')
        truncated.text = "0"
        difficult = SubElement(object, 'difficult')
        difficult.text = "0"

        bndbox = SubElement(object, 'bndbox')
        xmin = SubElement(bndbox, 'xmin')
        xmin.text = str(x1)
        ymin = SubElement(bndbox, 'ymin')
        ymin.text = str(y1)
        xmax = SubElement(bndbox, 'xmax')
        xmax.text = str(x2)
        ymax = SubElement(bndbox, 'ymax')
        ymax.text = str(y2)

    # just for beautifying xml file, which can skip
    pretty_xml(annotation, '\t', '\n')
    # for saving xml file
    xml = ElementTree.ElementTree(annotation)

    xml_path = img_path.replace('/images/', '/annotations/').replace('.jpg', '.xml')
    os.makedirs(os.path.dirname(xml_path), exist_ok=True)
    xml.write(xml_path, encoding='utf-8')
    # shutil.copyfile(img_path, xml_path.replace('.xml', '.jpg'))

    return xml_path