本文针对用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的代码
