Yolov5数据集混合训练

导师非让训练一个应急场景balabala数据集
那就拼拼凑凑训练下“人+安全帽+火灾”三个分类的数据集吧

准备

首先准备了两份数据集，一份是一位大神的安全帽检测，另一份是火灾检测。
这两份的数据集都是开源的，基本是只给了weights和images的，但是没有给labels，所有都是要自己批处理检测后进行混合训练的。

步骤

1.批处理生成labels的.txt文件（helmet数据集）

python detect.py --save-txt --source 数据集的路径 --weights 标有你想要检测目标的权重文件位置

注意，权重文件必须包含你想要的检测对象，比如要留下person，那么用yolov5x就可以了；要识别安全帽，可以借用大神训练好的helmet数据集（链接https://pan.baidu.com/s/1UbFkGm4EppdAU660Vu7SdQ）

2.这是会是生成的一张图片对应一个txt文件，那么写个脚本取出其中的txt文件

import os

YOLOV5_LABEL_ROOT = r"D:\dataset\yolov5\inference\output\fire\\"  # yolov5 导出的推理图片的 txt
DATASET_LABEL_ROOT = r"D:\dataset\converg\helmet-dataset\VOC2028\txt\\"  # 数据集的路径

if __name__ == '__main__':
    yolo_file = os.listdir(YOLOV5_LABEL_ROOT)

    # 遍历文件里面有 .txt 结尾的
    for file_name in yolo_file:

        # 判断 txt 文件才进行读取
        if not file_name.endswith(".txt"):
            continue

        file_path = YOLOV5_LABEL_ROOT + file_name
        with open(file_path, "r+") as f:
            for line in f.readlines():
                # #只需要提取 0 -> 检测目标映射 的数据
                if line.split()[0] != '0':
                    continue
                line = line[1:]
                line = '3'+line
                data_path = DATASET_LABEL_ROOT + file_name
                print(data_path)
                # 汇总到数据集的标注文件
                with open(data_path, "a") as fd:
                    fd.write(line)

但是两个数据集融合的时候，标签也需要改变，比如要把如下的txt的第一列0换成3

再写个脚本进行批处理

import os

YOLOV5_LABEL_ROOT = r"D:\dataset\converg\fire_dataset\VOC2020\txt\\"  # yolov5 导出的推理图片的 txt
DATASET_LABEL_ROOT = r"D:\dataset\converg\helmet-dataset\VOC2028\fire_labels\\"  # 数据集的路径

if __name__ == '__main__':
    yolo_file = os.listdir(YOLOV5_LABEL_ROOT)

    # 遍历文件里面有 .txt 结尾的
    for file_name in yolo_file:

        # 判断 txt 文件才进行读取
        if not file_name.endswith(".txt"):
            continue

        file_path = YOLOV5_LABEL_ROOT + file_name
        with open(file_path, "r+") as f:
            for line in f.readlines():
                line = line[1:]
                line = '3'+line
                data_path = DATASET_LABEL_ROOT + file_name
                print(data_path)
                # 汇总到数据集的标注文件
                with open(data_path, "a") as fd:
                    fd.write(line)

第二步
把数据集的images和label变为train和val的树状结构

第三步
这里准备工作就正式结束了
修改下coco128.yaml和yolov5s.yaml的参数个数以及集合就可以了

data.yaml

yolov5s.yaml

第四步
放到服务器训练

CUDA_VISIBLE_DEVICES=0,1(可选的gpu) python train.py --img 640 --data data/data.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt --batch-size 16 --epochs 100

训练结果

预测一张：

另外找的两张图拼接的，没啥问题，ok

需要数据集或者交流的小伙伴欢迎私信