import os
import xml.etree.ElementTree as ET
from xml.dom import minidom
from PIL import Image

def convert_txt_to_xml(txt_path, image_path, xml_path):
    # 打开图像，获取尺寸及通道数
    with Image.open(image_path) as img:
        img_width, img_height = img.size
        mode = img.mode
        if mode == "RGB":
            depth = 3
        elif mode == "L":
            depth = 1
        else:
            depth = 3  # 默认3

    # 读取TXT文件中的标注信息
    with open(txt_path, "r") as f:
        lines = f.readlines()

    # 类别编号到标签名称的映射
    id_to_label = {
        0: "light",
        1: "red",
        2: "yellow",
        3: "green"
    }

    # 构建XML结构
    annotation = ET.Element("annotation", verified="yes")

    folder_elem = ET.SubElement(annotation, "folder")
    folder_elem.text = os.path.basename(os.path.dirname(image_path))

    filename_elem = ET.SubElement(annotation, "filename")
    image_filename = os.path.basename(image_path)
    filename_elem.text = image_filename

    path_elem = ET.SubElement(annotation, "path")
    path_elem.text = os.path.abspath(image_path)

    source_elem = ET.SubElement(annotation, "source")
    database_elem = ET.SubElement(source_elem, "database")
    database_elem.text = "Unknown"

    size_elem = ET.SubElement(annotation, "size")
    width_elem = ET.SubElement(size_elem, "width")
    width_elem.text = str(img_width)
    height_elem = ET.SubElement(size_elem, "height")
    height_elem.text = str(img_height)
    depth_elem = ET.SubElement(size_elem, "depth")
    depth_elem.text = str(depth)

    segmented_elem = ET.SubElement(annotation, "segmented")
    segmented_elem.text = "0"

    # 逐行处理TXT标注
    for line in lines:
        parts = line.strip().split()
        if len(parts) != 5:
            print(f"跳过格式不正确的行: {line.strip()}")
            continue
        class_id_str, x_center_norm_str, y_center_norm_str, width_norm_str, height_norm_str = parts
        try:
            class_id = int(class_id_str)
            x_center_norm = float(x_center_norm_str)
            y_center_norm = float(y_center_norm_str)
            width_norm = float(width_norm_str)
            height_norm = float(height_norm_str)
        except Exception as e:
            print(f"转换数据出错 {line.strip()}: {e}")
            continue

        if class_id not in id_to_label:
            print(f"未知标签编号 {class_id}，跳过该行")
            continue

        label = id_to_label[class_id]

        # 计算绝对坐标
        x_center = x_center_norm * img_width
        y_center = y_center_norm * img_height
        bbox_width = width_norm * img_width
        bbox_height = height_norm * img_height

        xmin = int(x_center - bbox_width / 2)
        ymin = int(y_center - bbox_height / 2)
        xmax = int(x_center + bbox_width / 2)
        ymax = int(y_center + bbox_height / 2)

        # 构建object节点
        object_elem = ET.SubElement(annotation, "object")
        name_elem = ET.SubElement(object_elem, "name")
        name_elem.text = label
        pose_elem = ET.SubElement(object_elem, "pose")
        pose_elem.text = "Unspecified"
        truncated_elem = ET.SubElement(object_elem, "truncated")
        truncated_elem.text = "0"
        difficult_elem = ET.SubElement(object_elem, "difficult")
        difficult_elem.text = "0"

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

    # 美化输出的XML字符串
    xml_str = ET.tostring(annotation, encoding="utf-8")
    parsed_str = minidom.parseString(xml_str)
    pretty_xml_str = parsed_str.toprettyxml(indent="    ")

    with open(xml_path, "w", encoding="utf-8") as f:
        f.write(pretty_xml_str)

def batch_convert_txt_to_xml(txt_dir, image_dir, output_dir, image_ext=".jpg"):
    """
    批量将TXT转换为XML：
    txt_dir: TXT文件所在目录
    image_dir: 对应图像所在目录（图片文件名与TXT同名，后缀可指定）
    output_dir: XML文件保存目录
    image_ext: 图像文件扩展名（默认 .jpg）
    """
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    for file in os.listdir(txt_dir):
        if file.lower().endswith(".txt"):
            txt_path = os.path.join(txt_dir, file)
            base_name = os.path.splitext(file)[0]
            image_filename = base_name + image_ext
            image_path = os.path.join(image_dir, image_filename)
            if not os.path.exists(image_path):
                print(f"图像文件不存在: {image_path}，跳过 {txt_path}")
                continue
            xml_filename = base_name + ".xml"
            xml_path = os.path.join(output_dir, xml_filename)
            try:
                convert_txt_to_xml(txt_path, image_path, xml_path)
                print(f"成功转换: {txt_path} --> {xml_path}")
            except Exception as e:
                print(f"处理 {txt_path} 时出错: {e}")

if __name__ == "__main__":
    # 请修改以下目录为实际路径
    txt_directory = "path_to_txt_directory"          # TXT文件所在目录
    image_directory = "path_to_image_directory"        # 图像文件所在目录
    output_directory = "path_to_output_xml_directory"  # 输出XML文件保存目录
    batch_convert_txt_to_xml(txt_directory, image_directory, output_directory, image_ext=".jpg")