import os
import xml.etree.ElementTree as ET


def convert_xml_to_txt(xml_file, txt_file):
    tree = ET.parse(xml_file)
    root = tree.getroot()

    # 获取图像的宽度和高度
    size = root.find('size')
    image_width = float(size.find('width').text)
    image_height = float(size.find('height').text)

    with open(txt_file, 'w') as f:
        for obj in root.findall('object'):
            name = obj.find('name').text
            if name == 'person':  # 只处理person类别
                bndbox = obj.find('bndbox')
                xmin = float(bndbox.find('xmin').text)
                ymin = float(bndbox.find('ymin').text)
                xmax = float(bndbox.find('xmax').text)
                ymax = float(bndbox.find('ymax').text)

                # 归一化坐标
                x_center = (xmin + xmax) / 2.0 / image_width
                y_center = (ymin + ymax) / 2.0 / image_height
                width = (xmax - xmin) / image_width
                height = (ymax - ymin) / image_height

                # 写入txt文件，格式为：class x_center y_center width height
                f.write(f"0 {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")


def process_folder(xml_folder, output_folder):
    # 确保输出文件夹存在
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # 遍历文件夹中的所有xml文件
    for xml_file in os.listdir(xml_folder):
        if xml_file.endswith('.xml'):
            xml_path = os.path.join(xml_folder, xml_file)
            txt_filename = os.path.splitext(xml_file)[0] + '.txt'
            txt_path = os.path.join(output_folder, txt_filename)
            convert_xml_to_txt(xml_path, txt_path)
            print(f"Processed: {xml_path} -> {txt_path}")


if __name__ == "__main__":
    xml_folder = r"C:\Users\DELL\Desktop\VOC2007\Annotations"  # 替换为你的XML文件夹路径
    output_folder = r"C:\Users\DELL\Desktop\VOC2007\txt"  # 替换为你的输出TXT文件夹路径

    process_folder(xml_folder, output_folder)