基于yolov5+fastreid+deepsort的TensorRT目标跟踪

Object Tracking with TensorRT

Introduction

• 代码位置：https://github.com/linghu8812/yolov5_fastreid_deepsort_tensorrt
• 视频结果：https://www.bilibili.com/video/BV1qg411K74p?spm_id_from=333.999.0.0&vd_source=421ea4aa0115ac996e20dc1d2cda21d4

This is an implementation for object tracking in cplusplus code. The object detector uses yolov5s model. The idea of deepsort is adopted in object tracking. The implementation of sort refers to the sort-cpp. When extracting object features, it is extracted through the fast-reid trained model, and the person ReID uses mobilenetv2. The purpose of using these lightweight models is to ensure the real-time efficiency of video processing. The model inference base on TensorRT engine.

How to run?

0. Build environments

The TenosrRT environments build from Dockerfile, run with the following command.

docker build -t tensorrt_tracker:0.1.0_rc .

Following yolov5 and fast-reid requirements file to install their depends packages.

1. Transform PyTorch weights to ONNX

• Transform yolov5 weights

Use this yolov5 repo to transform yolov5 *.pt weights to ONNX models. Run the following command

git clone https://github.com/linghu8812/yolov5.git
python3 export.py ---weights weights/yolov5s.pt --batch-size 1 --imgsz 640 --include onnx --simplify

A pretrained yolov5 ONNX detection model can be downloaded form here, link: https://pan.baidu.com/s/1RUz7Xk78lvKCeZNk_BBvoQ, code: jung. download this model and put it to the weights folder.

• Transform fastreid weights

Use official fast-reid to transform PyTorch weights to ONNX model. Run the following command

https://github.com/JDAI-CV/fast-reid.git
python3 tools/deploy/onnx_export.py --config-file configs/Market1501/mgn_R50-ibn.yml --name mgn_R50-ibn --output outputs/onnx_model --batch-size 32 --opts MODEL.WEIGHTS market_mgn_R50-ibn.pth

A pretrained fast-reid ONNX detection model can be downloaded form here, link: https://pan.baidu.com/s/19TuHxxuVYLBzie5_Vu0cCQ, code: 1e35. download this model and put it to the weights folder.

2. Get video for inference ready

Put video file for inference to samples folder. Here is a video demo for inference can be used: https://pan.baidu.com/s/1Yyh1lwmzNl_gjvNz9EVI5w, code: fpi0.

3. Build project

Run the following command

git clone git@github.com:linghu8812/tensorrt_tracker.git
mkdir build && cd build
cmake ..
make -j

4. Run project

Run the following command

./object_tracker ../configs/config.yaml ../samples/test.mpg

results demo:

Reference

• yolov5: https://github.com/ultralytics/yolov5
• FastReID: A Pytorch Toolbox for General Instance Re-identification: https://arxiv.org/abs/2006.02631
• fast-reid: https://github.com/JDAI-CV/fast-reid
• Simple Online and Realtime Tracking: https://arxiv.org/abs/1602.00763
• sort-cpp: https://github.com/mcximing/sort-cpp
• Simple Online and Realtime Tracking with a Deep Association Metric: https://arxiv.org/abs/1703.07402
• tensorrt_inference: https://github.com/linghu8812/tensorrt_inference