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

  • Post author:
  • Post category: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日


版权声明:本文为weixin_43490422原创文章,遵循 CC 4.0 BY-SA 版权协议,转载请附上原文出处链接和本声明。