yolov5的xml转txt

import os
import xml.etree.ElementTree as ET

# 定义类别名称和ID映射关系
classes = {'1': 0}

# 定义函数将坐标转换为YOLOv5格式
def convert(size, box):
    dw = 1.0/size[0]
    dh = 1.0/size[1]
    x = (box[0] + box[1])/2.0 - 1
    y = (box[2] + box[3])/2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)

# 定义函数将XML文件转换为YOLOv5格式的TXT文件
def convert_annotation(xml_file, txt_file):
    tree = ET.parse(xml_file)
    root = tree.getroot()
    size = (int(root.find('size').find('width').text), int(root.find('size').find('height').text))
    with open(txt_file, 'w') as f:
        for obj in root.iter('object'):
            cls = '1'  # 只有一个类别
            cls_id = classes[cls]
            xmlbox = obj.find('bndbox')
            b = (int(xmlbox.find('xmin').text), int(xmlbox.find('xmax').text), int(xmlbox.find('ymin').text), int(xmlbox.find('ymax').text))
            bb = convert(size, b)
            f.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')

# 遍历XML文件夹并转换每个XML文件
xml_folder = 'D:/xmltxt/xml'
txt_folder = 'D:/xmltxt/txt'
for xml_file in os.listdir(xml_folder):
    if xml_file.endswith('.xml'):
        xml_path = os.path.join(xml_folder, xml_file)
        txt_path = os.path.join(txt_folder, xml_file.replace('.xml', '.txt'))
        convert_annotation(xml_path, txt_path)