noahmr / zed-yolov5

Object detection with YOLOv5, TensorRT and Stereolabs ZED

ZED-YOLOv5

This project shows how to perform YOLOv5 object detection with TensorRT and the Stereolabs ZED stereo depth camera. The code uses the yolov5-tensorrt library, and is based on one of the examples from this same library. The demo picture above shows its use at VDL RobotSports, an industrial team based in the Netherlands participating in the RoboCup Middle Size League.

Setup

Platforms

• Modern Linux distros
• NVIDIA L4T (Jetson platform)

Dependencies

• yolov5-tensorrt
• OpenCV
• ZED SDK (tested with version 3.6.1)

The software can be compiled using CMake and a modern C++ compiler (e.g. GCC) with support for C++14, using the following steps:

mkdir build
cd build
cmake ..
make

This will build the application detect.

Usage

1. Export your YOLOv5 model to ONNX using the official YOLOv5 guide.

2. Build the TensorRT engine using the build_engine tool from the yolov5-tensorrt library:

    build_engine --model YOLOV5_MODEL.onnx --output yolov5.engine
   

   where you replace YOLOV5_MODEL.onnx with the path to the ONNX file generated in step 1.

3. Run the detect application from this project:

    ./build/detect --engine yolov5.engine --classes coco_names.txt --gui
   

   this will load the TensorRT engine, open the ZED device, and display the results in a dialog (based on OpenCV highgui). To close the dialog again, press the 'e' key. If you are using a model trained on a custom dataset, you can change "coco_names.txt" to your own list of class names.

Notes

• It appears that since this pull request, the export.py script in the official YOLOv5 repository can now also directly export to a TensorRT engine, so you could try that instead of steps 1 and 2. However, this has not been tested.

• If you are running the application with an SVO, you might need to change the resolution in the code so that it matches the resolution at which the SVO was recorded. This is the line

    initParameters.camera_resolution = sl::RESOLUTION::HD720;
  

  in detect.cpp

• Currently, only a C++ implementation is available. In the future, a Python implementation will be added as well

Command-line options

Basic usage:

./detect --engine ENGINE_PATH

Arguments:

• -h --help: shows the help menu
• --engine: [Mandatory] path to the TensorRT engine
• --classes: [Optional] path to the file containing the class names
• --gui: [Optional] Display the results using a GUI (requires OpenCV highgui)
• --svo: [Optional] path to a ZED SVO file (to be used instead of a live sensor)

Additional Resources

• AI at RobotSports (Kaggle)

About

This project, as well as the yolov5-tensorrt library, are actively developed at VDL RobotSports, an industrial team based in the Netherlands participating in the RoboCup Middle Size League.

License

This software is licenced under the MIT license, see LICENCE.md.