XML标签格式转换为YOLO TXT格式

针对的是多边形（<polygon>）来描述对象的边界，而不是传统的矩形框（<bndbox>）

import xml.etree.ElementTree as ET
import os
from pathlib import Path

# 解析VOC格式的XML文件，提取目标框的标签和坐标
def parse_voc_xml(xml_file):
    tree = ET.parse(xml_file)
    root = tree.getroot()

    # 获取图像的宽度和高度
    size = root.find('size')
    img_width = int(size.find('width').text)
    img_height = int(size.find('height').text)

    bboxes = []
    for obj in root.findall('object'):
        label = obj.find('name').text  # 获取目标的标签
        polygon = obj.find('polygon')  # 获取多边形的坐标

        # 获取多边形的所有顶点坐标
        x_coords = []
        y_coords = []
        for i in range(1, 5):  # 假设最多4个点（也可以扩展）
            x_coords.append(int(polygon.find(f'x{i}').text))
            y_coords.append(int(polygon.find(f'y{i}').text))

        # 计算边界框（最小矩形框）的坐标
        xmin = min(x_coords)
        xmax = max(x_coords)
        ymin = min(y_coords)
        ymax = max(y_coords)

        # 将 VOC 格式的坐标转换为 YOLO 格式
        x_center = (xmin + xmax) / 2
        y_center = (ymin + ymax) / 2
        width = xmax - xmin
        height = ymax - ymin

        # 归一化坐标
        x_center /= img_width
        y_center /= img_height
        width /= img_width
        height /= img_height

        bboxes.append((label, x_center, y_center, width, height))
    
    return bboxes, img_width, img_height

# 将标签保存为 YOLO 格式
def save_yolo_labels(image_filename, bboxes, output_dir):
    # YOLO 标签文件的路径
    txt_filename = os.path.join(output_dir, Path(image_filename).stem + '.txt')
    
    with open(txt_filename, 'w') as f:
        for label, x_center, y_center, width, height in bboxes:
            # YOLO 格式为：<class_id> <x_center> <y_center> <width> <height>
            # 假设标签的类别是数字（如果是字符串，则需要将类别转换为数字映射）
            f.write(f"{label} {x_center} {y_center} {width} {height}\n")

# 主函数：将 VOC XML 文件转换为 YOLO 格式
def convert_voc_to_yolo(voc_xml_dir, output_dir, class_mapping):
    # 如果输出目录不存在，则创建
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 遍历所有 XML 文件
    for xml_file in Path(voc_xml_dir).glob('*.xml'):
        print(f"Processing {xml_file}")
        # 解析 VOC XML 文件
        bboxes, img_width, img_height = parse_voc_xml(xml_file)
        
        # 将标签转化为 YOLO 格式，且需要使用数字类别而不是字符串
        bboxes = [(class_mapping[label], x_center, y_center, width, height) for label, x_center, y_center, width, height in bboxes]
        
        # 保存 YOLO 格式的标签
        save_yolo_labels(xml_file.stem + '.jpg', bboxes, output_dir)  # 假设图像文件与XML文件同名，后缀为 .jpg

# 主程序入口
if __name__ == "__main__":
    # 输入目录：包含VOC XML标注文件的目录
    voc_xml_dir = './original/labels'  # 替换为 VOC XML 文件目录路径
    # 输出目录：YOLO 格式标签文件的存储目录
    output_dir = './original/save'  # 替换为 YOLO 标签文件存储目录路径

    # 类别映射：将VOC中的标签名映射为数字类别
    class_mapping = {
        'car': 0,  
        'truck': 1,
        'bus': 2,
        'motorcycle': 3,
        'van': 4,
        'freight_car': 5,
    }

    # 转换 VOC 格式标签为 YOLO 格式
    convert_voc_to_yolo(voc_xml_dir, output_dir, class_mapping)
# 'car', 'truck', 'bus', 'motorcycle', 'van', 'freight_car'