ModelBox开发案例 - 使用YOLO v3做口罩检测
ModelBox开发案例 - 使用YOLO v3做口罩检测
本案例将使用YOLO v3模型,实现一个简单的口罩检测应用,最终效果如下所示:
本案例所需资源(代码、模型、测试数据等)已做成模板放到华为云上,查看和下载模板可以使用如下命令:
- Windows PC版本请使用
solution.bat
工具:
PS ███\modelbox>: .\solution.bat -l
...
Solutions name:
mask_det_yolo3
...
结果中的mask_det_yolo3即为口罩检测模板,可使用如下命令下载模板:
PS ███\modelbox>: .\solution.bat -s mask_det_yolo3
...
- Linux开发板版本请使用
solution.py
脚本:
rock@rock-3a:~/███/modelbox$ ./solution.py -l
...
Solutions name:
mask_det_yolo3
...
结果中的mask_det_yolo3即为口罩检测模板,可使用如下命令下载模板:
rock@rock-3a:~/███/modelbox$ ./solution.py -s mask_det_yolo3
...
solution.bat
/solution.py
工具的参数中,-l
代表list
,即列出当前已有的模板名称;-s
代表solution-name
,即下载对应名称的模板。下载下来的模板资源,将存放在ModelBox
核心库的solution
目录下。
如果对ModelBox AI应用开发还不熟悉,请先阅读ModelBox 端云协同AI开发套件(RK3568)上手指南,或者ModelBox 端云协同AI开发套件(博时特EC02)上手指南。
模型准备
本案例使用的是YOLOv3_ResNet18
模型,在下载的资源包中,已经包含了转换好的rknn模型或者onnx模型,可以在对应的平台上直接使用。
原始模型是通过TensorFlow框架训练得到的,我们可以用AI Gallery中的物体检测YOLOv3_ResNet18算法和口罩检测小数据集,自行训练出该模型。
得到TensorFlow Frozen Graph格式的模型后,如果想体验rknn模型的转换过程,rknpu版本可参考RK1808模型转换验证案例,rknpu2版本可参考RK3568模型转换验证案例。
应用开发
打开VS Code,连接到ModelBox sdk所在目录或者远程开发板,开始进行口罩检测应用开发。下面以RK3568版本为例进行说明,其他版本与之类似。
1)创建工程
使用create.py
创建mask_det
工程(Windows PC版本请使用create.bat
):
rock@rock-3a:~/███/modelbox$ ./create.py -t server -n mask_det -s mask_det_yolo3
sdk version is modelbox-rk-aarch64
success: create mask_det in /home/rock/███/modelbox/workspace
将会用口罩检测模板创建出一个mask_det工程。
2)查看推理功能单元
AI应用的核心是模型推理部分,mask_det工程的推理功能单元在工程目录下的model/mask_infer
文件夹中,目录结构如下(以RK3568开发板版本为例):
其中yolo3_resnet18_mask_det_288x512-rknpu2.rknn
是转换好的rknn模型,mask_infer.toml
是该模型的ModelBox功能单元配置文件,其内容如下:
# Copyright (c) Huawei Technologies Co., Ltd. 2022. All rights reserved.
[base]
name = "mask_infer" # 功能单元名称
device = "rknpu" # 功能单元运行的设备类型
version = "1.0.0" # 功能单元版本号
type = "inference" # 功能单元类型,推理功能单元此处为固定值inference
virtual_type = "rknpu" # 推理引擎类型,RK3399pro/RK1808设备为rknpu,RK3568设备为rknpu2
group_type = "Inference" # 功能单元分组信息, 推理功能单元默认为Inference
is_input_contiguous = "false" # rk芯片推理时的固定属性
description = "yolo3 model for mask detection" # 功能单元的描述信息
entry = "./yolo3_resnet18_mask_det_288x512_rknpu.rknn" # 模型文件路径,默认在当前路径下
# 模型输入节点描述:yolo3模型有1个输入,uint8格式的图片数据
[input]
[input.input1]
name = "data"
type = "uint8"
device = "rknpu" # 输入数据位于哪种设备
# 模型输出节点描述:yolo3模型有3个输出,不同大小的float类型特征向量
[output]
[output.output1]
name = "yolo/output1"
type = "float"
[output.output2]
name = "yolo/output2"
type = "float"
[output.output3]
name = "yolo/output3"
type = "float"
可以看到该模型有3个输出节点,即YOLO v3模型输出的3个feature map,需要从中解码出检测框。
3)其他逻辑功能单元
后处理功能单元负责从模型推理结果中解码出检测框,在工程目录下的etc/flowunit/yolo3_post
文件夹中:
解码过程的核心逻辑在yolo3_utils.py
文件中,可以查阅YOLO v3模型细节阅读代码。
画图功能单元在工程目录下的etc/flowunit/draw_mask_bbox
文件夹中:
画图的核心逻辑在draw_mask_bbox.py
文件的draw_mask_info
函数中:
def draw_mask_info(self, image, bboxes):
'''在图中画出口罩佩戴信息'''
thickness = 2
font_scale = 1
text_font = cv2.FONT_HERSHEY_SIMPLEX
for bbox in bboxes:
label_index = int(bbox[5])
# 以头肩部为处理对象
if self.labels[label_index] != 'head':
continue
x_min, y_min, x_max, y_max = bbox[0], bbox[1], bbox[2], bbox[3]
# 根据头肩部找到匹配的人脸框
face_bbox = self.find_max_cover_bbox(
bbox, bboxes, 'face', self.face_cover_ratio)
if not face_bbox: # 没找到对应的人脸,输出'unknown'
yellow = (255, 255, 0)
cv2.rectangle(image, (x_min, y_min),
(x_max, y_max), yellow, thickness)
cv2.putText(image, 'unknown', (x_min, y_min-20),
text_font, font_scale, yellow, thickness)
continue
# 根据人脸框找到匹配的口罩框
mask_bbox = self.find_max_cover_bbox(
face_bbox, bboxes, 'mask', self.mask_cover_ratio)
if not mask_bbox: # 没找到对应的口罩框,输出'no mask'
red = (0, 0, 255)
cv2.putText(image, 'no mask', (x_min, y_min-20),
text_font, font_scale, red, thickness)
cv2.rectangle(image, (x_min, y_min),
(x_max, y_max), red, thickness)
else: # 找到对应的口罩框,输出'has mask'
green = (0, 255, 0)
cv2.putText(image, 'has mask', (x_min, y_min-20),
text_font, font_scale, green, thickness)
cv2.rectangle(image, (x_min, y_min),
(x_max, y_max), green, thickness)
cv2.rectangle(image, (mask_bbox[0], mask_bbox[1]),
(mask_bbox[2], mask_bbox[3]), green, thickness)
return image
可以看到,针对每个人,该模型会尝试检测出head(头肩部)、face和mask三个检测框。如果face检测框与mask检测框的重合度大于某个阈值,就判为佩戴口罩;否则,就判为没有佩戴口罩;如果没有检测到face检测框,就会显示Unknown,表示未知。
4)查看流程图
模型推理和配套的功能单元准备好后,我们就可以串联出流程图进行测试了,口罩检测工程的默认流程图为graph/mask_det.toml
(以RK3568开发板版本为例):
# Copyright (c) Huawei Technologies Co., Ltd. 2022. All rights reserved.
[driver]
# 功能单元的扫描路径,包含在[]中,多个路径使用,分隔
# ${HILENS_APP_ROOT} 表示当前应用的实际路径
# ${HILENS_MB_SDK_PATH} 表示ModelBox核心库的实际路径
dir = [
"${HILENS_APP_ROOT}/etc/flowunit",
"${HILENS_APP_ROOT}/etc/flowunit/cpp",
"${HILENS_APP_ROOT}/model",
"${HILENS_MB_SDK_PATH}/flowunit",
]
skip-default = true
[profile]
# 通过配置profile和trace开关启用应用的性能统计
profile = false # 是否记录profile信息,每隔60s记录一次统计信息
trace = false # 是否记录trace信息,在任务执行过程中和结束时,输出统计信息
dir = "${HILENS_DATA_DIR}/mb_profile" # profile/trace信息的保存位置
[flow]
desc = "mask detection example using yolo3 for local video or rtsp video stream" # 应用的简单描述
[graph]
format = "graphviz" # 流程图的格式,当前仅支持graphviz
graphconf = """digraph mask_det {
node [shape=Mrecord]
queue_size = 4
batch_size = 1
# 定义节点,即功能单元及其属性
input1[type=input, flowunit=input, device=cpu, deviceid=0]
data_source_parser[type=flowunit, flowunit=data_source_parser, device=cpu, deviceid=0]
video_demuxer[type=flowunit, flowunit=video_demuxer, device=cpu, deviceid=0]
video_decoder[type=flowunit, flowunit=video_decoder, device=rknpu, deviceid=0, pix_fmt="bgr"]
image_resize[type=flowunit, flowunit=resize, device=rknpu, deviceid=0, image_width=512, image_height=288]
mask_detection[type=flowunit, flowunit=mask_infer, device=rknpu, deviceid=0]
yolo3_post[type=flowunit, flowunit=yolo3_post, device=cpu, deviceid=0]
draw_mask_bbox[type=flowunit, flowunit=draw_mask_bbox, device=cpu, deviceid=0]
video_out[type=flowunit, flowunit=video_out, device=rknpu, deviceid=0]
# 定义边,即功能间的数据传递关系
input1:input -> data_source_parser:in_data
data_source_parser:out_video_url -> video_demuxer:in_video_url
video_demuxer:out_video_packet -> video_decoder:in_video_packet
video_decoder:out_video_frame -> image_resize:in_image
image_resize:out_image -> mask_detection:data
mask_detection:"yolo/output1" -> yolo3_post:in_feat3
mask_detection:"yolo/output2" -> yolo3_post:in_feat2
mask_detection:"yolo/output3" -> yolo3_post:in_feat1
video_decoder:out_video_frame -> draw_mask_bbox:in_image
yolo3_post:out_data -> draw_mask_bbox:in_bbox
draw_mask_bbox:out_image -> video_out:in_video_frame
}"""
该流程图对于某个视频流,经过视频解码、图像缩放、口罩检测推理、检测框后处理、画图等一系列操作后,将结果保存下来。
口罩检测工程的任务配置文件bin/mock_task.toml
中设置了输入输出源,内容为:
# 任务输入配置,当前支持以下几种输入方式:
# 1. rtsp摄像头或rtsp视频流:type="rtsp", url="rtsp://xxx.xxx"
# 2. 设备自带摄像头或者USB摄像头:type="url",url="${摄像头编号}" (需配合local_camera功能单元使用)
# 3. 本地视频文件:type="url",url="${视频文件路径}" (请使用${HILENS_APP_ROOT}宏,表示当前应用的实际路径)
# 4. http服务:type="url", url="http://xxx.xxx"(指的是任务作为http服务启动,此处需填写对外暴露的http服务地址,需配合httpserver类的功能单元使用)
[input]
type = "url"
# url = "0"
url = "${HILENS_APP_ROOT}/data/mask_test.mp4"
# 任务输出配置,当前支持以下几种输出方式:
# 1. rtsp视频流:type="local", url="rtsp://xxx.xxx"
# 2. 本地屏幕:type="local", url="0:xxx" (设备需要接显示器,系统需要带桌面)
# 3. 本地视频文件:type="local",url="${视频文件路径}"(请使用${HILENS_APP_ROOT}宏,表示当前应用的实际路径)
# 4. http服务:type="webhook", url="http://xxx.xxx"(指的是任务产生的数据上报给某个http服务,此处需填写上传的http服务地址)
[output]
type = "local"
# url = "0:mask_det" # "rtsp://192.168.2.2:8554/outstream"
url = "${HILENS_APP_ROOT}/hilens_data_dir/mask_test_result.mp4"
测试视频为data/mask_test.mp4
,该流程图使用这一视频进行口罩检测,检测结果绘制后保存为hilens_data_dir/mask_test_result.mp4
文件。
5)运行应用
在mask_det
工程路径下执行build_project.sh
进行工程构建(以RK3568开发板版本为例):
rock@rock-3a:~/███/modelbox/workspace/mask_det$ ./build_project.sh
build success: you can run main.sh in ./bin folder
rock@rock-3a:~/███/modelbox/workspace/mask_det$
然后执行bin/main.sh
(Windows PC版本请使用bin\main.bat
)运行应用,运行结束后在hilens_data_dir
目录下生成了mask_test_result.mp4
文件,可以下载到PC端查看。
- 点赞
- 收藏
- 关注作者
评论(0)