yolov5 pt转onnx_工程案例详解：YOLOv5在建筑工地中安全帽佩戴检测的应用

↑ 点击 蓝字  关注极市平台 来源丨 AI深度前沿视线 编辑丨极市平台

极市导读

 

本文结合YOLO v5的官方教程，分成了五个章节非常详细的讲解了YOLO v5在建筑工地中安全帽佩戴检测的应用。附有相关步骤的代码详解以及权重文件链接。>>极市双11现金福利，前往文末即可查看惊喜！

一、YOLO v5训练自己数据集教程 

• 1.1 创建自己的数据集配置文件 

• 1.2 创建每个图片对应的标签文件 

• 1.3 文件放置规范 

• 1.4 聚类得出先验框(可选) 

• 1.5 选择一个你需要的模型 

• 1.6 开始训练 

• 1.7 看训练之后的结果 

二、侦测 

三、检测危险区域内是否有人 

• 3.1 危险区域标注方式 

• 3.2 执行侦测 

• 3.3 效果：在危险区域里面的人体会被 红色框 选出来 

四、生成 ONNX 

五、增加数据集的分类

该项目是使用  YOLOv5 v2.x  来训练在智能工地安全领域中头盔目标检测的应用,先来一波演示！

指标

yolov5s 为基础训练，epoch = 50

分类	P	R	mAP0.5
总体	0.884	0.899	0.888
人体	0.846	0.893	0.877
头	0.889	0.883	0.871
安全帽	0.917	0.921	0.917

对应的权重文件： https://pan.baidu.com/share/init?surl=ELPhtW-Q4G8UqEr4YrV_5A ，提取码:   b981

yolov5m 为基础训练，epoch = 100

分类	P	R	mAP0.5
总体	0.886	0.915	0.901
人体	0.844	0.906	0.887
头	0.9	0.911	0.9
安全帽	0.913	0.929	0.916

对应的权重文件： https://pan.baidu.com/share/init?surl=0hlKrgpxVsw4d_vHnPHwEA ，提取码:   psst

yolov5l 为基础训练，epoch = 100

分类	P	R	mAP0.5
总体	0.892	0.919	0.906
人体	0.856	0.914	0.897
头	0.893	0.913	0.901
安全帽	0.927	0.929	0.919

对应的权重文件： https://pan.baidu.com/share/init?surl=iMZkRNXY1fowpQCcapFDqw ，提取码:  a66e

1 YOLOv5训练自己数据集教程

使用的数据集：Safety-Helmet-Wearing-Dataset ，感谢这位大神的开源数据集！ https://github.com/njvisionpower/Safety-Helmet-Wearing-Dataset

本文结合 YOLOv5官方教程 来写

环境准备

首先确保自己的环境：

Python >= 3.7Pytorch == 1.5.x

训练自己的数据

提示：

关于增加数据集分类的方法，请看【5. 增加数据集的分类】

---

  1.1 创建自己的数据集配置文件

因为我这里只是判断 【人没有带安全帽】、【人有带安全帽】、【人体】 3个类别 ，基于  data/coco128.yaml  文件，创建自己的数据集配置文件  custom_data.yaml

# 训练集和验证集的 labels 和 image 文件的位置train: ./score/images/trainval: ./score/images/val# number of classesnc: 3# class namesnames: ['person', 'head', 'helmet']

  1.2 创建每个图片对应的标签文件

使用标注工具类似于 Labelbox 、CVAT 、精灵标注助手 标注之后，需要生成每个图片对应的  .txt  文件，其规范如下：

• 每一行都是一个目标

• 类别序号是零索引开始的(从0开始)

• 每一行的坐标 class x_center y_center width height 格式

• 框坐标必须采用归一化的 xywh格式(从0到1)。如果您的框以像素为单位，则将x_center和width除以图像宽度，将y_center和height除以图像高度。代码如下：

import numpy as npdef convert(size, box):    """    将标注的 xml 文件生成的【左上角x,左上角y,右下角x，右下角y】标注转换为yolov5训练的坐标    :param size: 图片的尺寸：[w,h]    :param box: anchor box 的坐标 [左上角x,左上角y,右下角x,右下角y,]    :return: 转换后的 [x,y,w,h]    """    x1 = int(box[0])    y1 = int(box[1])    x2 = int(box[2])    y2 = int(box[3])    dw = np.float32(1. / int(size[0]))    dh = np.float32(1. / int(size[1]))    w = x2 - x1    h = y2 - y1    x = x1 + (w / 2)    y = y1 + (h / 2)    x = x * dw    w = w * dw    y = y * dh    h = h * dh    return [x, y, w, h]

生成的  .txt  文件放置的名字是图片的名字，放置在 label 文件夹中，例如：

./score/images/train/00001.jpg  # image./score/labels/train/00001.txt  # label

生成的 .txt 例子

1 0.1830000086920336 0.1396396430209279 0.13400000636465847 0.159159163013100621 0.5240000248886645 0.29129129834473133 0.0800000037997961 0.168168172240257261 0.6060000287834555 0.29579580295830965 0.08400000398978591 0.17717718146741391 0.6760000321082771 0.25375375989824533 0.10000000474974513 0.213213218376040460 0.39300001866649836 0.2552552614361048 0.17800000845454633 0.28228228911757470 0.7200000341981649 0.5570570705458522 0.25200001196935773 0.42942943982779980 0.7720000366680324 0.2567567629739642 0.1520000072196126 0.23123123683035374

  1.3 文件放置规范

文件树如下

  1.4 聚类得出先验框(Yolov5 内部已做适配，可选)

使用代码  ./data/gen_anchors/clauculate_anchors.py  ，修改数据集的路径

FILE_ROOT = r"xxx" # 根路径ANNOTATION_ROOT = r"xxx"  # 数据集标签文件夹路径ANNOTATION_PATH = FILE_ROOT + ANNOTATION_ROOT

跑完会生成一个文件  anchors.txt ，里面有得出的建议先验框：

Best Accuracy = 79.72%Best Anchors = [[14.74, 27.64], [23.48, 46.04], [28.88, 130.0], [39.33, 148.07], [52.62, 186.18], [62.33, 279.11], [85.19, 237.87], [88.0, 360.89], [145.33, 514.67]]

  1.5 选择一个您需要的模型

在文件夹  ./models  下选择一个你需要的模型然后复制一份出来，将文件开头的  nc =  修改为数据集的分类数，下面是借鉴  ./models/yolov5s.yaml 来修改的

# parametersnc: 3  # number of classes     <============ 修改这里为数据集的分类数depth_multiple: 0.33  # model depth multiplewidth_multiple: 0.50  # layer channel multiple# anchorsanchors: # <============ 根据 ./data/gen_anchors/anchors.txt 中的 Best Anchors 修改，需要取整(可选)  - [14,27, 23,46, 28,130]   - [39,148, 52,186, 62.,279]   - [85,237, 88,360, 145,514]# YOLOv5 backbonebackbone:  # [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 headhead:  [[-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)  ]

   

1.6 开始训练

这里选择了  yolov5s  模型进行训练，权重也是基于  yolov5s.pt  来训练

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

其中， yolov5s.pt  需要自行下载放在本工程的根目录即可，下载地址 官方权重

  1.7 看训练之后的结果

训练之后，权重会保存在  ./runs  文件夹里面的每个  exp  文件里面的  weights/best.py  ，里面还可以看到训练的效果

2 推断

侦测图片会保存在  ./inferenct/output/  文件夹下 运行命令：

python detect.py --source   0  # webcam                            file.jpg  # image                             file.mp4  # video                            path/  # directory                            path/*.jpg  # glob                            rtsp://170.93.143.139/rtplive/470011e600ef003a004ee33696235daa  # rtsp stream                            http://112.50.243.8/PLTV/88888888/224/3221225900/1.m3u8  # http stream

例如使用我的  s  权重检测图片，可以运行以下命令，侦测图片会保存在  ./inferenct/output/  文件夹下

python detect.py --source 图片路径 --weights ./weights/helmet_head_person_s.pt

3 检测危险区域内是否有人

   3.1 危险区域标注方式

我这里使用的是 精灵标注助手 标注，生成了对应图片的 json 文件

   3.2 执行侦测

侦测图片会保存在  ./inferenct/output/  文件夹下 运行命令：

python area_detect.py --source ./area_dangerous --weights ./weights/helmet_head_person_s.pt

   3.3 效果：在危险区域里面的人体会被 红色框 选出来

4 生成ONNX

    4.1 安装 onnx 库

pip install onnx

    4.2 执行生成

python ./models/export.py --weights ./weights/helmet_head_person_s.pt --img 640 --batch 1

onnx  和  torchscript  文件会生成在  ./weights  文件夹中

5 增加数据集的分类

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

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

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

修改  ./data/gen_data/merge_data.py  中的自己数据集标签所在的路径，执行这个python脚本，会进行  person  类型的合并 教程来源： https://github.com/PeterH0323/Smart_Construction

推荐阅读

• 【CV101录播】技术方案详解1：人流密度/河道污染/安全帽识别获奖算法

• 深入浅出的Yolov5核心基础知识完整讲解

• 详细记录u版YOLOv5目标检测ncnn实现

双11极市回馈送你现金红包，领取即可提现，快来！ 福利一：邀好友翻倍概率，瓜分现金666元 关注“极市平台”，回复关键词“双11”，参与抽奖即可瓜分666元现金红包。邀请好友参与抽奖助力，抽中概率翻倍！双11当天晚11点11分自动开奖~ 福利二：邀3位好友，即领现金红包

邀请3位好友参与抽奖助力，即可瓜分100元！数量有限，先到先得！

添加极市小助手微信(ID : cvmart2)，备注：姓名-学校/公司-研究方向-城市(如：小极-北大-目标检测-深圳)，即可申请加入极市目标检测/图像分割/工业检测/人脸/医学影像/3D/SLAM/自动驾驶/超分辨率/姿态估计/ReID/GAN/图像增强/OCR/视频理解等技术交流群：每月大咖直播分享、真实项目需求对接、求职内推、算法竞赛、干货资讯汇总、与 10000+来自港科大、北大、清华、中科院、CMU、腾讯、百度等名校名企视觉开发者互动交流~ △长按添加极市小助手 △长按关注极市平台，获取 最新CV干货 觉得有用麻烦给个在看啦~