Object Detection Using Python And YOLO-NAS

Introduction

Hello! 😃

In this tutorial I will show you how to use Python and YOLO-NAS to detect objects in both images and videos. 😎

YOLO-NAS is the faster and more accurate then it's predecessors.

I have done a tutorial about object detection before but this one is a lot more easier and beginner friendly IMO. 😸

Creating The Virtual Environment

First thing we need to do is create and activate the virtual environment, this can be done with the following commands:

python3 -m venv env
source env/bin/activate

Done! Next we need to create a requirements file and install the dependencies.

Installing The Dependencies

Create a requirements.txt file and add the following:

imutils
torchinfo
torch
super-gradients

To install the dependencies, run the following command:

pip install -r requirements.txt

Now we can finally start coding. 😄

Detecting Objects In Images

First I will show you how to detect objects in images, open a file called image.py and add the following imports:

import gc
import os
import torch
from torchinfo import summary
from super_gradients.training import models
import argparse

Next we need to create a method to detect objects in the image, the code to do this is as follows:

def detect_objects(image):
    device = torch.device("cpu")
    model = models.get("yolo_nas_s", pretrained_weights="coco").to(device)
    out = model.predict(image, conf=0.6)

    out.save("predictions");

    del model
    gc.collect()
    torch.cuda.empty_cache()

The above method takes an image patch, detects objects and then saves the output to the predictions directory. Finally we clean up. 🧹

Finally we need a main method:

if __name__ == "__main__":
    os.makedirs("predictions", exist_ok=True)

    ap = argparse.ArgumentParser()
    ap.add_argument("-i", "--image", required = True, help = "Path to input file")
    args = vars(ap.parse_args())

    detect_objects(args["image"])

The main method takes an image argument and then passes it to the detect_objects method defined earlier.

Finally the command can be run via:

python image.py -i [path to image file]

Once done you should have something similiar to the following:

Detected Objects

Cool! 😆 Next we will do videos

Detecting Objects In Videos

Detecting objects in videos is exactly the same as the above code except instead of an image file we pass a video file.

Since I've explained what the code does above I will provide the full code here:

import gc
import os
import torch
from torchinfo import summary
from super_gradients.training import models
import argparse

def detect_objects(video):
    device = torch.device("cpu")
    model = models.get("yolo_nas_s", pretrained_weights="coco").to(device)
    out = model.predict(video)          

    out.save("predictions") 

    del model
    gc.collect()
    torch.cuda.empty_cache()

if __name__ == "__main__":
    os.makedirs("predictions", exist_ok=True)

    ap = argparse.ArgumentParser()
    ap.add_argument("-v", "--video", required = True, help = "path to input video")
    args = vars(ap.parse_args())

    detect_objects(args["video"])

Notice that I changed the image argument to a video argument.

This can run via:

python video.py -v [path to video file]

The processed video should now appear in your current directory.

Conclusion

Here I have shown how to use python and YOLO to detect objects in images and video using the new YOLO NAS.

I hope this has taught you something as I had a lot of fun trying out the new YOLO. 😄

Note if you want to use GPU instead of CPU, replace cpu with cuda.

As always you can find the sample code for the above via my Github: https://github.com/ethand91/yolonas-object-detect

Happy Coding! 😎

Like me work? I post about a variety of topics, if you would like to see more please like and follow me. Also I love coffee.

If you are looking to learn Algorithm Patterns to ace the coding interview I recommend the following course