Jetson Xavier NX和Jetson Nano上的YOLOv5和DeepSort的C++实现

这个仓库使用YOLOv5和DeepSort来跟踪人头，可以在Jetson Xavier NX和Jetson Nano上运行。在Jetson Xavier NX上，当图像中包含70多个头部时，它可以达到10 FPS的性能（你可以尝试Python版本，当使用Python版本时，你会发现在Jetson Xavier NX上非常慢，DeepSort可能需要近1秒的处理时间）。

感谢B.Q Long提供Windows版本的CMakeLists.txt。如果你想在Windows上运行这个仓库，你可以使用CMakeLists_deepsort-tensorrt_win10.txt和CMakeLists_yolov5-deepsort-tensorrt_win10.txt。

你可以在哔哩哔哩、YouTube和YouTube上观看视频演示。

要求

Jetson Nano或Jetson Xavier NX
Jetpack 4.5.1
Python3（默认版本，Jetson Nano或Jetson Xavier NX默认带有TensorRT 7.1.3.0的Python3）
TensorRT 7.1.3.0
PyTorch 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

如果你在这个项目中遇到问题，可以查看这篇文章。

即将推出

Int8优化
IOU跟踪
更快速且占用更少内存

速度

从读取图像到完成DeepSort的整个处理时间（包括每张图像的预处理和后处理）注意！DeepSort跟踪的对象数量是70+，不是单个或10-20人，而是70+。所有结果都是在Jetson Xavier NX上获得的。

骨干网络	TensorRT优化前（不含跟踪）	TensorRT优化前（含跟踪）	TensorRT优化后（仅检测）	TensorRT优化后（检测+跟踪）	FPS（检测+跟踪）
YOLOv5s_416	100ms	0.9s	10-15ms	100-150ms	8 ~ 9
YOLOv5s-640	120ms	1s	18-20ms	100-150ms	8 ~ 9

构建和运行

git clone https://github.com/RichardoMrMu/yolov5-deepsort-tensorrt.git
cd yolov5-deepsort-tensorrt
// 在运行cmake和make之前，请将./src/main.cpp中的char* yolo_engine = "";和char* sort_engine = "";修改为你自己的路径
mkdir build 
cmake ..
make

如果你在cmake和make过程中遇到错误，请查看这篇文章或查看注意事项。

数据集

如果你需要训练自己的头部检测模型，可以使用SCUT-HEAD数据集，这个数据集包含头部边界框标注，可以免费下载。

模型

你需要两个模型，一个是YOLOv5模型用于检测，从tensorrtx生成。另一个是DeepSort模型用于跟踪。你应该以相同的方式生成这两个模型。

生成YOLOv5模型

对于YOLOv5检测模型，我选择了YOLOv5s，并选择了yolov5s.pt->yolov5s.wts->yolov5s.engine的转换路径。注意，使用的模型可以从YOLOv5和DeepSort获取，如果你需要使用自己的模型，可以按照"运行自定义模型"的说明进行操作。你也可以查看tensorrtx官方README。以下是deepsort.onnx和deepsort.engine文件，你可以在百度网盘和https://github.com/RichardoMrMu/yolov5-deepsort-tensorrt/releases/tag/yolosort找到。

模型	链接
百度网盘	百度网盘链接密码：`z68e`

获取YOLOv5仓库

注意，这里使用官方预训练模型。我使用的是YOLOv5-5，v5.0版本。因此，如果你训练自己的模型，请确保你的YOLOv5代码是v5.0版本。

git clone -b v5.0 https://github.com/ultralytics/yolov5.git
cd yolov5
mkdir weights
cd weights
// 下载 https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s.pt
wget https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s.pt

获取tensorrtx

对于YOLOv5 v5.0，从YOLOv5 release v5.0下载.pt文件，执行git clone -b v5.0 https://github.com/ultralytics/yolov5.git和git clone -b yolov5-v5.0 https://github.com/wang-xinyu/tensorrtx.git，然后按照tensorrtx/yolov5-v5.0中的说明运行。

获取xxx.wts模型

cp {tensorrtx}/yolov5/gen_wts.py {ultralytics}/yolov5/
cd yolov5 
python3 gen_wts.py -w ./weights/yolov5s.pt -o ./weights/yolov5s.wts
// 将生成一个'yolov5s.wts'文件

你可以在yolov5/weights/目录下找到yolov5s.wts模型。

构建tensorrtx/yolov5并获取TensorRT引擎

cd tensorrtx/yolov5
// 如果你的模型是在自定义数据集上训练的,请在 yololayer.h 中更新 CLASS_NUM
mkdir build
cd build
cp {ultralytics}/yolov5/yolov5s.wts {tensorrtx}/yolov5/build
cmake ..
make
// yolov5s
sudo ./yolov5 -s yolov5s.wts yolov5s.engine s
// 测试你的引擎文件
sudo ./yolov5 -d yolov5s.engine ../samples

然后你就得到了 yolov5s.engine,你可以将 yolov5s.engine 放入我的项目中。例如

cd {yolov5-deepsort-tensorrt}
mkdir resources
cp {tensorrtx}/yolov5/build/yolov5s.engine {yolov5-deepsort-tensorrt}/resources

获取 deepsort 引擎文件你可以从这个drive 链接获取 deepsort 预训练模型 ckpt.t7 就可以了。

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
// 测试
./demo ../resource/deepsort.engine ../resources/track.txt

完成所有5个步骤后,你就可以得到 yolov5s.engine 和 deepsort.engine。

在获取 yolov5s.engine 和 deepsort.engine 的过程中你可能会遇到一些问题,你可以在 github 上提问或查看这篇 csdn 文章。

<details> <summary>不同版本的 yolov5</summary>

目前,tensorrt 支持 yolov5 v1.0(仅 yolov5s)、v2.0、v3.0、v3.1、v4.0 和 v5.0。

对于 yolov5 v5.0,从 yolov5 release v5.0 下载 .pt 文件,git clone -b v5.0 https://github.com/ultralytics/yolov5.git 和 git clone https://github.com/wang-xinyu/tensorrtx.git,然后按照当前页面的使用说明进行操作。
对于 yolov5 v4.0,从 yolov5 release v4.0 下载 .pt 文件,git clone -b v4.0 https://github.com/ultralytics/yolov5.git 和 git clone -b yolov5-v4.0 https://github.com/wang-xinyu/tensorrtx.git,然后按照 tensorrtx/yolov5-v4.0 中的使用说明进行操作。
对于 yolov5 v3.1,从 yolov5 release v3.1 下载 .pt 文件,git clone -b v3.1 https://github.com/ultralytics/yolov5.git 和 git clone -b yolov5-v3.1 https://github.com/wang-xinyu/tensorrtx.git,然后按照 tensorrtx/yolov5-v3.1 中的使用说明进行操作。
对于 yolov5 v3.0,从 yolov5 release v3.0 下载 .pt 文件,git clone -b v3.0 https://github.com/ultralytics/yolov5.git 和 git clone -b yolov5-v3.0 https://github.com/wang-xinyu/tensorrtx.git,然后按照 tensorrtx/yolov5-v3.0 中的使用说明进行操作。
对于 yolov5 v2.0,从 yolov5 release v2.0 下载 .pt 文件,git clone -b v2.0 https://github.com/ultralytics/yolov5.git 和 git clone -b yolov5-v2.0 https://github.com/wang-xinyu/tensorrtx.git,然后按照 tensorrtx/yolov5-v2.0 中的使用说明进行操作。
对于 yolov5 v1.0,从 yolov5 release v1.0 下载 .pt 文件,git clone -b v1.0 https://github.com/ultralytics/yolov5.git 和 git clone -b yolov5-v1.0 https://github.com/wang-xinyu/tensorrtx.git,然后按照 tensorrtx/yolov5-v1.0 中的使用说明进行操作。

</details> <details> <summary>配置</summary>

从命令行参数中选择模型 s/m/l/x/s6/m6/l6/x6。
输入形状在 yololayer.h 中定义。
类别数量在 yololayer.h 中定义,如果使用自己的模型,不要忘记调整这个参数。
INT8/FP16/FP32 可以通过 yolov5.cpp 中的宏来选择,INT8 需要更多步骤,请先按照"如何运行"进行操作,然后再参考下面的"INT8 量化"。
GPU id 可以在 yolov5.cpp 中的宏中选择。
NMS 阈值在 yolov5.cpp 中。
边界框置信度阈值在 yolov5.cpp 中。
Batch size 在 yolov5.cpp 中。

</details>

运行自定义模型

你可能需要训练自己的模型并将训练好的模型转换为 TensorRT。你可以按照以下步骤进行操作。

训练自定义模型你可以按照官方 wiki在你的数据集上训练自己的模型。例如,我选择 yolov5-s 来训练我的模型。
转换自定义模型就像 TensorRT 官方指南中说的那样。当你按照"生成 yolov5 模型"获得 yolov5 和 tensorrt 仓库后,下一步是将你的 PyTorch 模型转换为 TensorRT。在此之前,你需要将 yololayer.h 文件的第 20、21 和 22 行(CLASS_NUM、INPUT_H、INPUT_W)更改为你自己的参数。

// 之前 
static constexpr int CLASS_NUM = 80; // 20
static constexpr int INPUT_H = 640;  // 21  yolov5 的输入高度和宽度必须能被 32 整除。
static constexpr int INPUT_W = 640; // 22

// 之后
// 如果你的模型是 2 分类且图像尺寸为 416*416
static constexpr int CLASS_NUM = 2; // 20
static constexpr int INPUT_H = 416;  // 21  yolov5 的输入高度和宽度必须能被 32 整除。
static constexpr int INPUT_W = 416; // 22

cd {tensorrtx}/yolov5/
// 如果你的模型是在自定义数据集上训练的,请在 yololayer.h 中更新 CLASS_NUM
mkdir build
cd build
cp {ultralytics}/yolov5/yolov5s.wts {tensorrtx}/yolov5/build
cmake ..
make
sudo ./yolov5 -s [.wts] [.engine] [s/m/l/x/s6/m6/l6/x6 或 c/c6 gd gw]  // 将模型序列化为计划文件
sudo ./yolov5 -d [.engine] [图像文件夹]  // 反序列化并运行推理，[图像文件夹]中的图像将被处理。
// 以yolov5s为例
sudo ./yolov5 -s yolov5s.wts yolov5s.engine s
sudo ./yolov5 -d yolov5s.engine ../samples
// 以自定义模型为例，yolov5.yaml中depth_multiple=0.17，width_multiple=0.25
sudo ./yolov5 -s yolov5_custom.wts yolov5.engine c 0.17 0.25
sudo ./yolov5 -d yolov5.engine ../samples

通过这种方式，你可以获得自己的tensorrt yolov5模型。尽情享用吧！

其他项目
- yolov5-deepsort-tensorrt
- 面部情绪识别
- yolov5火焰烟雾检测（Python版）
- yolov5头盔检测
- yolov5火焰烟雾检测
- 网页版视线估计
- yolov5头盔检测（Python版）
- yolov5反光衣检测（Python版）
- yolov5反光衣检测
- yolov5吸烟检测
- yolov5口罩检测
- yolov5口罩检测（Python版）
- yolov5烟雾检测（Python版）
- deepsort-tensorrt