xml转txt格式（训练YOLO v5）

xml转txt格式（训练YOLO v5）

import os
import xml.etree.ElementTree as ET

classes = ["D00", "D10", "D20", "D40"]  # 自己数据库中的类别，可以不包括xml文件中标注的全部类别

def corner2center(size, gt):
	# 读取xml文件中的xywh
	dw = 1./size[0]
	dh = 1./size[1]
	x = (gt[0] + gt[1]) / 2.0
	y = (gt[2] + gt[3]) / 2.0
	w = gt[1] - gt[0]
    h = gt[3] - gt[2]
    return x*dw, y*dh, w*dw, h*dh

def annotation2txt(image):
	# 读取xml文件，将xml文件放入当前路径下的Annotations文件夹中
	input_xml = open("Annotations/%s.xml" % image)  
	# 生成要转化成的txt文件，将其写入到txt文件夹下
	output_txt = open("txt/%s.txt" % image)
	tree = ET.parse(input_xml)
	root = tree.getroot()
	size = root.find("size")
	w = int(size.find("width").text)
    h = int(size.find("height").text)
	
	for obj in root.iter("object"):
        cls = obj.find("name").text
        if cls not in classes:
            continue
		cls_id = classes.index(cls)
		xmlbox = obj.find("bndbox")
        location = (float(xmlbox.find("xmin").text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
        location_0_1 = corner2center((w, h), location)
        output_txt.write(str(cls_id) + " " + " ".join([str(i) for i in location_0_1]) + "\n")

if __name__ == "__main__":
	xmls = os.listdir("Annotations")
    if not os.path.exists("label/"):
        os.makedirs("label/")
    for xml in xmls:
        xml = xml.split(".")[0]
        in_file = open(".\Annotations\%s.xml" % xml)
        i = in_file.readlines(0)
        if len(i) is 0:
            continue
        else:
            convert_annotation(xml)