Export your custom YOLOv7 PyTorch model to ONNX

 

Another post starts with you beautiful people💓. In my previous post, we learned about training a custom dataset with an official Pytorch-based YOLOv7 object detector. If you have not seen that post, I recommend you check it once. The link is here. Once we achieve the best model, the next important step is to use that model. Sometimes you may need to use multiple ML models which you have trained on different-different ML frameworks like PyTorch, TensorFlow, Caffe, etc. In production or the real world, the trained model can be deployed as Rest API, or integrated with a web application without changing its form but what if you need to use that within your mobile device as an Android or iOS app as well as you want to use it with an embedded system like Nvidia Jetson 💣? Here comes the problem of interoperability. In this post, we are going to learn how can we export our custom YOLOv7 model to ONNX format.

ONNX is an open format built to represent machine learning models. ONNX defines a common set of operators - the building blocks of machine learning and deep learning models - and a common file format to enable AI developers to use models with a variety of frameworks, tools, runtimes, and compilers. For our export, I will use the same Google Colab platform where I had trained my custom yolov7 model. Let's first import some required dependencies-

Since Google Colab has already PyTorch installed so we don't need to install it. During this work, my Colab had the following version of Python and PyTorch-

Next, we will download the official code repo of YOLOv7 as below-

Once the code base is downloaded, let's check a test image with our trained YOLOv7 model so that in the last of this post we can compare it with the exported ONNX model inference. The command to do the inference is as below-

! python detect.py --weights best.pt --img-size 1280 --source test_image.jpg

Here, just replace best.pt with your custom YOLOv7 weight file and  test_image.jpg with your test image. After successfully running of above command the resulting image will be saved in the path- runs/detect/exp/test_image.jpg. Let's check the outcome by opening the image below-

And it looks like as below-

So till now, everything looks fine. Let's move to our export. For this purpose, the official code repo that we had already downloaded has a script named export.py.

We will use the same to export the model into ONNX format as below-

You can check export.py to understand it's each argument. After successfully running of above command you will see the following output in the console-

Now, the next most important step is to do the inference with this exported model. Let's understand this with below code snippets one by one-

In the above code snippet, first, we have loaded the ONNX model, our test image, required libraries for inference, CPU or GPU selection, and the main class of the ONNX runtime as a session. In the next code snippet, we will create a function to do the following things with the given input image-

1. Resize and pad image while meeting stride-multiple constraints
2. Scale ratio
3. Compute padding
4. divide padding into 2 sides
5. Resize and find the top, bottom, left, and right coordinates
6. Add border and return the image with the scale ratio and sides

Next, we will create a list of our target object names- in my case, it is wheat, color scheme for the bounding boxes. Then we will apply the above-created function- letterbox() to the image, will convert the image to a NumPy array and then we will normalize it for further processing. Next, we will extract the output from the model as below-

Now, we are ready for making the inference. First, we will fetch the output of the ONNX inference session and then we will iterate it to get the image, bounding box coordinates, confidence score, and the class name, and then we will use this information to draw the bounding boxes around the detected class with its name as below-

  That's it! Our inference script is ready for our ONNX file and once executed it had shown me the output image as below-

If you compare this image with the previous output image you will find although the confidence score of the detected class has been dropped yet the accuracy of the exported model is the same. It means our exported ONNX model is working as expected. One more point to observe in our inference script is that we have not imported the PyTorch library which we used to train the custom dataset with YOLOv7. That is the interoperability of the ONNX format. Now our production deployment team can choose any ML framework they are comfortable with and using this unified inference framework (ONNX), the exported model can be deployed to anywhere. You can find the complete Colab Notebook here.

So why wait now☝ Just copy the notebook, export your custom YOLOv7 model to ONNX format and deploy it to production💥. That's it for today guys. In my next post, I will share to export the model into TensorRT format till then👉 Go chase your dreams, have an awesome day, make every second count, and see you later in my next post😇