项目简介:
代码地址:
github地址:https://github.com/RichardoMrMu/yolov5-fire-smoke-detect gitee地址:https://gitee.com/mumuU1156/yolov5-fire-smoke-detect 欢迎start或提出你的issue、改进建议。
本项目使用yolov5实现自然状态中火源或者烟雾的检测,使用c++实现,并用tensorrt加速,尝试了不同大小的模型,同时使用int8加速,在Jetson Xavier nx上整个项目的推理时间在30ms左右,即33FPS的速度。
You can see video play in BILIBILI, or YOUTUBE.
如果你想从头开始尝试训练模型,你可以查看 yolov5-fire-smoke-detect-python这个仓库,并阅读他的readme,帮助你是用python去训练一个自己的烟雾和火焰检测模型,并使用c++去部属到边缘设备上。
效果
如果你想在预训练模型上进行微调,可以尝试yolov5的官方pretrained模型,然后在对应数据集上进行训练,根据教程将pytorch的模型转为tensorrt使用的engine文件。
Speed
Backbone | before TensorRT | TensortRT(detection) | FPS(detection) |
Yolov5s-640-float16 | 100ms | 60-70ms | 14 ~ 18 |
Yolov5m-640-float16 | 120ms | 70-75ms | 13 ~ 14 |
Yolov5s-640-int8 | 30-40ms | 25 ~ 33 | |
Yolov5m-640-int8 | 50-60ms | 16 ~ 20 |
Achieve and Experiment
- Int8.
- yolov5-s
- yolov5-m
Comming soon
DataSet
本案例使用数据集包含MIT协议互联网下载图片和MIT协议公开数据集,共有6675张图片,分别检测烟雾和火灾
数据集图片格式是VOC数据格式,VOC数据是每个图像文件对应一个同名的xml文件,xml文件内包含对应图片的基本信息,比如文件名、来源、图像尺寸以及图像中包含的物体区域信息和类别信息等。
xml文件中包含以下字段:
- filename,表示图像名称。
- size,表示图像尺寸。包括:图像宽度、图像高度、图像深度。
<size>
<width>500</width>
<height>375</height>
<depth>3</depth>
</size>
- object字段,表示每个物体。包括:
标签 | 说明 |
name | 物体类别名称 |
pose | 关于目标物体姿态描述(非必须字段) |
truncated | 如果物体的遮挡超过15-20%并且位于边界框之外,请标记为 |
difficult | 难以识别的物体标记为 |
bndbox子标签 | (xmin,ymin) 左上角坐标,(xmax,ymax) 右下角坐标, |
将这6675张图片按9:1比例随机切分,切分后包含6008张图片的训练集和667张图片的验证集,提供处理好的数据集,下载地址。包含以下文件夹和文件:images
,annotations
,labels.txt
, train_list.txt
和 val_list.txt
,分别图片、xml标注文件、存储类别信息、训练样本列表、验证样本列表。训练样本列表和验证样本列表的每一行格式为:图片路径 对应的xml路径
,例如images/fire_00348.jpg annotations/fire_00348.xml
。
最终数据集文件组织结构为:
├── dataset
├── annotations
│ ├── fire_000001.xml
│ ├── fire_000002.xml
│ ├── fire_000003.xml
│ | ...
├── images
│ ├── fire_000001.jpg
│ ├── fire_000003.jpg
│ ├── fire_000003.jpg
│ | ...
├── label_list.txt
├── train.txt
└── valid.txt
项目使用
这个项目中,需要检测的yolov5,我选择yolov5s模型,并通过yolov5s.pt->yolov5s.wts->yolov5s.engine
,获取engine file。同时你也可以参考tensorrtx official readme。
0.Python Code
你可以下载这个yolov5-fire-smoke-detect-python的代码进行训练,推理等操作,也可以直接下载模型。
python仓库的操作步骤,请参看该仓库的README.md。
1. requirement
请注意,下列流程,我在Jetson Xavier nx
和Jetson nano
上完整实现过,你同样可以在ubuntu系统上实现,如果想在windows上测试,需要你自己摸索。
Jetson nano or Jetson Xavier nx
Jetpack 4.5.1
python3 with default(jetson nano or jetson xavier nx has default python3 with tensorrt 7.1.3.0 )
tensorrt 7.1.3.0
torch 1.8.0
torchvision 0.9.0
torch2trt 0.3.0
onnx 1.4.1
opencv-python 4.5.3.56
protobuf 3.17.3
scipy 1.5.4
2. 下载yolov5-fire-smoke-detect
如果你已经有了yolov5s.engine,那么仅下载{yolov5-fire-smoke-detect}
仓库,然后修改{yolov5-deepsort-tensorrt}/src/main.cpp
中char* yolo_engine = "";
为自己模型的位置 ,执行编译链接即可。如果没有这个文件,请参考3往后的步骤。
// 下载c++项目代码
git clone https://github.com/RichardoMrMu/yolov5-fire-smoke-detect
// {}代表该项目的绝对路径
cd {yolov5-fire-smoke-detect}
mkdir build
// 修改{yolov5-fire-smoke-detect}/src/main.cpp中deepsort
//char* yolo_engine = ""; 为自己模型的位置
// frame = capture.open("");将open中换成mp4所在的路径,或者如果使用usb摄像头,就写成capture.open(0)
// 然后执行编译链接即可
cmake ..
make
// make成功后,会在build文件夹下生成一个yolosort二进制文件 运行即可
./yolofire // 运行二进制文件
3. 将pytorch的模型转为tensorrt支持的模型文件
因为python
自身速度的问题和pytorch原生模型速度的问题,如果将检测部分的yolov5模型和track模块特征提取的模型转为tensorrt,速度会快10倍左右。
- 下载yolov5官方代码和模型
请注意,我下面使用的yolov5代码版本是yolov5-5,即v5.0版本,所以,如果你想自己训练模型替换engine file,请确保的yolov5代码版本是v5.0。
// 下载v5.0
git clone -b v5.0 https://github.com/ultralytics/yolov5.git
cd yolov5
mkdir weights
cd weights
// download https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s.pt
// 这里使用yolov5s模型
wget https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s.pt
- 下载并配置tensorrt
git clone https://github.com/wang-xinyu/tensorrtx
- 将pytorch的yolov5s.pt模型文件转为wst文件
将tensorrtx下的get_wts.py文件拷贝到yolov5文件夹下,然后将下载好的或者训练好的yolov5s.pt模型转换为wts模型
cp tensorrtx/gen_wts.py yolov5/
cd yolov5
python3 gen_wts.py -w ./weights/yolov5s.pt -o ./weights/yolov5s.wts
// a file 'yolov5s.wts' will be generated.
你可以在 yolov5/weights/
找到yolov5s.wts
模型
- 编译
tensorrtx/yolov5
项目并获取tensorrt engine file
此处请注意,如果你是自己训练的模型,那么要将yololayer.h
中的CLASS_NUM
,INPUT_H
,INPUT_W
修改为自己的参数。即文件的20,21,22行。因为此处是二分类,即火焰和烟,因此CLASS_NUM为2。
cd tensorrtx/yolov5
// update CLASS_NUM as 2 in yololayer.h
mkdir build
cd build
cp {ultralytics}/yolov5/yolov5s.wts {tensorrtx}/yolov5/build
cmake ..
make
// yolov5s
sudo ./yolov5 -s yolov5s.wts yolov5s.engine s
// test your engine file
sudo ./yolov5 -d yolov5s.engine ../samples
然后你就得到了yolov5.engine文件,你需要将这个文件放到我的项目中,你可以创建一个{yolov5-fire-smoke-detect}/resources
文件夹放模型。
cd {yolov5-fire-smoke-detect}
mkdir resources
cp {tensorrtx}/yolov5/build/yolov5s.engine {yolov5-fire-smoke-detect}/resources
经过以上的步骤,获得了yolov5.engine,将这个文件放到相应的位置。即可编译运行该项目。
INT8 Quantization int8模型量化
int8量化和之前的float16有些许差别,tensorrt的int8量化需要图片的calibration,正如tensorrtx readme所说。
对于一般的yolo模型训练,你可能需要下载coco_calid.zip
,你可以通过这个google drive url or BAIDUYUN — 密码是a9wh
. 随后解压到{project}/build/
.
然后你需要将yolov5.cpp
文件中第十行的USE_FLOAT16
改成USE_INT8
。
然后重新编译运行:
cmake ..
make
// yolov5s
sudo ./yolov5 -s yolov5s.wts yolov5s-int8.engine s
// testyour engine file
sudo ./yolov5 -d yolov5s-int8.engine ../samples
这样你就得到了int8量化的yolov5模型了,你可以尝试一下,是否加速了。
如果你遇到问题,请在github项目yolov5-fire-smoke-detect的issue中提问题,或者在这篇文章下评论。