NVIDIA之TLT迁移学习口罩实时检测

0 背景

在医院、机场、车站等场所中,需要对行人进行实时口罩检测。本文基于NVIDIA TLT工具,介绍口罩检测模型的训练及部署方法,实现在嵌入式端实时处理的效果。在嵌入式上部署推理性能表如下,可以看到裁剪后的模型,即使在 Jetson Nano平台上也能跑搭配20帧,做到实时处理的效果。

  • 注意:这个帧率是使用deepstream获得的端到端帧率,其中nano是使用的FP16精度,Xavier NX和AGX使用的是int8精度
  • NX和AGX有两种推理模块,分别是GPU和DLA,当在DLA上推理的时候,GPU是空闲的,可以处理其它任务

流程如下

官方博客:https://developer.nvidia.com/blog/implementing-a-real-time-ai-based-face-mask-detector-application-for-covid-19/

--------------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.txt

2 数据集准备

2.1 下载数据集

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
    │   └───images

2.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 4

4.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 $KEY

4.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 $KEY

5 部署

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 $KEY

5.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 运行检测即可。

 
已标记关键词 清除标记
©️2020 CSDN 皮肤主题: 游动-白 设计师:白松林 返回首页