YOLO2VOC:将YOLO格式的数据转为xml格式的数据

当我们使用训练好的模型进行辅助标记的时候会有将YOLO检测到的数据转为VOC格式的xml的需求

import xmltodict
from xml.dom import minidom
import copy 

# 待填充数据
frame_id = 3 # 视频帧编号
class_list = ['conveyor', 'refueller', 'aircraft', 'lounge', 'dining car', 'front of baggage car', 'tractor'] #类别信息
bbox_list = [[5, 824, 747, 912, 809],
 [3, 882, 549, 1365, 631],
 [4, 536, 768, 1023, 988],
 [1, 846, 687, 983, 747],
 [2, 2, 418, 126, 588],
 [0, 402, 847, 696, 987],
 [5, 844, 688, 984, 750]]


# 准备相关模板
xml_base ="""
<annotation>
    <folder>img</folder>
    <filename></filename>
    <path></path>
    <size>
        <width>1632</width>
        <height>1080</height>
        <depth>3</depth>
    </size>
</annotation>

"""
obj_base ="""
    <object>
        <name></name>
        <bndbox>
            <xmin></xmin>
            <ymin></ymin>
            <xmax></xmax>
            <ymax></ymax>
        </bndbox>
    </object>
"""

# 读取模板并填充基本信息
xml_dict = xmltodict.parse(xml_base)
frame_id_format = "{:>05d}".format(frame_id)
filename = frame_id_format+".jpg"
path = "./img/"+filename
xml_dict["annotation"]["filename"] = filename
xml_dict["annotation"]["path"] = path

# 填入obj数据
xml_dict["annotation"]["object"] = []
obj_template = xmltodict.parse(obj_base)["object"]
for bbox in bbox_list:
    tmp_obj = copy.deepcopy(obj_template)
    cls_id,x1,y1,x2,y2 = bbox
    tmp_obj["name"] = class_list[cls_id]
    tmp_obj["bndbox"]["xmin"] = x1
    tmp_obj["bndbox"]["ymin"] = y1
    tmp_obj["bndbox"]["xmax"] = x2
    tmp_obj["bndbox"]["ymax"] = y2
    xml_dict["annotation"]["object"].append(tmp_obj)

xmlstr = xmltodict.unparse(xml_dict)
xml_pretty_str = minidom.parseString(xmlstr).toprettyxml()


print(xml_pretty_str)