将 VOC 格式 XML 转换为 YOLO 格式 TXT

目录

1. 导入必要的模块

2. 定义类别名称

3. 设置文件路径

完整代码

---

1. 导入必要的模块

import os
import xml.etree.ElementTree as ET

os：用于文件和目录操作，例如创建目录、遍历文件等。

xml.etree.ElementTree：用于解析XML文件，从中提取信息。

2. 定义类别名称

class_names = ['nest', 'balloon', 'kite', 'trash']

这是一个列表，定义了数据集中所有物体的类别名称。类别名称的顺序非常重要，因为它们的索引（从0开始）将作为YOLO格式中的class_id。

3. 设置文件路径

xmlpath = 'C:/Users/10431/Desktop/4517+VOC/Annotations/Annotations/'
txtpath = 'C:/Users/10431/Desktop/4517+VOC/Annotations/yolo/'

xmlpath：VOC格式的XML文件所在的目录路径。

txtpath：转换后的YOLO格式TXT文件将保存的目录路径。

完整代码

import os
import xml.etree.ElementTree as ET

# 定义类别名称
class_names = ['nest', 'balloon', 'kite', 'trash']

# 设置输入和输出路径
xmlpath = 'C:/Users/10431/Desktop/4517+VOC/Annotations/Annotations/'
txtpath = 'C:/Users/10431/Desktop/4517+VOC/Annotations/yolo/'

# 如果输出目录不存在，则创建
if not os.path.exists(txtpath):
    os.makedirs(txtpath)

# 收集所有 XML 文件
files = [os.path.join(root, file) for root, _, files in os.walk(xmlpath) for file in files if file.endswith('.xml')]
number = len(files)
print(f"找到 {number} 个 XML 文件")

# 遍历并转换每个 XML 文件
for i, xml_file_path in enumerate(files):
    # 提取文件名并构建输出路径
    name = os.path.splitext(os.path.basename(xml_file_path))[0]
    txt_file_path = os.path.join(txtpath, name + '.txt')

    # 解析 XML 文件
    with open(xml_file_path, 'r') as xml_file:
        tree = ET.parse(xml_file)
        root = tree.getroot()
        w = int(root.find('size').find('width').text)  # 图像宽度
        h = int(root.find('size').find('height').text)  # 图像高度

    # 写入 TXT 文件
    with open(txt_file_path, 'w') as f_txt:
        content = ""
        first = True
        for obj in root.iter('object'):
            # 获取类别和边界框信息
            class_name = obj.find('name').text
            class_num = class_names.index(class_name)  # 类别的索引
            xmlbox = obj.find('bndbox')
            x1 = int(xmlbox.find('xmin').text)
            x2 = int(xmlbox.find('xmax').text)
            y1 = int(xmlbox.find('ymin').text)
            y2 = int(xmlbox.find('ymax').text)

            # 转换为 YOLO 格式
            x_center = (x1 + x2) / 2 / w
            y_center = (y1 + y2) / 2 / h
            width = (x2 - x1) / w
            height = (y2 - y1) / h

            # 构建 YOLO 格式的标注行
            line = f"{class_num} {x_center} {y_center} {width} {height}"
            content += line if first else f"\n{line}"
            first = False

        # 写入内容到 TXT 文件
        f_txt.write(content)
        print(f"已将 {name}.xml 转换为 {name}.txt")

print("转换完成！")