Jetson Nano 2GB部署NanoDet实现火焰检测

前言

距离Nvidia推出Jetson Nano开发板已经有一段时间
在GTC2020上NV再次推出了定价$59的2GB版Nano成为树莓派有利竞争对手

同样的价格下自带128个cuda和免费DLI课程让新版Jetson Nano非常适合部署边缘AI项目
在拿到开发板的第一时间我决定做一个小项目对Nano进行体验
本文将结合机器视觉与AI实现实时火焰检测，输出real-time bounding box与概率

网络选择

动态目标检测有很多成熟的网络可供选择，如YOLO、SSD、RCNN等
这些网络效果很好但算力开销大，不适合在Jetson Nano 2GB上实时运行
通过搜索我发现了一个最近开源的快速目标检测网络NanoDet
详见作者的Github
对于320 x 320的输入FLOPS仅为0.72B非常适合嵌入式部署

环境搭建

参考NanoDet作者给出的依赖，需要的环境为

Cython
termcolor
numpy
torch>=1.3
torchvision
tensorboard
pycocotools
matplotlib
pyaml
opencv-python
tqdm

在Jetson Nano部署主要难点是pytorch安装，使用Nvidia官方预编译whl文件，详见
Nvidia Developer Forums
注意torch和vison版本匹配，我的环境是 PyTorch v1.6 + torchvision v0.7.0

数据标注

火焰检测不属于voc标准分类，因此需要自己进行标注训练
使用图片搜索引擎爬取火焰图片，清洗筛选后留下397张构建数据集
数据标注使用开源工具labelImg进行标注，保存为voc格式

按照8:2构建训练集与测试集，结构如下

fire
├── train
│   ├── ann
│   │   ├── 1.xml
│   │   └── 2.xml
│   └── img
│       ├── 1.jpg
│       └── 2.jpg
└── val
    ├── ann
    │   └── 1.xml
    └── img
        └── 1.jpg

修改配置

修改网络的yml配置文件
默认配置为80类目标检测，
自定义保存位置、识别种类、图像宽高、训练集、测试集，修改以下部分：

save_dir: ./fire
num_classes: 1
class_names: &class_names ['fire']
train:
  name: xml_dataset
  img_path: ./fire/train/img
  ann_path: ./fire/train/ann
  input_szie: [320,320]
val:
  name: xml_dataset
  img_path: ./fire/val/img
  ann_path: ./fire/val/ann
  input_szie: [320,320]

注意：默认数据集为coco格式，需要改为xml_dataset

训练网络

这里使用PC端RTX2080进行训练

python3 train.py fire/nanodet_fire.yml

配置batch_size = 80 total_epochs = 160 耗时2.5小时

结果权重文件保存在./fire/model_last.pth

部署推理

得到权重文件即可在Jetson Nano上进行实时推理

pyhton3 livecam.py

使用ShuffleNetV2作为backbone首次运行需要下载一个预训练模型
首次部署时确保Jetson Nano连接到网络，后续可离线使用

实测Jetson Nano 2GB输入640x480图像实时检测可以跑到16FPS

视频代码

演示视频已上传至B站：https://www.bilibili.com/video/BV1np4y1z7kF/
项目代码Github地址：https://github.com/azureology/jetson-nano-fire-detection