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. ๐ธ
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.
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. ๐
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:
Cool! ๐ Next we will do 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.
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