[Target Detection] Construct an XML file based on the inference results of the detection model

The data set labeling of target detection is generally marked by the labelimg plug-in. There are many sample types that are easy to label. For these data sets that are easy to label, we often have to label them when iterating the model, which leads to a waste of time. We assume It is when you already have a model that has been trained based on a certain task. At this time, you have new training samples. If you don’t want to laboriously label them one by one, you can use the currently trained model line to these unlabeled samples. Perform inference and save the results of each sample to an xml file.

Advantages: Although not 100% of all labeling is correct, it can help you label a lot of things. At this time, if there is an error, you only need to adjust it. You don’t need to label a large area. The labelimg plugin is an adaptive matching image With the XML file, you only need to have the same name.

Not much to say, share the xml file storage format designed in advance:

import os
from xml.etree import ElementTree as ET
def create_xml(bbox_data,img_info, save_path='./newxml/'):
'''
bbox_data:检测模型推理出来的bbox预测框位置信息
img_info:这样应该是字典里面存储图片名,图片路径,图片宽高,图片通道数
save_path:xml保存路径
'''
        root = ET.Element("annotation")
        folder = ET.SubElement(root, "folder")
        folder.text = 'data'
        filename = ET.SubElement(root, "filename")
        filename.text =img_info["file_name"] 
        path=ET.SubElement(root,"path")
        path.text=img_info["url"]
        source=ET.SubElement(root,"source")
        database=ET.SubElement(source,"database")
        database.text='Unknown'
        size=ET.SubElement(root,"size")
        width=ET.SubElement(size,"width")
        height=ET.SubElement(size,"height")
        depth=ET.SubElement(size,"depth")
        width.text=str(img_info["width1"])
        height.text=str(img_info["height1"])
        depth.text=str(img_info["cn"])
        segmented=ET.SubElement(root,"segmented")
        segmented.text='0'
        # Create object elements
        for bbox in bbox_data:
            object_ = ET.SubElement(root, "object")
            # Create name element
            name = ET.SubElement(object_, "name")
            name.text = "crack"
            # Create bbox elements
            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 = ET.SubElement(object_, "bndbox")
            xmin = ET.SubElement(bndbox, "xmin")
            xmin.text = str(int(bbox[0]))
            ymin = ET.SubElement(bndbox, "ymin")
            ymin.text = str(int(bbox[1]))
            xmax = ET.SubElement(bndbox, "xmax")
            xmax.text = str(int(bbox[2]))
            ymax = ET.SubElement(bndbox, "ymax")
            ymax.text = str(int(bbox[3]))
        tree = ET.ElementTree(root)
        # Save xml tree
        info=img_info["file_name"].split('.')[0]
        savepath=save_path+f"{info}.xml"
        tree.write(savepath, encoding="utf-8", xml_declaration=True) 

xml display:

The saved xml file can be automatically matched using the labelimg plug-in

Guess you like

Origin blog.csdn.net/qq_44992785/article/details/129652409