yolo的数据集txt转xml格式

对于一些用于yolo系列算法的训练中，常会用到yolo格式进行标注，yolo格式多保存为txt格式，但是用于数据增强很多代码采用xml格式才能进行数据增强，所以这个代码可以有效提供txt转成xml格式，方便使用。 

import os
import xml.etree.ElementTree as ET


def convert_txt_to_xml(txt_folder_path, xml_output_folder):
    txt_files = os.listdir(txt_folder_path)
    for txt_file in txt_files:
        if not txt_file.endswith(".txt"):
            continue

        txt_path = os.path.join(txt_folder_path, txt_file)
        xml_output_path = os.path.join(xml_output_folder, os.path.splitext(txt_file)[0] + ".xml")

        with open(txt_path, 'r') as f:
            lines = f.readlines()

        root = ET.Element("annotation")

        # Parse image file name from txt path
        image_name = os.path.splitext(txt_file)[0] + ".jpg"

        # Add image filename to XML
        filename = ET.SubElement(root, "filename")
        filename.text = image_name

        # Add image size to XML
        size = ET.SubElement(root, "size")
        width = ET.SubElement(size, "width")
        height = ET.SubElement(size, "height")
        depth = ET.SubElement(size, "depth")
        width.text = "416"  # Replace with actual image width
        height.text = "416"  # Replace with actual image height
        depth.text = "3"  # Replace with actual image depth (number of channels)

        for line in lines:
            line = line.strip().split()
            if len(line) != 5:
                continue

            class_id = line[0]
            try:
                x_center = float(line[1])
                y_center = float(line[2])
                width = float(line[3])
                height = float(line[4])
            except ValueError:
                continue

            xmin = x_center - (width / 2)
            ymin = y_center - (height / 2)
            xmax = x_center + (width / 2)
            ymax = y_center + (height / 2)

            object_elem = ET.SubElement(root, "object")
            name = ET.SubElement(object_elem, "name")
            bndbox = ET.SubElement(object_elem, "bndbox")

            name.text = class_id

            xmin_elem = ET.SubElement(bndbox, "xmin")
            ymin_elem = ET.SubElement(bndbox, "ymin")
            xmax_elem = ET.SubElement(bndbox, "xmax")
            ymax_elem = ET.SubElement(bndbox, "ymax")

            xmin_elem.text = str(xmin)
            ymin_elem.text = str(ymin)
            xmax_elem.text = str(xmax)
            ymax_elem.text = str(ymax)

        tree = ET.ElementTree(root)
        tree.write(xml_output_path)


# Example usage
txt_folder_path = "path/to/txt_folder"
xml_output_folder = "path/to/txt_folder"

convert_txt_to_xml(txt_folder_path, xml_output_folder)