DEV Community

Cover image for Object Detection Using Python And YOLO-NAS
Ethan
Ethan

Posted on • Originally published at ethan-dev.com

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
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

To install the dependencies, run the following command:

pip install -r requirements.txt
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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()
Enter fullscreen mode Exit fullscreen mode

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"])
Enter fullscreen mode Exit fullscreen mode

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]
Enter fullscreen mode Exit fullscreen mode

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"])
Enter fullscreen mode Exit fullscreen mode

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

This can run via:

python video.py -v [path to video file]
Enter fullscreen mode Exit fullscreen mode

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.

โ€œBuy Me A Coffeeโ€

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

Top comments (0)