0 背景
在医院、机场、车站等场所中，需要对行人进行实时口罩检测。本文基于NVIDIA TLT工具，介绍口罩检测模型的训练及部署方法，实现在嵌入式端实时处理的效果。在嵌入式上部署推理性能表如下，可以看到裁剪后的模型，即使在 Jetson Nano平台上也能跑搭配20帧，做到实时处理的效果。
- 注意：这个帧率是使用deepstream获得的端到端帧率，其中nano是使用的FP16精度，Xavier NX和AGX使用的是int8精度
- NX和AGX有两种推理模块，分别是GPU和DLA，当在DLA上推理的时候，GPU是空闲的，可以处理其它任务
流程如下
--------------20210201更新-------------------
NGC解决方案
1 安装
首先按照我之前的文章《NVIDIA之TLT迁移学习训练自己的数据集》部署一下容器环境，进入容器环境中
git clone https://github.com/NVIDIA-AI-IOT/face-mask-detection.git
cd tlt-ds-face_mask_detect
python -m pip install -r requirements.txt2 数据集准备
2.1 下载数据集
- Faces with Mask
- Kaggle Medical Mask Dataset Download Link
- MAFA - MAsked FAces Download Link（链接有问题）
- Faces without Mask
- FDDB Dataset Download Link
- WiderFace Dataset Download Link
MAFA官方数据集出现问题，可通过网盘下载（链接: https://pan.baidu.com/s/1qiflOq65yMNNTXyt8020GA 提取码: 4fz6）
下载好之后按照现编的目录结构解压
data
├───MAFA
│ ├───MAFA-Label-Test
│ │ LabelTestAll.mat
│ │ readme-test.txt
│ │
│ ├───MAFA-Label-Train
│ │ LabelTrainAll.mat
│ │ readme-train.txt
│ │
│ │
│ ├───test-images
│ │ └───images
│ │ test_00000001.jpg
│ │ .........
│ │ .........
│ │ test_00004935.jpg
│ │
│ ├───train-images
│ └───images
│ train_00000001.jpg
│ .........
│ .........
│ train_00025876.jpg
│
│
├───FDDB
│ ├───2002
│ ├───2003
│ ├───FDDB-folds
│ FDDB-fold-01-ellipseList.txt
│ FDDB-fold-01.txt
│ ..........
│ ..........
│ FDDB-fold-10-ellipseList.txt
│ FDDB-fold-10.txt
│
│
├───Kaggle
│ ├───Medical Mask Dataset
│ ├───images
│ │ -1x-1.jpg
│ │ .........
│ │ .........
│ │ w1240-p16x9-fa978043deff83fed485af12d16e39c61398fc30.jpg
│ │ W37H3GWNO5EV5HTJQNP4KLIRGI.jpg
│ │
│ └───labels
│ -1x-1.xml
│ .........
│ .........
│ w1240-p16x9-fa978043deff83fed485af12d16e39c61398fc30.xml
│ W37H3GWNO5EV5HTJQNP4KLIRGI.xml
│
├───Wider
├───wider_face_split
│ readme.txt
│ wider_face_test.mat
│ wider_face_test_filelist.txt
│ wider_face_train.mat
│ wider_face_train_bbx_gt.txt
│ wider_face_val.mat
│ wider_face_val_bbx_gt.txt
│
├───WIDER_train
│ └───images
├───WIDER_val
│ └───images2.2 格式转换
使用下边的方法转换数据集格式，注意要根据自己数据集的真实路径进行指定
python data2kitti.py --kaggle-dataset-path=data/Kaggle/medical-masks-dataset --mafa-dataset-path=data/MAFA --fddb-dataset-path=data/FDDB --widerface-dataset-path=data/Wider --kitti-base-path=data/kitti/ --train --category-limit=20000
输出如下
Directory Already Exists
Directory Already Exists
Kaggle Dataset: Total Mask faces: 4154 and No-Mask faces:790
Total Mask Labelled:4154 and No-Mask Labelled:790
Directory Already Exists
Directory Already Exists
/workspace/tlt-experiments/face-mask-detection/data_utils/mafa2kitti.py:51: RuntimeWarning: overflow encountered in ubyte_scalars
bbox = [_bbox_label[0], _bbox_label[1], _bbox_label[0]+_bbox_label[2], _bbox_label[1]+_bbox_label[3]]
MAFA Dataset: Total Mask faces: 1846 and No-Mask faces:232
Total Mask Labelled:6000 and No-Mask Labelled:1022
Directory Already Exists
Directory Already Exists
FDDB Dataset: Mask Labelled:0 and No-Mask Labelled:2845
Total Mask Labelled:6000 and No-Mask Labelled:3867
WideFace: Total Mask Labelled:0 and No-Mask Labelled:2134
----------------------------
Final: Total Mask Labelled:6000
Total No-Mask Labelled:6001
----------------------------
注意：原始脚本中默认包含6000个戴口罩和不带口罩的目标，大概包含2000张图片，如果想多使用一些数据集，可以使用 --category-limit=10000或其它数字来指定数据集大小
3 首次训练
3.1 训练模型
启动 jupyter notebook
jupyter notebook --ip 0.0.0.0 --allow-root设置环境变量
%env KEY=YjlxOTRkaHRjYWI2Z2NxN2cwOXBlZjh1OTQ6ZTE2YjdkNzctMmQ0OS00MDZhLTgzMGMtNjc5ZTIyZGNkNzA1
%env USER_EXPERIMENT_DIR=/workspace/tlt-experiments/face-mask-detection/data
%env DATA_DOWNLOAD_DIR=/workspace/tlt-experiments/face-mask-detection/data
%env SPECS_DIR=./tlt_specs 训练
tlt-train detectnet_v2 -e $SPECS_DIR/detectnet_v2_train_resnet18_kitti.txt \
-r $USER_EXPERIMENT_DIR/experiment_dir_unpruned \
-k $KEY \
-n resnet18_detector \
--gpus 4训练到自己设置的回合数后，可以验证一下模型
3.2 验证模型
!tlt-evaluate detectnet_v2 -e $SPECS_DIR/detectnet_v2_train_resnet18_kitti.txt\
-m $USER_EXPERIMENT_DIR/experiment_dir_unpruned/weights/resnet18_detector.tlt \
-k $KEY输出如下
Validation cost: 0.000681
Mean average_precision (in %): 83.2432
class name average precision (in %)
------------ --------------------------
mask 87.3164
no-mask 79.17表明训练到一个不错的模型
4 再次训练
可以对首次训练的模型进行裁剪，然后再次训练，得到一个体积更小的模型，有利于在终端部署
# Create an output directory if it doesn't exist.
!mkdir -p /workspace/tlt-experiments/face-mask-detection/data/experiment_dir_pruned
!tlt-prune -m /workspace/tlt-experiments/face-mask-detection/data/experiment_dir_unpruned/weights/resnet18_detector.tlt \
-o /workspace/tlt-experiments/face-mask-detection/data/experiment_dir_pruned/resnet18_nopool_bn_detectnet_v2_pruned.tlt \
-eq union \
-pth 0.8 \
-k $KEY运行后在 experiment_dir_pruned 路径下生成 resnet18_nopool_bn_detectnet_v2_pruned.tlt，用于下一步的训练
4.1 修改配置文件
修改 detectnet_v2_retrain_resnet18_kitti.txt 配置文件中对应内容
4.2 训练
# Retraining using the pruned model as pretrained weights
tlt-train detectnet_v2 -e $SPECS_DIR/detectnet_v2_retrain_resnet18_kitti.txt \
-r /workspace/tlt-experiments/face-mask-detection/data/experiment_dir_retrain \
-k $KEY \
-n resnet18_detector_pruned --gpus 44.3 验证精度
tlt-evaluate detectnet_v2 -e $SPECS_DIR/detectnet_v2_retrain_resnet18_kitti.txt \
-m $USER_EXPERIMENT_DIR/experiment_dir_retrain/weights/resnet18_detector_pruned.tlt \
-k $KEY4.4 可视化结果
tlt-infer detectnet_v2 -e $SPECS_DIR/detectnet_v2_inference_kitti_tlt.txt \
-o /workspace/tlt-experiments/face-mask-detection/data/tlt_infer_testing_pruned \
-i /workspace/tlt-experiments/face-mask-detection/data/test_images_thermal \
-k $KEY5 部署
5.1 导出模型
mkdir -p $USER_EXPERIMENT_DIR/experiment_dir_final
# Removing a pre-existing copy of the etlt if there has been any.
import os
output_file=os.path.join(os.environ['USER_EXPERIMENT_DIR'],
"experiment_dir_final/resnet18_detector_thermal.etlt")
if os.path.exists(output_file):
os.system("rm {}".format(output_file))
tlt-export detectnet_v2 \
-m $USER_EXPERIMENT_DIR/experiment_dir_retrain/weights/resnet18_detector_pruned.tlt \
-o $USER_EXPERIMENT_DIR/experiment_dir_final/resnet18_detector_thermal.etlt \
-k $KEY5.2 deepstream部署
下载相关配置文件
git clone https://github.com/NVIDIA-AI-IOT/face-mask-detection.git修改ds_configs中的 config_infer_primary_masknet_gpu.txt 和 deepstream_app_source1_camera_masknet_gpu.txt
然后使用 deepstream-app 运行检测即可。
1154
