yolo数据集txt转xml格式转换

在网上找了一大堆打代码，发现都会报错，自己弄了一个，只需要改文件地址和标签名就行了。在设置路径那里把jpg文件夹和txt文件夹的地址更换一下，创建一个xml文件夹，然后根据自己的对应标签在标签映射那里改一下，运行代码就可以了。

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

# 设置路径
input_dir = 'path/to/your/yolo/txt/files'
output_dir = 'path/to/save/xml/files'
img_dir = 'path/to/your/jpg/images'

# 创建输出目录
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

# 标签映射
label_map = {
    0: 'label_name_0',
    1: 'label_name_1',
    2: 'label_name_2',
    # 添加更多标签
}

def create_xml(image_path, boxes, output_path):
    # 读取图片的宽和高
    img = cv2.imread(image_path)
    height, width, channels = img.shape

    annotation = ET.Element('annotation')
    ET.SubElement(annotation, 'folder').text = os.path.basename(os.path.dirname(image_path))
    ET.SubElement(annotation, 'filename').text = os.path.basename(image_path)
    ET.SubElement(annotation, 'path').text = image_path
    source = ET.SubElement(annotation, 'source')
    ET.SubElement(source, 'database').text = 'Unknown'
    
    size = ET.SubElement(annotation, 'size')
    ET.SubElement(size, 'width').text = str(width)
    ET.SubElement(size, 'height').text = str(height)
    ET.SubElement(size, 'depth').text = str(channels)
    ET.SubElement(annotation, 'segmented').text = '0'

    for box in boxes:
        obj = ET.SubElement(annotation, 'object')
        ET.SubElement(obj, 'name').text = box[0]
        ET.SubElement(obj, 'pose').text = 'Unspecified'
        ET.SubElement(obj, 'truncated').text = '0'
        ET.SubElement(obj, 'difficult').text = '0'
        bndbox = ET.SubElement(obj, 'bndbox')
        ET.SubElement(bndbox, 'xmin').text = str(box[1])
        ET.SubElement(bndbox, 'ymin').text = str(box[2])
        ET.SubElement(bndbox, 'xmax').text = str(box[3])
        ET.SubElement(bndbox, 'ymax').text = str(box[4])

    xml_str = ET.tostring(annotation, 'utf-8')
    pretty_xml_as_string = minidom.parseString(xml_str).toprettyxml(indent="  ")
    with open(output_path, 'w') as f:
        f.write(pretty_xml_as_string)

def yolo_to_xml(txt_file, img_dir, output_dir):
    with open(txt_file, 'r') as f:
        lines = f.readlines()
    
    boxes = []
    for line in lines:
        class_id, x_center, y_center, width, height = map(float, line.strip().split())
        class_name = label_map[int(class_id)]  # 使用标签映射
        img_path = os.path.join(img_dir, os.path.basename(txt_file).replace('.txt', '.jpg'))
        img = cv2.imread(img_path)
        img_height, img_width, _ = img.shape
        x_center *= img_width
        y_center *= img_height
        width *= img_width
        height *= img_height
        xmin = int(x_center - width / 2)
        ymin = int(y_center - height / 2)
        xmax = int(x_center + width / 2)
        ymax = int(y_center + height / 2)
        boxes.append((class_name, xmin, ymin, xmax, ymax))
    
    output_path = os.path.join(output_dir, os.path.basename(txt_file).replace('.txt', '.xml'))
    create_xml(img_path, boxes, output_path)

# 批量转换
for txt_file in os.listdir(input_dir):
    if txt_file.endswith('.txt'):
        yolo_to_xml(os.path.join(input_dir, txt_file), img_dir, output_dir)