【深度学习】检测结果转成XML

检测图像标注

• 在图像标注过程中，为了减少大量标注工作，可以预先训练小模型，然后用该模型对未标注图像进行推理，将检测框角点数据转成XML。
• 再利用labelimg对未标注图像进行加载，同时加载生成的XML，可以减少一部分数据标注工作

现有xml

<?xml version="1.0" ?><annotation>
	<folder>fold</folder>
	<filename>000000.jpg</filename>
	<path>C:\Users\LENOVO\Desktop\000000.jpg</path>
	<source>
		<database>Unknown</database>
	</source>
	<size>
		<width>512</width>
		<height>512</height>
		<depth>3</depth>
	</size>
	<segmented>0</segmented>
</annotation>

增加检测框数据后的xml

<?xml version='1.0' encoding='utf-8'?>
<annotation>
	<folder>coalsize</folder>
	<filename>123.jpg</filename>
	<path>E:\coalsize\123.jpg</path>
	<source>
		<database>Unknown</database>
	</source>
	<size>
		<width>512</width>
		<height>512</height>
		<depth>3</depth>
	</size>
	<segmented>0</segmented>
	<object>
		<name>obj</name>
		<pose>Unspecified</pose>
		<truncated>0</truncated>
		<difficult>0</difficult>
		<bndbox>
			<xmin>0</xmin>
			<ymin>0</ymin>
			<xmax>0</xmax>
			<ymax>0</ymax>
		</bndbox>
	</object>
	<object>
		<name>obj</name>
		<pose>Unspecified</pose>
		<truncated>0</truncated>
		<difficult>0</difficult>
		<bndbox>
			<xmin>5</xmin>
			<ymin>5</ymin>
			<xmax>5</xmax>
			<ymax>5</ymax>
		</bndbox>
	</object>
</annotation>

具体实现

import os
from xml.etree.ElementTree import ElementTree,Element,SubElement

def prettyXml(element, indent, newline, level = 0): 
    # 判断element是否有子元素
    if element:
        # 如果element的text没有内容      
        if element.text == None or element.text.isspace():     
            element.text = newline + indent * (level + 1)      
        else:    
            element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * (level + 1)    
    # 此处两行如果把注释去掉，Element的text也会另起一行 
    #else:     
        #element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * level    
    temp = list(element) # 将elemnt转成list    
    for subelement in temp:    
        # 如果不是list的最后一个元素，说明下一个行是同级别元素的起始，缩进应一致
        if temp.index(subelement) < (len(temp) - 1):     
            subelement.tail = newline + indent * (level + 1)    
        else:  # 如果是list的最后一个元素， 说明下一行是母元素的结束，缩进应该少一个    
            subelement.tail = newline + indent * level   
        # 对子元素进行递归操作 
        prettyXml(subelement, indent, newline, level = level + 1)   

class DetectToXML(object):
   def __init__(self):
      self.base_xml_path = "data/000000.xml"
      self.save_name = "test"
      self.detect_roi = [100,100,100,2]
      self.img_name = "123.jpg"
      self.img_path = "E:\\coalsize\\"
      self.tree = ElementTree()
      self.root = self.tree.parse(self.base_xml_path)
   def change_text(self):
      self.root[0].text = "coalsize"
      self.root[1].text = self.img_name
      self.root[2].text = self.img_path + self.img_name
   def add_object(self):
      _object = SubElement(self.root, 'object')
      object_con = {"name":"obj","pose":"Unspecified","truncated":0,"difficult":0,
                     "bndbox":{"xmin":self.detect_roi[0],"ymin":self.detect_roi[1],
                     "xmax":self.detect_roi[2],"ymax":self.detect_roi[3]}}
      for k in object_con.keys():
         tmp = SubElement(_object, k)
         if k == "bndbox":
            for m in object_con[k].keys():
               tmp_ = SubElement(tmp, m)
               tmp_.text = str(object_con[k][m])
            break
         tmp.text = str(object_con[k])
   def save_xml(self):
      prettyXml(self.root, '\t', '\n') 
      self.tree.write(self.save_name+".xml", encoding="utf-8", xml_declaration=True)


def make_xml(detect_roi = [[0,0,0,0],[5,5,5,5]],save_name = "test"):
   dm = DetectToXML()
   dm.save_name = save_name
   if len(detect_roi) == 0:
      return 0
   dm.change_text()
   for xyxy in detect_roi:
      dm.detect_roi = xyxy
      dm.add_object()
   dm.save_xml()

if __name__ == '__main__':
   make_xml()