voc数据集转yolov5数据格式

from xml.etree import ElementTree as ET
import os

xml_path = r"F:\traffic_signs_yolov5\Annotations"
label_path = r"F:\traffic_signs_yolov5\labels\train"
image_path = r"F:\traffic_signs_yolov5\images\train"

class_name = []

for xml_file in os.listdir(xml_path):
    xml2txt_name = os.path.splitext(xml_file)[0] + '.txt'

    tree = ET.parse(os.path.join(xml_path,xml_file))
    root = tree.getroot()
    sign_object = root.findall('object')
    img_size = root.find("size")
    img_width = int(img_size.find("width").text)
    img_height = int(img_size.find("height").text)

    if sign_object:
        for sign in sign_object:
            name = sign.find('name')
            sign_name = name.text
            if sign_name not in class_name:
                class_name.append(sign_name)

            class_id = class_name.index(sign_name)

            bndbox = sign.find('bndbox')
            xmin = float(bndbox[0].text)
            ymin = float(bndbox[1].text)
            xmax = float(bndbox[2].text)
            ymax = float(bndbox[3].text)

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

            with open(os.path.join(label_path,xml2txt_name), 'a') as f:
                f.write("{} {} {} {} {}\n".format(class_id, x_center, y_center, width, height))

    else:
        #with open(os.path.join(label_path,xml2txt_name),'a') as f:
        #    f.write("")
        remove_img_name = os.path.splitext(xml_file)[0] + '.jpg'
        if os.path.exists(os.path.join(image_path,remove_img_name)):
            os.remove(os.path.join(image_path,remove_img_name))
            print("remove {}".format(remove_img_name))


print("class_name:",class_name)