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上观看视频演示。

要求

1. Jetson Nano或Jetson Xavier NX
2. Jetpack 4.5.1
3. Python3（默认版本，Jetson Nano或Jetson Xavier NX默认带有TensorRT 7.1.3.0的Python3）
4. TensorRT 7.1.3.0
5. PyTorch 1.8.0
6. Torchvision 0.9.0
7. Torch2trt 0.3.0
8. ONNX 1.4.1
9. OpenCV-Python 4.5.3.56
10. Protobuf 3.17.3
11. SciPy 1.5.4

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

即将推出

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

速度

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

构建和运行

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找到。

1. 获取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

2. 获取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中的说明运行。

3. 获取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模型。

4. 构建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

5. 获取 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 文章。

运行自定义模型

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

1. 训练自定义模型 你可以按照官方 wiki在你的数据集上训练自己的模型。例如,我选择 yolov5-s 来训练我的模型。
2. 转换自定义模型 就像 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