使用Python将yolov5的推理保存的txt文本格式转换为XML格式

目标检测是计算机视觉中的基础任务，可以让机器在图像中识别和定位物体。为了高效地标注和存储目标检测数据，我们使用了各种不同的数据格式。其中一种常用的格式是XML结构，通常用于VOC数据集格式。在本文中，我们将探讨一个Python脚本，它可以将目标检测标注从特定的文本文件格式转换为XML格式。

问题描述

假设您有一个图像数据集，同时还有相应的文本文件，其中包含目标检测的标注信息。每个文本文件都提供了有关图像中存在的物体的信息。例如，在每个文本文件中，每行可能包含物体类别、边界框坐标等信息。然而，您希望将这些标注信息转换为XML格式，以便更好地组织和管理数据。

Python脚本实现

以下是一个使用Python编写的脚本，用于将目标检测标注从文本格式转换为XML格式：

# 导入所需库
import os
import cv2
import shutil
from xml.dom.minidom import Document# 定义物体类别
cls_names = ['人', '狗']def makexml(picPath, txtPath, xmlPath):# 删除已有的XML保存路径，并创建新的路径if os.path.exists(xmlPath):shutil.rmtree(xmlPath)os.makedirs(xmlPath)# 获取图片文件列表files = os.listdir(picPath)# 遍历每个图片文件for i, name in enumerate(files):ss = os.path.join(picPath, name)print(ss)img = cv2.imread(ss)Pheight, Pwidth, Pdepth = img.shape# 创建XML文档对象xmlBuilder = Document()annotation = xmlBuilder.createElement("annotation")  # 创建annotation标签xmlBuilder.appendChild(annotation)# 添加folder标签folder = xmlBuilder.createElement("folder")folderContent = xmlBuilder.createTextNode("VOC2007")folder.appendChild(folderContent)annotation.appendChild(folder)# 添加filename标签filename = xmlBuilder.createElement("filename")filenameContent = xmlBuilder.createTextNode(name)filename.appendChild(filenameContent)annotation.appendChild(filename)# 添加size标签size = xmlBuilder.createElement("size")width = xmlBuilder.createElement("width")widthContent = xmlBuilder.createTextNode(str(Pwidth))width.appendChild(widthContent)size.appendChild(width)height = xmlBuilder.createElement("height")heightContent = xmlBuilder.createTextNode(str(Pheight))height.appendChild(heightContent)size.appendChild(height)depth = xmlBuilder.createElement("depth")depthContent = xmlBuilder.createTextNode(str(Pdepth))depth.appendChild(depthContent)size.appendChild(depth)annotation.appendChild(size)# 读取对应的文本文件txt = os.path.join(txtPath, name[0:-3] + "txt")if not os.path.exists(txt):# 若文本文件不存在，则直接保存XML文件并处理下一张图片print(txt)f = open(os.path.join(xmlPath, name[0:-4] + '.xml'), 'w', encoding="utf-8")xmlBuilder.writexml(f, indent='\t', newl='\n', addindent='\t', encoding='utf-8')f.close()continue# 若文本文件存在，则读取文本内容并添加相应的标签txtFile = open(txt)txtList = txtFile.readlines()for i in txtList:oneline = i.strip().split(" ")object = xmlBuilder.createElement("object")picname = xmlBuilder.createElement("name")nameContent = xmlBuilder.createTextNode(cls_names[int(oneline[0])])picname.appendChild(nameContent)object.appendChild(picname)# 添加pose、truncated、difficult等标签# 添加bndbox标签，并计算bounding box的坐标# 将以上标签添加至annotation标签下pose = xmlBuilder.createElement("pose")poseContent = xmlBuilder.createTextNode("Unspecified")pose.appendChild(poseContent)object.appendChild(pose)truncated = xmlBuilder.createElement("truncated")truncatedContent = xmlBuilder.createTextNode("0")truncated.appendChild(truncatedContent)object.appendChild(truncated)difficult = xmlBuilder.createElement("difficult")difficultContent = xmlBuilder.createTextNode("0")difficult.appendChild(difficultContent)object.appendChild(difficult)bndbox = xmlBuilder.createElement("bndbox")xmin = xmlBuilder.createElement("xmin")mathData = int(((float(oneline[1])) * Pwidth + 1) - (float(oneline[3])) * 0.5 * Pwidth)xminContent = xmlBuilder.createTextNode(str(mathData))xmin.appendChild(xminContent)bndbox.appendChild(xmin)ymin = xmlBuilder.createElement("ymin")mathData = int(((float(oneline[2])) * Pheight + 1) - (float(oneline[4])) * 0.5 * Pheight)yminContent = xmlBuilder.createTextNode(str(mathData))ymin.appendChild(yminContent)bndbox.appendChild(ymin)xmax = xmlBuilder.createElement("xmax")mathData = int(((float(oneline[1])) * Pwidth + 1) + (float(oneline[3])) * 0.5 * Pwidth)xmaxContent = xmlBuilder.createTextNode(str(mathData))xmax.appendChild(xmaxContent)bndbox.appendChild(xmax)ymax = xmlBuilder.createElement("ymax")mathData = int(((float(oneline[2])) * Pheight + 1) + (float(oneline[4])) * 0.5 * Pheight)ymaxContent = xmlBuilder.createTextNode(str(mathData))ymax.appendChild(ymaxContent)bndbox.appendChild(ymax)object.appendChild(bndbox)annotation.appendChild(object)# 保存生成的XML文件f = open(os.path.join(xmlPath, name[0:-4] + '.xml'), 'w', encoding="utf-8")xmlBuilder.writexml(f, indent='\t', newl='\n', addindent='\t', encoding='utf-8')f.close()# 调用函数进行转换
makexml("../data/images/", "../runs/detect/exp19/labels/", "../data/all/annotations/")