将xml文件转yolov5训练数据txt标签文件

将xml文件转yolov5训练数据txt标签文件

提示：文章写完后，目录可以自动生成，如何生成可参考右边的帮助文档

---

前言

相信有很多同学跟我一样，在学习YOLOv5时，全网的寻找数据集资源，然后好不容易找到了一两个图片数据集，还自带标签文件，很开心，结果打开labels文件夹一看，全是.xml格式的文件，但是yolov5训练需要.txt格式文件，这个时候就需要写个脚本转化一下了，毕竟手动打标签——可真是太痛苦啦！！！（眼睛酸，手僵硬，很无聊）
本人网上了解了一下xml文件和txt文件的格式区别，发现一些规律，只需提取xml中的一些关键数据，通过计算，就能转化为txt文件需要的内容了，以下展开介绍。

---

一、准备数据集

这是我下载的数据集目录结构，xml文件在Annotation文件夹中，图片文件在JPEGImage文件夹中。
然后打开一个xml文件，观察其格式。
我们需要：
size——图片大小
object——对象，包括对象的name，和4个数值——表示标签框选的左上角和右下角的坐标数值

为了获取这些数据，可以使用正则表达式，使用python的re库，下面介绍代码部分。

二、转化代码部分

将xml文件转txt文件需要知道他们的内容格式。
.xml文件：

.txt文件内容格式：

每行内容表示为：类别序号 图框中心点x坐标 图框中心点y坐标 图框宽度 图框高度 （需要注意的是后面4个数值均为0-1之间的浮点数，经归一算法处理）
每一行表示一个标签图框，第一个数表示检测类别，后面4个表示图框位置与大小信息。

我们需要获取.xml中的图框信息，使用正则表达式，提前需要的内容，实现代码如下：

#定义函数，获取需要的内容，需要输入参数xml_path——.xml文件父级目录，name——.xml文件名
def convert_annotation(xml_path,name):
	#拼接得到文件路径
    xml_name=xml_path+"/"+name
    #打开文件
    with open(xml_name,"r",encoding="utf-8")as f1:
        text=f1.read().replace("\n","")			#读取文件内容，并且去除换行
        text=text.replace(" ","")				#去除空格，便于后期正则提取
    # print(text)
    #正则表达式，获取图片的宽度和高度
    img_size=re.findall("<width>([0-9]+)</width>.*?<height>([0-9]+)</height>",text)[0]
    print(img_size)
    #正则表达式，获取所有类别的名字，以及他的标签图框的两个坐标的4个值，分别是min的x、y,和max的x,y坐标值
    find_datas=re.findall("<object>.*?<name>([a-z|A-Z]*?)</name>.*?<xmin>([0-9]+?)</xmin>.*?<ymin>([0-9]+?)</ymin>.*?<xmax>([0-9]+?)</xmax>.*?<ymax>([0-9]+?)</ymax>",text)
    # print(find_datas)

获取的坐标值需要进行归一化处理，有利于后期图像训练时的计算。

#定义函数，接收两个数组，Imgsize——储存图片宽和高，box——储存4个坐标值
def convert(Imgsize,box):
    width,height=Imgsize
    xmin,ymin,xmax,ymax=box
    dw=max(xmin,xmax)-min(xmin,xmax)#框的宽
    dh=max(ymin,ymax)-min(ymin,ymax)#框的高
    x_=(xmin+xmax)/2.0#框的中心点x坐标
    y_=(ymin+ymax)/2.0#框的中心点y坐标

    #归一化，round(x,6)表示结果保留6位小数
    cx=round(x_/width,6)
    cy=round(y_/height,6)
    w=round(dw/width,6)
    h=round(dh/height,6)

    return [cx,cy,w,h]

完整代码如下：

import os
from os import listdir,getcwd
import re

classes=[]					#列表，储存标签类别
classes_dict={}				#字典，储存标签类别及对应的序号，yolo标签文件中用序号表示各个类别

def convert(Imgsize,box):
    width,height=Imgsize
    xmin,ymin,xmax,ymax=box
    dw=max(xmin,xmax)-min(xmin,xmax)#框的宽
    dh=max(ymin,ymax)-min(ymin,ymax)#框的高
    x_=(xmin+xmax)/2.0#框的中心点x坐标
    y_=(ymin+ymax)/2.0#框的中心点y坐标

    #归一化
    cx=round(x_/width,6)
    cy=round(y_/height,6)
    w=round(dw/width,6)
    h=round(dh/height,6)

    return [cx,cy,w,h]


def save_txt(namepath,text):			#将内容保存至文本文件中，即成功转化为yolo的标签文件
    with open(namepath,'w')as f:
        f.write(text)

def convert_annotation(xml_path,name):
    xml_name=xml_path+"/"+name
    with open(xml_name,"r",encoding="utf-8")as f1:
        text=f1.read().replace("\n","")
        text=text.replace(" ","")
    # print(text)
    img_size=re.findall("<width>([0-9]+)</width>.*?<height>([0-9]+)</height>",text)[0]
    print(img_size)
    find_datas=re.findall("<object>.*?<name>([a-z|A-Z]*?)</name>.*?<xmin>([0-9]+?)</xmin>.*?<ymin>([0-9]+?)</ymin>.*?<xmax>([0-9]+?)</xmax>.*?<ymax>([0-9]+?)</ymax>",text)
    # print(find_datas)
    savetext=""
    for item in find_datas:			#由于每张图片中可能存在多个类别或标签，将同张图片中所有标签内容保存在一个txt文件中
        class_=item[0]				#类别名
        if class_ not in classes:	#判断此类别是否已存在于classes列表中
            classes.append(class_)	#如果不存在，则添加至列表中
            classes_dict[class_]=len(classes)-1		#同时，为此类别标记序号，即第几个类别，从0开始标号

        print(item)
        imgsize=[int(img_size[0]),int(img_size[1])]			#正则获取的内容为文本格式，需要转化为数字
        box=[int(item[1]),int(item[2]),int(item[3]),int(item[4])]
        site=convert(imgsize,box)							#整理内容，归一化数值
        # print(class_,site)
        savetext+="{0} {1} {2} {3} {4}".format(classes_dict[class_],site[0],site[1],site[2],site[3])		#按格式拼接内容
        savetext+="\n"			#文本换行
    # print(savetext)
    name=name.split(".")[0]		#获取文件名，不包含文件后缀
    save_txt(labels_p+"/"+name+".txt",savetext.strip())			#保存标签文件
    

if __name__=="__main__":
    root_path=os.getcwd()			#获取当前目录绝对路径
    print(root_path)
    labels_p=root_path+"/"+"Labels"	#定义标签储存的文件夹路径
    try:
        os.makedirs(labels_p)		#创建文件夹
    except:pass
    xml_path=root_path+'/'+'Annotations'		#.xml文件的父级文件夹路径
    xml_list=sorted(listdir(xml_path))			#获取所有.xml文件名，并保存至列表中，同时排序
    print(xml_list[:20])						#打印前20个内容，查看结果，验证是否正确
    for name in xml_list:						#历遍列表
        print(name)
        convert_annotation(xml_path,name)		#开始转化
    
    print(classes,classes_dict)
    with open(labels_p+"/classes.txt","w")as f2:	#保存类别标签文件
        text=""
        for t in classes:
            text+=t
            text+="\n"
        print(text)
        f2.write(text)

运行结果：

总结

准备好数据集，在数据集目录下放置此脚本，打开编辑器，修改脚本中文件夹路径，这一行：

xml_path=root_path+'/'+'Annotations'		#.xml文件的父级文件夹路径

即可，运行脚本，转化完成。

---

初次编写文章，有许多不足，如果有什么问题或者哪里不懂的可评论区交流。
感谢大家的点赞！

---

上面代码可在这下载：transition.py

此外，我还编写了自动随机提取数据集的脚本：
get_dataset.py 用于将准备好的数据集提取生成YOLO使用的文件夹目录格式。如下：
dataset
—images
—.—train
—.—val
—.—test
—labels
—.—train
—.—val
—.—test
大家可自行下载。