yolov5训练安全帽数据集之 xml文件生成txt

    本文针对用yolov5训练安全帽数据集，由于初始的安全帽数据集数据较少，因此我找了三个数据集进行了融合，其中两个数据集的xml文件中分别以’0’，'1’代表没戴安全帽和戴安全帽，另一个数据集通过‘white’，‘red’，‘blue’等颜色来表示戴了安全帽以及安全帽的颜色，‘none’来表示没有带。
第一种xml格式如下图：

第二种如下：

除此之外，还有一些问题比如不同的数据及都是用数字来命名的，因此会重名，yolov5需要的信息是x，y，w，h，xml中提供的是x1，y1，x2，y2，那么我们如何解决呢 ？

第一步：对不同的数据集进行改名，将本代码和数据集文件夹放在同一文件下，并且我的数据集中图片和xml标注文件都是放在一起的，这样操作更方便点

import os

# 1. 获取一个要重命名的文件夹的名称：
folder_name = "anquanmao/"

# 2. 获取那个文件夹中所有的文件名字：
file_names = os.listdir(folder_name)

# 3. 循环用新名字代替旧名字（xml）
for name in file_names:
    if name.split('.')[-1] == 'xml':
        old_file_name = os.path.join(folder_name, name)
        #这里我选择在其中一个数据集图片名称前边加上‘new’，大家可以根据需要自己修改
        new_file_name = os.path.join(folder_name, 'new'+ name.split('.')[0].split('_')[0]+'.xml')
        os.rename(old_file_name, new_file_name)
# 4. 循环用新名字代替旧名字（jpg）        
    if name.split('.')[-1] == 'jpg':
        old_file_name = os.path.join(folder_name, name)
        new_file_name = os.path.join(folder_name, 'new'+ name.split('.')[0].split('_')[0]+'.xml')
        os.rename(old_file_name, new_file_name)

另一个数据集用同样的方式处理，然后我们就可以将三个数据集放在一起啦，注意是全部（包括图片和xml标注文件），然后生成训练需要的txt文件。

import glob
import os
xml_path = 'your_path'
#定义从xml获取信息的函数
def _read_anno(filename):

    import xml.etree.ElementTree as ET
    tree = ET.parse(filename)
    #获取宽w和高h
    a = tree.find('size')
    w,h = [int(a.find('width').text),
           int(a.find('height').text)]

    objects = []
    #这里是针对错误xml文件，图片的w和h都为0，这样的xml文件可以直接忽视，返回空列表
    if w == 0:
        return []
    #这里需要根据需要修改，因为我训练的目的是判断是否戴了头盔，因此从xml获取的name为none或者0的label都为0，其他的颜色或者1都为1    
    for obj in tree.findall('object'):
    	#获取name，我上边的实例图片中的红色区域
        name = obj.find('name').text
        #修改label，这里是不同数据集大融合的关键
        if name == 'none' or name == 0:
            label = 0
        else:
            label = 1
		#读取检测框的左上、右下角点的坐标
        bbox = obj.find('bndbox')
        x1, y1, x2, y2 = [int(bbox.find('xmin').text),
                          int(bbox.find('ymin').text),
                          int(bbox.find('xmax').text),
                          int(bbox.find('ymax').text)]
		#这里也很关键，yolov5需要中心点以及宽和高的标注信息，并且进行归一化，下边label后边的四个值即是归一化后保留4位有效数字的x，y，w，h
        obj_struct = [label,round((x1+x2)/(2.0*w),4), round((y1+y2)/(2.0*h),4), round((x2-x1)/(w),4),round((y2-y1)/(h),4)]

        objects.append(obj_struct)


    return objects
#接下来是写入txt文件中
if __name__ == '__main__':
    #定义一个空的字符串
    t = ''
    #获取所有的xml文件路径
    allfilepath = []
    for file in os.listdir(xml_path):
        if file.endswith('.xml'):
            file = os.path.join(xml_path,file)
            allfilepath.append(file)
        else:
            pass
     #生成需要的对应xml文件名的txt
    for file in allfilepath:
        txt_path = file.split('.')[0] + '.txt'
        result = _read_anno(file)
        #跳过空列表
        if len(result)==0:
            continue
        #写入信息，注意每次循环结束都把t重新定义，result是一个二维列表（行数为目标个数，列对应label和位置信息），为了避免读取出错（还有一个原因是我菜），我们一个一个的写入。
        with open(txt_path,'w') as f:
            for line in result:
                for a in line:
                    t = t+str(a)+' '
                f.writelines(t)
                f.writelines('\n')
                t =''

ok，txt文件都已经生成了，之后我会再写一个生成训练集train.txt和验证集val.txt的代码