python-将yolo格式的标签转换成xml格式标签

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

def yolo_to_voc_bbox(x, y, w, h, image_w, image_h):
    """
    将Yolo格式的边界框转换为VOC格式的边界框。
    """
    center_x = x * image_w
    center_y = y * image_h
    box_w = w * image_w
    box_h = h * image_h
    xmin = int(center_x - (box_w / 2))
    ymin = int(center_y - (box_h / 2))
    xmax = int(center_x + (box_w / 2))
    ymax = int(center_y + (box_h / 2))
    return xmin, ymin, xmax, ymax

def convert_yolo_to_xml(yolo_anns, image_filename, image_width, image_height, output_folder):
    """
    将Yolo格式的注释转换为VOC XML格式。
    yolo_anns: Yolo格式的注释列表，每个元素为[class_id, x_center, y_center, width, height]
    image_filename: 图像文件名
    image_width: 图像宽度
    image_height: 图像高度
    output_folder: 输出XML文件的目录
    """
    annotation = ET.Element('annotation')
    ET.SubElement(annotation, 'folder').text = 'images'
    ET.SubElement(annotation, 'filename').text = os.path.basename(image_filename)
    size = ET.SubElement(annotation, 'size')
    ET.SubElement(size, 'width').text = str(image_width)
    ET.SubElement(size, 'height').text = str(image_height)
    ET.SubElement(size, 'depth').text = '3'

    for ann in yolo_anns:
        obj = ET.SubElement(annotation, 'object')
		# ET.SubElement(obj, 'name').text = str([ann[0])  # 假设类别ID直接作为类别名称
        ET.SubElement(obj, 'name').text = str(CLASS_DICT[ann[0]])  # 假设类别ID直接作为类别名称
        bndbox = ET.SubElement(obj, 'bndbox')
        xmin, ymin, xmax, ymax = yolo_to_voc_bbox(*ann[1:], image_width, image_height)
        ET.SubElement(bndbox, 'xmin').text = str(xmin)
        ET.SubElement(bndbox, 'ymin').text = str(ymin)
        ET.SubElement(bndbox, 'xmax').text = str(xmax)
        ET.SubElement(bndbox, 'ymax').text = str(ymax)

    xml_str = ET.tostring(annotation, encoding='utf8').decode('utf8')
    xml_pretty_str = minidom.parseString(xml_str).toprettyxml(indent="  ")
    
    output_filename = os.path.splitext(os.path.basename(image_filename))[0] + '.xml'
    output_path = os.path.join(output_folder, output_filename)
    with open(output_path, 'w') as f:
        f.write(xml_pretty_str)
    print(f"XML saved to {output_path}")


def get_yolo_annotations(yolo_file):
    """读取yolo标注格式的txt文件"""
    with open(yolo_file, 'r', encoding='utf-8') as f:
        data = f.readlines()
    yolo_annotations = []
    for i in data:
        line = [x for x in i.strip().split(' ')]
        yolo_annotations.append([int(line[0]), float(line[1]),float(line[2]),float(line[3]),float(line[4])])
    return(yolo_annotations)

###单张
# image_file = './dataset/images/1.jpg'	#图像路径
# img = cv2.imread(image_file)
# image_width, image_height = img.shape[1], img.shape[0]
# yolo_file = "./dataset/yololabel/1.txt"	#yolo标签路径
# yolo_annotations = get_yolo_annotations(yolo_file)
# # # 示例使用
# # yolo_annotations = [
# #     [0, 0.3, 0.4, 0.1, 0.2],  # 类别ID, 中心X, 中心Y, 宽比例, 高比例
# #     [1, 0.6, 0.1, 0.15, 0.1]
# # ]
# output_folder = '/workspace/data/BiaoJi/biandianzhan_shuxian/digital_yolov5_dataset/Annotations'
# os.makedirs(output_folder, exist_ok=True)
# convert_yolo_to_xml(yolo_annotations, image_file, image_width, image_height, output_folder)

# 自定义的标签字典
CLASS_DICT = {0:'LEDD_0', 1:'LEDD_1', 2:'LEDD_2', 3:'LEDD_3', 4:'LEDD_4', 5:'LEDD_5', \
              6:'LEDD_6', 7:'LEDD_7', 8:'LEDD_8', 9:'LEDD_9', 10:'LEDD_-'} 

###批量
imgpath = '/workspace/data/BiaoJi/biandianzhan_shuxian/digital_yolov5_dataset/images'
yolotxtpath = "/workspace/data/BiaoJi/biandianzhan_shuxian/digital_yolov5_dataset/yololabel"
output_folder = '/workspace/data/BiaoJi/biandianzhan_shuxian/digital_yolov5_dataset/Annotations'
os.makedirs(output_folder, exist_ok=True)

for imgname in os.listdir(imgpath):
    if '.jpg' not in imgname:
        continue
    img = cv2.imread(os.path.join(imgpath, imgname))
    image_width, image_height = img.shape[1], img.shape[0]

    yolo_file = os.path.join(yolotxtpath, imgname.replace('.jpg', '.txt'))
    yolo_annotations = get_yolo_annotations(yolo_file)

    convert_yolo_to_xml(yolo_annotations, imgname, image_width, image_height, output_folder)