登录 订阅

用Yolov5,实现训练安全帽+反光衣模型

认真猫咪

• 7 min read
用Yolov5,实现训练安全帽+反光衣模型

1、下载yolov5

官方github 官方地址下载，或者git下来
ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite (github.com)

git clone https://github.com/ultralytics/yolov5  # clone
cd yolov5
pip install -r requirements.txt  # install

确保环境符合要求

    Python>=3.7
    Pytorch==1.5.x
    PyQt5==5.15.3
    PyQtChart==5.15.3
    PyQt5-tools
    GPUtil

2、安装运行组件

命令行进入所在文件夹，使用PIP一次性加载所有的需求组件。

pip install -r requirements.txt

3、训练的教程

官方教程连接

  1. 训练教程所需的素材
    1. 数据说明文件《xxxx.yaml》文件
    2. 两个文件夹：images、labels
      ![[Pasted image 20240312221356.png]]
images 文件夹下面三个文件夹。labels也是一样3个文件夹。每个文件夹的文件是一样的，只是后缀不同，分别对应每张图片的标注信息。
如 :
images/train/0001.jpg
labels/train/0001.txt
![[Pasted image 20240312221459.png]]
  1. txt的数据结构（参照官方教程。
    一行5个数据。分别为
    与上面YAML对应。
    目标类型编号（0、1）| 横坐标 | 纵坐标 | 高度 | 宽度
  2. 数据转换。有时有些博主给的数据标注格式是。Annotations。这个也是类似的。无法直接满足yolov5的训练要求。需要进行转换。可以用博主的py方法。gen_head_helmet.py,来将 VOC 的数据集转换成 YOLOv5 训练需要用到的格式。
  3. 自己做数据，可以使用使用标注工具类似于 Labelbox 、CVAT 、精灵标注助手 标注之后，需要生成每个图片对应的 .txt 文件，其规范如下：
  • 每一行都是一个目标
  • 类别序号是零索引开始的（从0开始）
  • 每一行的坐标 class x_center y_center width height 格式
  • 框坐标必须采用归一化的 xywh格式（从0到1）。如果您的框以像素为单位，则将x_center和width除以图像宽度，将y_center和height除以图像高度。代码如下：
    文件的放置规范
    ![[Pasted image 20240312222733.png]]
  1. 数据清洗：由于网上很多博主，都只有安全帽训练后的模型、反光衣训练后的模型、没有，同时识别反光衣和安全帽的模型。这时需要将数据进行合并。
  2. 把反光衣、安全帽的数据变成上面的，分别放置到对应文件夹。这样的。
# 我的情况是，人体、安全帽、头像，已经有一套官方的模型训练数据了。我想把反光衣的加进去。
反光衣的模型，包含两个检测目标。分别为。反光衣、其他反光衣。
所以我做了以下工作，是他们合并起来
yaml我定义为以下内容：
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]

train: ./data\VOC2028\Safety_Helmet_Train_dataset/score/images/train

val: ./data\VOC2028\Safety_Helmet_Train_dataset/score/images/val

  

# number of classes

nc: 5

  

# class names

names: ['person', 'head', 'helmet','reflective_clothes','other_clothes']


1、先将反光衣的txt，按照YAML。反光衣分别顺序为3、4
2、原始的txt标注文件是，对应是0和1.所有写一个方法，把他们全部换成3和4.详情看merge_data.py
with open(file_path, "r") as f:

            for line in f.readlines():

  

                # 只需要提取 0 -> person 的数据

                if line.split()[0] != '0':

                    continue

  

                data_path = DATASET_LABEL_ROOT + file_name

                print(data_path)

                # 汇总到数据集的标注文件

                with open(data_path, "a") as fd:

                    fd.write(line)
3、将对应jpg图片，丢到/images/train。对应的txt,丢到/labels/train。
4、只有训练数据是不行的，val文件夹，对应的是校正数据，要将正确的对应的，丢进去。 又写了个方法，找出所有的图片只有争取的情况下的图片标注。findval.py 。
5、找到了校验的txt，再把校验的图片也找出来。根据txt文件名称，复制对应的图片，到目标文件夹。详情见findval-img.py。
  1. 数据合并：把所有的数据合并起来，就得到了，反光衣、安全帽，一起检测的模型数据，放置指定文件目录后，开始训练。在文件夹 ./models 下选择一个你需要的模型然后复制一份出来，将文件开头的 nc =  修改为数据集的分类数，下面是借鉴 ./models/yolov5s.yaml来修改的。这里参照博主的不用修改太多内容。
    第一行改一下。
# parameters
nc: 3  # number of classes     <============ 修改这里为数据集的分类数
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple

# anchors
anchors:
  - [10,13, 16,30, 33,23]  # P3/8
  - [30,61, 62,45, 59,119]  # P4/16
  - [116,90, 156,198, 373,326]  # P5/32

# YOLOv5 backbone
backbone:
  # [from, number, module, args]
  [[-1, 1, Focus, [64, 3]],  # 0-P1/2
   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
   [-1, 3, BottleneckCSP, [128]],
   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
   [-1, 9, BottleneckCSP, [256]],
   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
   [-1, 9, BottleneckCSP, [512]],
   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
   [-1, 1, SPP, [1024, [5, 9, 13]]],
   [-1, 3, BottleneckCSP, [1024, False]],  # 9
  ]

# YOLOv5 head
head:
  [[-1, 1, Conv, [512, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 6], 1, Concat, [1]],  # cat backbone P4
   [-1, 3, BottleneckCSP, [512, False]],  # 13

   [-1, 1, Conv, [256, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 4], 1, Concat, [1]],  # cat backbone P3
   [-1, 3, BottleneckCSP, [256, False]],  # 17

   [-1, 1, Conv, [256, 3, 2]],
   [[-1, 14], 1, Concat, [1]],  # cat head P4
   [-1, 3, BottleneckCSP, [512, False]],  # 20

   [-1, 1, Conv, [512, 3, 2]],
   [[-1, 10], 1, Concat, [1]],  # cat head P5
   [-1, 3, BottleneckCSP, [1024, False]],  # 23

   [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  ]

4、完成训练

修改对应的文件所在未知就和了。

----我的命令---

python train.py --img 640  --batch 16 --epochs 10 --data data\VOC2028\Safety_Helmet_Train_dataset/custom_data.yaml  --cfg ./models/custom_yolov5.yaml --weights yolov5s.pt

在runs文件夹train下就能找到你的模型训练结果啦。使用best.PT文件就是刚刚训练好的结果啦。

给数据加上人员标注信息。反向标注。
关于增加数据集分类的方法：

SHWD 数据集里面的 person 指的是头（head），没有 人体 的类别，先将现有的自己的数据集执行脚本生成 yolov5 需要的标签文件 .txt，之后再用 yolov5x.pt 加上 yolov5x.yaml ，使用指令检测出人体

python detect.py --save-txt --source 自己数据集的文件目录 --weights ./weights/yolov5x.pt

-----我的----

python detect.py --save-txt --source D:\ZSH\Development\AI\Vison\yolov5-master\yolov5-master\data\VOC2028\Safety_Helmet_Train_dataset\score\images\train --weights ./weights/yolov5x.pt

yolov5 会推理出所有的分类，并在 inference/output 中生成对应图片的 .txt 标签文件；

修改 ./data/gen_data/merge_data.py 中的自己数据集标签所在的路径，执行这个python脚本，会进行 人体(person) 类型的合并

tips:为加快训练过程，请确保你的torch版本是对应你的显卡CUDA要求。
![[Pasted image 20240313001511.png]]

Sign up for more like this.

Enter your email
Subscribe