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:
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