0


Yolov5口罩佩戴实时检测项目(模型剪枝+opencv+python推理)

目录

0. 前言

如果只是想体验项目,请直接跳转到本文第2节,或者跳转到我的facemask_detect。

剪枝的代码可以查看我的github:yolov5-6.2-pruning
第1章是讲述如何得到第2章用到的onnx格式的模型文件(我的项目里直接提供了这个文件)。
第2章开始讲述如何使用cv2.dnn加载onnx文件并推理yolov5n模型。

1. 训练

本节内容:

下载mask_yolo数据集

1.1 获取口罩佩戴检测数据集

原数据集在FaceMaskDetection中有获取的方法。但是其是voc格式的数据集,用在yolov5中需要转成yolo格式的;而且有多处标注存在内容遗漏。

建议直接下载我生成的yolo格式的数据集,

获取链接

:mask_yolo。

1.2 训练环境配置

yolov5项目对pytorch的版本较为敏感,可能有些问题在某个特定版本会出现,有些问题则不会。

举例来说,pytorch 1.12版本调用yolov5中export.py文件生成onnx文件或者detect.py调用onnx文件,会出现问题,详见issue.

pytorch需要用

1.11

及更低的版本。

其余库参考yolov5的

requirements.txt

即可。

1.3 修改模型文件和数据集文件

本节内容:

clone yolov5项目

下载yaml文件并修改

1.3.1 使用的模型

本项目使用的是yolov5系列中最小的一个模型

yolov5n

,因为使用最常用的yolov5s会使得推理变慢,造成延时现象。并且只是检测是否佩戴口罩,可能并不需要太大的模型。

请前往github中clone最新的yolov5(当前为6.2版本)代码,地址:yolov5。

关于yolov5网络结构图可以参考我的另一篇博文:yolov5s 6.0结构图。

1.3.2 下载yaml文件并修改

  1. 下载数据集的yaml文件:mask.yaml,放入yolov5/data文件夹,并修改path路径为你下载的数据集存放的位置。在这里插入图片描述
  2. 下载模型的yaml文件:yolov5n_mask.yaml,放入yolov5/models文件夹。 关键参数nc=2对应:不佩戴口罩(face)和佩戴口罩(face_mask)。

1.4 训练

本节内容:

修改训练参数并训练

1.4.1 修改训练参数

以下三个参数分别对应预训练权重文件、模型文件、数据集文件,填上对应的路径即可。

默认

yolov5n.pt

程序会自动下载,如果下载失败,请使用该链接手动下载:yolov5n.pt, 然后放置在yolov5主目录下。

parser.add_argument('--weights',type=str, default=ROOT /'yolov5n.pt',help='initial weights path')# 模型yaml
parser.add_argument('--cfg',type=str, default=ROOT /'models/yolov5n_mask.yaml',help='model.yaml path')# 修改数据集yaml
parser.add_argument('--data',type=str, default=ROOT /'data/mask.yaml',help='dataset.yaml path')

修改

imgsz

参数,这个参数表示训练时图片缩放到什么尺寸,默认为640,建议改小一些,以加快训练速度,本项目这里设置为

320

.

parser.add_argument('--imgsz','--img','--img-size',type=int, default=320,help='train, val image size (pixels)')

windows端则需要修改

workers

参数设置为0,否则加载数据集可能会报错。

parser.add_argument('--workers',type=int, default=0,help='max dataloader workers (per RANK in DDP mode)')

epochs可以调小一些,比如100或者200.

parser.add_argument('--epochs',type=int, default=200,help='total training epochs')

其余参数可以参考我的train.py.

修改完毕后,运行train.py文件即可。

1.4.2 训练结果

训练结果可以在

runs/train

中查看,

exp

后面的数字最大的,就是当前训练对应的文件夹。

查看一下验证的结果,貌似还不错,用衣服遮住口鼻的也没有被误检为face_mask。
在这里插入图片描述

训练文件夹中的

weights

文件夹存放着权重文件,我们需要

best.pt

1.5 转换为onnx格式

使用yolov5中的

export.py

将pytorch模型转换为

onnx

格式(再次提示不要使用pytorch 1.12版本,使用1.11版本及更低版本)。

该py文件中用到了onnx,请使用

pip install onnx

安装。

使用如下命令转换,其中

runs/train/exp3/weights/best.pt

修改为你自己的训练权重路径。

python export.py --weights runs/train/exp3/weights/best.pt --include onnx

然后我们就在

best.pt

所在的文件夹得到了

best.onnx

文件

export会把三个检测头的输出合并到一起(以输入为

640×640

为例)

(80,80,255)->(3,80,80,85)->(19200,85)(40,40,255)->(3,40,40,85)->(4800,85)(20,20,255)->(3,20,20,85)->(1200,85)(19200+4800+1200,85)=(25200,85)

因此export后的模型detect只会输出一个(25200,85)的矩阵。

在这里插入图片描述

2. 使用口罩佩戴实时检测项目

本节内容:

项目介绍和使用


我的项目地址:facemask_detect

从github中clone下来,并根据对应的requirements.txt安装项目的虚拟环境。

该项目中的模型使用opencv-python进行onnx格式的模型推理,摆脱了对pytorch库的依赖,同时使用onnx格式会比pytorch格式进行推理速度要更加快。

项目的具体使用步骤参照项目中的README.md。

项目界面如下所示:
在这里插入图片描述

项目中的yolov5n_mask.onnx就是在1.4小节得到的

best.onnx

文件。

2.1 cv2.dnn推理yolov5n

如果只是想体验项目,本文到上面就结束了。

下面讲解如何使用cv2.dnn模块使用得到的onnx格式的模型文件。

这部分代码在项目中的

controller/utils/yolo_inference.py

中,

该部分可以大致分为

  1. 加载模型:cv2.dnn.readNetFromONNX
  2. 数据预处理:包括letterbox、BGR2RGB、归一化等
  3. 推理:net.forwad()
  4. 后处理操作:NMS、scale_boxes
  5. 绘制标注信息:cv2.rectangle

下面讲解其中的一些内容。

2.1.1 读取模型

opencv-python读取onnx格式的模型,需要使用

cv2.dnn.readNetFromONNX

函数。

net = cv2.dnn.readNetFromONNX('models/yolov5s_mask.onnx')

第三次提示之前pytorch的版本不要使用1.12版本,因为使用这个版本得到的onnx模型,上边的代码加载时会报错,请换成1.11或者更低的版本。

2.1.2 letterbox函数

yolov5在将图像数据输入到模型之前需要进行

缩放操作

,同时cv2.dnn中进行forward时输入的尺寸也要和训练时用到的尺寸保持一致,因为上边训练用到的是320×320的,所以模型的输入也得是320×320的,否则会报错:

cv2.error: OpenCV(4.5.5) D:\a\opencv-python\opencv-python\opencv\modules\dnn\src\layers\reshape_layer.cpp:107: error:(-215:Assertion failed) total(srcShape, srcRange.start, srcRange.end)== maskTotal in function 'cv::dnn::computeShapeByReshapeMask'

然而直接进行缩小,可能会导致图标中的目标变形,影响目标检测效果。
所以yolov5采用如下图的缩放技巧,以

(750,500)

缩放到

(640,640)

为例。

  1. 计算高度和宽度缩放比
  2. 使用最小的缩放比,保持缩放后高度和宽度仍是等比的
  3. 对于不足640的边,以灰色填充两边使得其达到640

通过这样的技巧既能够保持高宽比,又能达到正方形缩放的效果。
在这里插入图片描述

2.1.3 推理

img_trans = cv2.dnn.blobFromImage(image=img_data, scalefactor=1/255., swapRB=True)
net = cv2.dnn.readNetFromONNX('models/yolov5n_mask.onnx')
net.setInput(img_trans)# 输入数据,并获得输出,yolov5中将三个检测头的输出合并到了一起
outputs = net.forward()

使用opencv读取的图片数据需要进行预处理操作:

img_trans = cv2.dnn.blobFromImage(image=img_data, scalefactor=1/255., swapRB=True)
  • 归一化,将数据从0-255缩放至0-1:scalefactor=1 / 255.
  • 将RGB转换为BGR:swapRB=True

在推理之前需要设置输入数据:

net.setInput(img_trans)

然后直接使用

net.forward()

得到最后一层(即

Detect

层)的输出。

2.1.3 NMS

2.1.4 待续

后续还有些内容需要更新,待续。当前日期:2022年10月18日


本文转载自: https://blog.csdn.net/weixin_43490422/article/details/127148825
版权归原作者 little student 所有, 如有侵权,请联系我们删除。

“Yolov5口罩佩戴实时检测项目(模型剪枝+opencv+python推理)”的评论:

还没有评论