import os
import xml.etree.ElementTree as ET

//只需要改这三项
input_dir = r'H:/PYTHON_Proj/basketball/Annotations'
out_dir = r'H:/PYTHON_Proj/basketball/txt'
class_list = ["basketball"]


def file_name(input_dir):
    """列出输入目录下所有的XML文件名（不带扩展名）"""
    files = [os.path.splitext(f)[0] for f in os.listdir(input_dir) if f.endswith('.xml')]
    return files


def convert_coordinate(imgshape, bbox):
    """将边界框坐标从像素转换为归一化坐标"""
    xmin, ymin, xmax, ymax = bbox
    width, height = imgshape
    dw, dh = 1. / width, 1. / height
    x, y = (xmin + xmax) / 2.0, (ymin + ymax) / 2.0
    w, h = xmax - xmin, ymax - ymin
    x, y, w, h = x * dw, y * dh, w * dw, h * dh
    return x, y, w, h


def read_xml(file_name):
    """读取XML文件并转换为TXT格式"""
    xml_path = os.path.join(input_dir, f"{file_name}.xml")
    tree = ET.parse(xml_path)
    root = tree.getroot()

    size = root.find('size')
    width = int(size.find('width').text)
    height = int(size.find('height').text)
    imgshape = (width, height)

    txt_result = ''
    for obj in root.findall('object'):
        obj_name = obj.find('name').text
        if obj_name in class_list:
            obj_id = class_list.index(obj_name)
            bbox = obj.find('bndbox')
            xmin, ymin, xmax, ymax = map(float, [bbox.find(x).text for x in ['xmin', 'ymin', 'xmax', 'ymax']])
            x, y, w, h = convert_coordinate(imgshape, (xmin, ymin, xmax, ymax))
            txt = f"{obj_id} {x} {y} {w} {h}\n"
            txt_result += txt

    txt_path = os.path.join(out_dir, f"{file_name}.txt")
    with open(txt_path, 'w') as f:
        f.write(txt_result)


def main():
    filelist = file_name(input_dir)
    print(class_list)
    for n, i in enumerate(filelist, start=1):
        print(f"第{n}个xml成功转换\n")
        read_xml(i)

        # 写入类别文件
    with open(os.path.join(out_dir, "classes.txt"), 'w') as f:
        for cls in class_list:
            f.write(f"{cls}\n")


if __name__ == "__main__":
    main()