利用Python将XML标注文件转换为YOLO格式的TXT文件

在计算机视觉领域，目标检测任务中常用的训练数据格式之一是YOLO（You Only Look Once）使用的TXT文件格式。本文将介绍如何使用Python脚本将常见的Pascal VOC XML标注文件转换为YOLO所需的TXT文件格式，并展示如何可视化这些标注结果。

背景知识

Pascal VOC XML 格式

Pascal VOC 数据集使用XML格式来标注图像中的目标位置和类别信息。一个典型的XML文件示例如下：

<annotation>
    <folder>VOC2012</folder>
    <filename>image_0001.jpg</filename>
    <size>
        <width>640</width>
        <height>480</height>
        <depth>3</depth>
    </size>
    <object>
        <name>person</name>
        <bndbox>
            <xmin>150</xmin>
            <ymin>100</ymin>
            <xmax>300</xmax>
            <ymax>250</ymax>
        </bndbox>
    </object>
    <!-- 更多 object 元素 -->
</annotation>

YOLO TXT 格式

YOLO框架要求每个图像对应的标签文件是一个TXT文件，其中每一行代表一个目标框的信息，格式如下：

<class_id> <x_center> <y_center> <width> <height>

• <class_id>: 目标的类别ID，从0开始编号。
• <x_center>, <y_center>: 目标框中心点的归一化坐标（相对于图像宽度和高度的比例）。
• <width>, <height>: 目标框的归一化宽高比例。

实现步骤

1. 导入库

首先，我们需要导入必要的库：

import xml.etree.ElementTree as ET
import sys
import os.path
import cv2

2. 解析XML文件

定义一个 XmlParse 类来读取和解析XML文件：

class XmlParse:
    def __init__(self, file_path):
        self.tree = None
        self.root = None
        self.xml_file_path = file_path
 
    def ReadXml(self):
        try:
            self.tree = ET.parse(self.xml_file_path)
            self.root = self.tree.getroot()
        except Exception as e:
            print("parse xml failed!")
            sys.exit()
        finally:
            return self.tree
 
    def WriteXml(self, destfile):
        dses_xml_file = os.path.abspath(destfile)
        self.tree.write(dses_xml_file, encoding="utf-8", xml_declaration=True)

3. 转换函数

定义 xml2txt 函数，用于将XML文件转换为YOLO格式的TXT文件，并可视化标注结果：

def xml2txt(xml, labels, name_list, img_path):
    for i, j in zip(os.listdir(xml), os.listdir(img_path)):
        p = os.path.join(xml + '/' + i)  
        xml_file = os.path.abspath(p)  
        parse = XmlParse(xml_file)
        tree = parse.ReadXml()  
        root = tree.getroot()  

        W = float(root.find('size').find('width').text)
        H = float(root.find('size').find('height').text)

        fil_name = root.find('filename').text[:-4]
        if not os.path.exists(labels):  
            os.mkdir(labels)
        out = open(labels + './' + fil_name + '.txt', 'w+')
        for obj in root.iter('object'):
            x_min = float(obj.find('bndbox').find('xmin').text)
            x_max = float(obj.find('bndbox').find('xmax').text)
            y_min = float(obj.find('bndbox').find('ymin').text)
            y_max = float(obj.find('bndbox').find('ymax').text)

            xcenter = x_min + (x_max - x_min) / 2
            ycenter = y_min + (y_max - y_min) / 2
            w = x_max - x_min
            h = y_max - y_min

            xcenter = round(xcenter / W, 6)
            ycenter = round(ycenter / H, 6)
            w = round(w / W, 6)
            h = round(h / H, 6)

            class_dict = dict(zip(name_list, range(0, len(name_list))))
            class_name = obj.find('name').text
            if class_name not in name_list:
                pass
            else:
                class_id = class_dict[class_name]
                out.write(str(class_id) + " " + str(xcenter) + " " + str(ycenter) + " " + str(w) + " " + str(h) + "\n")

                # 显示图像并标注边界框
                m = os.path.join(img_path + '/' + j)
                block = cv2.imread(m)
                cv2.rectangle(block, pt1=(int((xcenter - w / 2) * W), int((ycenter - h / 2) * H)),
                              pt2=(int((xcenter + w / 2) * W), int((ycenter + h / 2) * H)),
                              color=(0, 255, 0), thickness=2)
                cv2.imshow('block', block)
                cv2.waitKey(300)

4. 主程序入口

最后，在主程序中调用上述函数进行批量转换：

def folder_Path():
    img_path = './images_new'
    xml_path = './xmls_new'  
    labels = './labels_new'  
    name_list = ['trafficlight', 'stop', 'speedlimit', 'crosswalk']  
    xml2txt(xml_path, labels, name_list, img_path)
 
if __name__ == '__main__':
    folder_Path()

运行结果

假设我们有以下目录结构：

.
├── images_new/
│   ├── image_0001.jpg
│   └── ...
├── xmls_new/
│   ├── image_0001.xml
│   └── ...
└── labels_new/

运行脚本后，labels_new 文件夹中会生成与图像对应的 .txt 文件，每行表示一个目标框的信息。同时，图像会被打开并在屏幕上显示标注的结果。

总结

通过本文介绍的方法，我们可以方便地将Pascal VOC格式的XML标注文件转换为YOLO所需的TXT文件格式，并且能够实时查看标注效果。这大大简化了数据准备的工作流程，有助于提高模型训练的效率。

希望这篇文章对您有所帮助！

---

如果您有任何问题或建议，请随时评论交流。感谢阅读！