Jetson 系列——基于yolov5和deepsort的多目标头部识别,跟踪,使用tensorrt和c++加速

项目简介:

代码地址:
github地址:https://github.com/RichardoMrMu/yolov5-deepsort-tensorrt
gitee地址:https://gitee.com/mumuU1156/yolov5-deepsort-tensorrt
欢迎start或提出你的issue、改进建议。
本项目使用yolov5+deepsort实现室内人头部追踪和计数,使用c++实现,并用tensorrt加速,在物体有70+左右的情况下,在Jetson Xavier nx上整个项目的推理时间在130ms左右,即7FPS的速度。你可以体验一下python版本的yolov5+deepsort,使用原生pytorch,当跟踪的物体达到70+的时候,deepsort一次推理的速度将到1s左右,根本无法实现实时推理。
You can see video play in BILIBILI, or YOUTUBE and YOUTUBE.

效果


这里我是用我自己训练的模型,识别人物的头部。你也可以尝试yolov5的官方pretrained模型,然后根据教程将pytorch的模型转为tensorrt使用的engine文件,即可进行1000种类别的检测。

Speed

Backbone before TensorRT without tracking before TensorRT with tracking TensorRT(detection + tracking) FPS(detection + tracking)
Yolov5s_416 100ms 0.9s 450ms 1.5 ~ 2
Yolov5s-640 120ms 1s 100-150ms 8 ~ 9

项目使用

这个项目中,需要两个模型,检测的yolov5和deepsort特征提取的ckpt.t7,对于yolov5,我选择yolov5s模型,并通过yolov5s.pt->yolov5s.wts->yolov5s.engine,获取engine file。deepsort模型在官方仓库里有训练好的模型deepsort。如果你要训练自己的模型,你可以参考custom model,待补充。同时你也可以参考tensorrtx official readme。
下面是我转换好的deepsort.onnx和deepsort.engine文件,可以直接下载使用。

模型 链接
百度云 BaiduYun url 密码:z68e

0.数据集

如果你想训练自己的头部检测模型,你可以使用SCUT-HEAD这个数据集,它是公开的,且可以随意下载的,有头部框的标注。

1. requirement

请注意,下列流程,我在Jetson Xavier nxJetson 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-deepsort-tensorrt

如果你已经有了yolov5s.engine和deepsort.engine,那么仅下载{yolov5-deepsort-tensorrt}仓库,然后修改{yolov5-deepsort-tensorrt}/src/main.cpp
char* yolo_engine = "";char* sort_engine = "";为自己模型的位置 ,执行编译链接即可。如果没有这两个文件,请参考3往后的步骤。

// 下载c++项目代码
git clone https://github.com/RichardoMrMu/yolov5-deepsort-tensorrt.git
// {}代表该项目的绝对路径
cd {yolov5-deepsort-tensorrt}
mkdir build
// 修改{yolov5-deepsort-tensorrt}/src/main.cpp中deepsort
//char* yolo_engine = ""; 和 char* sort_engine = "";为自己模型的位置
// frame = capture.open("");将open中换成mp4所在的路径,或者如果使用usb摄像头,就写成capture.open(0)
// 然后执行编译链接即可
cmake ..
make
// make成功后,会在build文件夹下生成一个yolosort二进制文件 运行即可
./yolosort // 运行二进制文件

3. 将pytorch的模型转为tensorrt支持的模型文件

因为python自身速度的问题和pytorch原生模型速度的问题,如果将检测部分的yolov5模型和track模块特征提取的模型转为tensorrt,速度会快10倍左右。

  1. 下载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
  1. 下载并配置tensorrt
git clone https://github.com/wang-xinyu/tensorrtx
  1. 将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模型

  1. 编译tensorrtx/yolov5项目并获取tensorrt engine file
    此处请注意,如果你是自己训练的模型,那么要将yololayer.h中的CLASS_NUMINPUT_H,INPUT_W修改为自己的参数。即文件的20,21,22行。
cd tensorrtx/yolov5
// update CLASS_NUM in yololayer.h if your model is trained on custom dataset
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-deepsort-tensorrt}/resources文件夹放模型。

cd {yolov5-deepsort-tensorrt}
mkdir resources
cp {tensorrtx}/yolov5/build/yolov5s.engine {yolov5-deepsort-tensorrt}/resources
  1. 获取deepsort特征提取的engine文件

你可以从deepsort中找到训练好的模型,你也可以直接下载下面的连接drive url,下载ckpt.t7就好了。我也放在了百度云一份,如果你连接有麻烦,可以下载百度云上的模型。

模型 链接
google drive drive url
百度云 BaiduYun url 密码:x5mj
  • 2021-9-16日补充
    对于deep_sort_pytorch,该仓库进行了更新,使得直接下载后,将exportOnnx.py复制到该仓库代码路径下,运行会报错。No module named 'fastreid'.,你可以根据该仓库的第五步提示,安装fastreid,也可以下载之前版本的deep_sort_pytorch,地址:
    deep_sort_pytorch官方仓库旧版本地址:https://github.com/ZQPei/deep_sort_pytorch/tree/d9027f9d230633fdab23fba89516b67ac635e378
    点击download zip下载代码即可。
    我的github备份地址:https://github.com/RichardoMrMu/deep_sort_pytorch
    对于我github上备份的代码,可以直接执行
git clone https://github.com/RichardoMrMu/deep_sort_pytorch.git
// 或者点击download zip

补充完毕。

git clone https://github.com/RichardoMrMu/deepsort-tensorrt.git
// 根据github的说明
cp {deepsort-tensorrt}/exportOnnx.py {deep_sort_pytorch}/
python3 exportOnnx.py
mv {deep_sort_pytorch}/deepsort.onnx {deepsort-tensorrt}/resources
cd {deepsort-tensorrt}
mkdir build
cd build
cmake ..
make 
.onnx2engine ../resources/deepsort.onnx ../resources/deepsort.engine
// test
./demo ../resource/deepsort.engine ../resources/track.txt

经过以上的步骤,获得了yolov5.engine和deepsort.engine,将这两个文件放到相应的位置。即可编译运行该项目。
如果你遇到问题,请在github项目yolov5-deepsort-tensorrt的issue中提问题,或者在这篇文章下评论。欢迎讨论。

Other Project

  • Jetson 系列——基于yolov5对是否吸烟的检测,部属于jetson xavier nx,使用tensorrt、c++和int8加速,可用于边缘计算

  • Jetson 系列——基于yolov5对是否带口罩的检测,部属于jetson xavier nx,使用tensorrt、c++和int8加速,可用于边缘计算

  • Jetson 系列——基于yolov5对反光衣和安全帽的检测,部属于jetson xavier nx,使用tensorrt、c++和int8加速,可用于边缘计算

  • DL Project系列——基于gazecapture,mediapipe和TF.js,Flask的视线估计和脸部特征点估计

  • Jetson 系列——基于yolov5对工地头盔的检测,部属于jetson xavier
    nx,使用tensorrt、c++和int8加速,可用于边缘计算

  • Jetson 系列——基于yolov5对火源或者烟雾的检测,使用tensorrt、c++和int8加速

  • 竞赛——人脸情绪识别挑战赛,科大讯飞&DataWhale
    第三名

  • Jetson 系列——基于deepsort的物体跟踪使用tensorrt和c++加速

  • Jetson 系列——基于yolov5和deepsort的多目标头部识别,跟踪,使用tensorrt和c++加速

你可能感兴趣的:(jetson,nane,or,xavier,实验,deep,learning,project,c++,pytorch,python,deepsort,yolo)