YOLOv5 | YOLOv7 object tracking (StrongSORT, ByteTrack, OCSORT)

Introduction

This work presents an approach for tracking people using thermal imaging cameras. For object detection was used YOLO family models (mostly YOLOv5 and YOLOv7 from YOLOU. The detections generated by the YOLO model are passed to one of three trackers: StrongSORT, OCSORT and ByteTrack. They can track any object that your YOLO model was trained to detect.

Installation

$ sudo apt-get update && sudo apt upgrade
$ sudo apt install ffmpeg

$ python3 -m pip install -e .

I also suggest to use this NVIDIA docker container or just install manually PyTorch>=1.13, CUDA>=11.7, TensorRT>=8.5, cuDNN>=8.5

If you run from Python IDE, make sure that it emulates terminal in output console.

All tested YOLOs and ReID weights (PyTorch and TensorRT) you can find here

Tracking

Firstly:

$ cd ./yolos_thermal_tracking

$ python3 inference.py --yolo-weights <weights path> --source <file(s) / folder path>

Tracking methods:

$ python3 inference.py --tracking-method ocsort
                                         bytetrack
                                         strongsort --reid-weights <ReID weights path>                

Tracking sources:

Tracking can be run on most video formats which FFMPEG supports.

$ python3 inference.py --source vid.mp4  # video
                                vid1.mp4, vid2.webm  # list of videos
                                path/  # directory

YOLO models:

$ python3 inference.py --yolo-weights yolov5m.pt
                                      yolov5m.engine
                                      yolov5s.pt
                                      yolov5s.engine
                                      yolov7s.pt (tiny version)
                                      yolov7s.engine (tiny version)
                                      yolov7.pt
                                      yolov7.engine

ReID models:

For StrongSort tracking method you need to choose a ReID model based from this ReID model zoo. These model can be further optimized for you needs by the reid_export.py script.

$ python3 inference.py --reid-weights osnet_x0_25_msmt17.pt
                                      osnet_x0_25_msmt17.engine
                                      osnet_x0_5_msmt17.pt
                                      osnet_x0_5_msmt17.engine

Filter tracked classes:

If you want to track only specific classes of objects, add their index after the --classes flag.

$ python3 inference.py --classes 0  # Only Pedestrians

Here is a list of all the possible objects that a YOLO model trained on MS COCO can detect.

Save results:

You can save results as video(s) in the specific folder

$ python3 inference.py --save-folder <save_path> 

Acknowledgements

Expand

• https://github.com/mikel-brostrom/Yolov5_StrongSORT_OSNet (Trackers implementation)
• https://github.com/jizhishutong/YOLOU (YOLO models implementation)
• https://github.com/WongKinYiu/yolov7
• https://github.com/ultralytics/yolov5
• https://github.com/ppogg/YOLOv5-Lite
• https://github.com/ceccocats/tkDNN
• https://github.com/pjreddie/darknet
• https://github.com/nebuly-ai/nebullvm
• https://github.com/luanshiyinyang/awesome-multiple-object-tracking