YOLOv7 with TensorRT on Jetson Nano (with Python script example)

Result of object detection with Nvidia Jetson Nano, YOLOv7, and TensorRT.

At the end of 2022, I started working on a project where the goal was to count cars and pedestrians. Because of privacy issues and lower data consumption, we needed to use edge devices. Our budget was tight, so we chose Nvidia Jetson Nano to run our object detection model. This is how I found out Nvidia doesn’t really care about their cheapest SoC. Documentation is lacking, software is outdated and there aren’t many useful tutorials. I hope you will finally get some useful information from this one. I will try to keep it as short as possible. No ChatGPT bullshit.

Object detection model

I will use a YOLOv7 object detection model in this tutorial. I chose it because it was a state-of-the-art object detection model at the time of developing my project in late 2022. Now you can use any other YOLO model (older or newer) as long as it’s supported by libraries we will use for conversion from PyTorch to the TensorRT model.

Jetson Nano Setup (non-optimized YOLOv7 deployment)

I already have a tutorial on how to prepare your Jetson Nano for YOLOv7 deployment. It’s a good starting point because it goes into detail on how to install all required libraries and deal with Python virtual environment on Jetson Nano. In the end, you’ll be able to run the YOLOv7 object detection model. The result is around 17 FPS (YOLOv7 Tiny with the input of 416x416) and 9 FPS (YOLOv7 Tiny with the input of 640x640). And this is just inference speed, without pre- and post-processing of the video. It’s slow. That’s why we are here.

Optimization using TensorRT

Let’s begin the fun part. Go to home directory and clone TensorRTx repository.

cd
git clone https://github.com/wang-xinyu/tensorrtx

TensorRTx is used to convert your PyTorch model to TensorRT engine model.

If you don’t have your custom weights, you can use regular YOLOv7 tiny weights from here. Weights should be in your cloned official YOLOv7 repository.

Go to cloned tensorrtx repository and copy script for generating .wts file from tensorrtx to yolov7 directory:

cd tensorrtx/yolov7
cp cp gen_wts.py {your-path-to-yolov7-directory}/yolov7
# Move to yolov7 directory
cd {your-path-to-yolov7-directory}/yolov7

Make sure you have your virtualenv activated to have access to the installed PyTorch library.

Convert .pt file to .wts file. This will take around 30 seconds.

python gen_wts.py -w yolov7-tiny.pt 

You should get new yolov7-tiny.wts file in your current directory.

Build tensorrtx. Go to tensorrtx directory:

cd
cd tensorrtx/yolov7

Go to /include directory and open config.h. Find variable kNumClass and check if a number of classes matches your model’s number of classes (change if you have the custom trained model). Building this will create a shared object library plugin file (.so) which is needed for inference. After that, you convert the weights file to a .engine file.

mkdir build
cd build
cp {your-path-to-yolov7-directory}/yolov7/yolov7.wts {your-path to-tensorrtx-directory}/yolov7/build
cmake ..
make

# Serialize model to engine file
sudo ./yolov7 -s [.wts] [.engine] [t/v7/x/w6/e6/d6/e6e gd gw]
# Example:
sudo ./yolov7 -s yolov7-tiny.wts yolov7-tiny.engine t

# Deserialize and run inference, the images in [image folder] will be processed.
sudo ./yolov7 -d [.engine] [image folder]
# Example:
sudo ./yolov7 -d yolov7-tiny.engine ../images

Check images that were generated (zidane.jpg and bus.jpg).

Now install pycuda and python-tensorrt.

Execute this command:

ln -s /usr/include/locale.h /usr/include/xlocale.h

Install pycuda with some added options:

pip3 install --global-option=build_ext --global-option="-I/usr/local/cuda/include" --global-option="-L/usr/local/cuda/lib64" pycuda

Open .bashrc file and add this to the end:

sudo nano ~/.bashrc

# Add this line to the end
export PYTHONPATH=/usr/lib/python3.6/dist-packages:$PYTHONPATH
# Save and exit
source ~/.bashrc

If you still have problems installing pycuda and tensorrt, check out this tutorial.

Run inference with YOLOv7 and TensorRT

I prepared a Python script to test this yolov7 and tensorrt. Actually, I found it somewhere on the internet and modified it based on my needs. You can find it in my repository. It’s called yolov7_trt_cam.py.

You can edit input and engine file inside Python code. Make sure you use the correct paths!

Inside

if __name__ == "__main__":

Find this piece of code and change your plugin library path and engine file path based on your needs.

# CHANGE THIS!!!
PLUGIN_LIBRARY = "/path-to-plugin-library/libmyplugins.so"
engine_file_path = "/path-to-engine-file/yolov7-tiny.engine"

If you want to setup source for video find:

class inferThread(threading.Thread):
    def __init__(self, yolov7_wrapper):
        threading.Thread.__init__(self)
        self.yolov7_wrapper = yolov7_wrapper
        
        # CHANGE THIS LOCATION OR USE YOUR CAMERA!!!
        self.cap = cv2.VideoCapture(
            "/opt/nvidia/deepstream/deepstream/samples/streams/sample_1080p_h264.mp4"
        )

Run yolov7_trt_cam.py and look at the results. Congrats!

NEED HELP WITH YOUR PROJECT?

Best way for any business-related stuff is to send me a DM through Linkedin. Ping me and we will see how we can develop your project!