用 yolov5s 训练口罩检测模型，并在 jetson nano 下部署测试

一 训练环境:

ubuntu 20.04  + pytorch 1.7 + torchvison 0.8 

gpu: gtx1080 8GB 

 

二 下载 yolov5 模型及权重:

https://github.com/ultralytics/yolov5

利用 weigths/download_weights.sh 来下载 yolov5s.pt , 必要的时候 科鞋上网

 

三 组织 训练数据 

3.1 数据集下载 (感谢学长):

https://pan.baidu.com/s/1ZJbzf5qF9GKEqyo0RskWiQ

密码：  53hn

3.2 准备 yolo 训练使用的数据集格式

把第一个 红框的数据集拿来训练 ， 我整理好了 yolo系列需要的数据格式，分享一下：

https://github.com/walletiger/facemask_annotations

 

一图来了解  ultralytics/yolovXX 系列模型需要的数据集格式：

把网上下载的 口罩数据集 train 和 val 数据解压 到 images/train , images/val

3.3 准备 facemask.yaml 文件

放到 yolov5/data目录下， 格式如下：

train: /workspace/data/facemark/facemask_annotations/facemask_labels_yolov5/dataset/images/train  # 128 images
val: /workspace/data/facemark/facemask_annotations/facemask_labels_yolov5/dataset/images/val   # 128 images

# number of classes
nc: 2

# class names
names: ['face', 'facemark']

四 训练 过程

python3 train.py --img 640 --batch 16 --epochs 160 --data facemask.yaml --weights yolov5s.pt  --cache-images

在 runs/train/expX 下可以看到 训练 收敛过程和 验证结果

 

results.txt 记录了每个 epoch 训练的 精度 、召回率、mAP等数据，

train_batchx.jpg 是阶段性验证结果 ，如下示例：

 

五 在jetson nano 下验证模型

5.1 下载 yolov5 代码

git clone https://github.com/ultralytics/yolov5

5.2 上传训练 好的模型：

我的在训练 主机 runs/exp3/weights/best.pt ， 放到  jetson nano yolov5/

5.3 测试一把 video 检测

 python3 detect.py --source /workspace/data/360p1.mp4 --weights best.pt --conf 0.3

准确率还可以， 速度不快 110~120 ms 一帧

据说 可以通过 转 onnx -> tensortRT 获得开挂般的性能。。 后续 再研究。

转 tensorRT :

https://blog.csdn.net/walletiger/article/details/111027113