将YOLO数据集的txt标签转为VOC数据集的xml标签

import os
import xml.etree.ElementTree as ET
from xml.dom.minidom import parseString

# Define paths
images_path = "/mnt/data/fall_person/images/val/"  # 图片路径
labels_path = "/mnt/data/fall_person/labels/val/"  # TXT标注文件路径
output_path = "/mnt/data/fall_person/anchor/"  # 生成的XML文件存储路径

# 创建output目录（如果不存在）
if not os.path.exists(output_path):
    os.makedirs(output_path)

# 获取图片文件列表
image_files = [f for f in os.listdir(images_path) if f.endswith('.jpg') or f.endswith('.png')]

def create_xml(filename, img_width, img_height, objects):
    annotation = ET.Element("annotation")

    folder = ET.SubElement(annotation, "folder").text = "images"
    filename_tag = ET.SubElement(annotation, "filename").text = filename

    size = ET.SubElement(annotation, "size")
    ET.SubElement(size, "width").text = str(img_width)
    ET.SubElement(size, "height").text = str(img_height)
    ET.SubElement(size, "depth").text = "3"  # Assuming RGB images

    for obj in objects:
        obj_tag = ET.SubElement(annotation, "object")
        ET.SubElement(obj_tag, "name").text = obj['name']
        ET.SubElement(obj_tag, "pose").text = "Unspecified"
        ET.SubElement(obj_tag, "truncated").text = "0"
        ET.SubElement(obj_tag, "difficult").text = "0"

        bndbox = ET.SubElement(obj_tag, "bndbox")
        ET.SubElement(bndbox, "xmin").text = str(obj['xmin'])
        ET.SubElement(bndbox, "ymin").text = str(obj['ymin'])
        ET.SubElement(bndbox, "xmax").text = str(obj['xmax'])
        ET.SubElement(bndbox, "ymax").text = str(obj['ymax'])

    return annotation

def txt_to_xml(image_file):
    filename = os.path.splitext(image_file)[0]
    label_file = os.path.join(labels_path, filename + ".txt")
    image_file_path = os.path.join(images_path, image_file)
    
    if not os.path.exists(label_file):
      return
    # 读取图片的宽高（可以使用PIL或OpenCV读取图片尺寸）
    from PIL import Image
    img = Image.open(image_file_path)
    img_width, img_height = img.size

    objects = []
    with open(label_file, 'r') as file:
        for line in file.readlines():
            parts = line.strip().split()
            class_name = parts[0]
            # YOLO标注格式：class_id, center_x, center_y, width, height
            center_x = float(parts[1]) * img_width
            center_y = float(parts[2]) * img_height
            width = float(parts[3]) * img_width
            height = float(parts[4]) * img_height

            xmin = int(center_x - width / 2)
            ymin = int(center_y - height / 2)
            xmax = int(center_x + width / 2)
            ymax = int(center_y + height / 2)

            obj = {
                'name': class_name,
                'xmin': xmin,
                'ymin': ymin,
                'xmax': xmax,
                'ymax': ymax
            }
            objects.append(obj)

    annotation = create_xml(filename, img_width, img_height, objects)
    dom = parseString(ET.tostring(annotation))
    with open(os.path.join(output_path, filename + ".xml"), "w") as f:
        f.write(dom.toprettyxml(indent="\t"))

# 处理所有图片
for image_file in image_files:
    txt_to_xml(image_file)

print("转换完成！")