How to count the number of objects in each class in an image which has the different objects detected using YOLO NAS model?

Question

I have, say, a fruits dataset and I have used YOLO NAS model to determine/detect the different types of fruits in an image. But now I want to count how many fruits are there in each category and display that on the screen. How to go about this?

I have tried searching for information as much as possible but I couldn't get something to do with YOLO NAS and their documentation is not giving me any info on the same. Any help is appreciated. Atleast pointing me to an article which explains the same is also helpful. It is difficult to get something for an image and not a video.

Answer 1

from super_gradients.training import models
import cv2
from collections import Counter

yolo_nas_l = models.get("yolo_nas_l", pretrained_weights="coco")


image = cv2.imread("image.jpg")
result = list(yolo_nas_l.predict(image, conf=0.35))[0]

class_names = result.class_names

counter = Counter(list(result.prediction.labels.astype("int")))
output = list(counter.items())
output = list(map(lambda x: (class_names[x[0]],x[1]),output))

print(output) # [('person', 8), ('bottle', 6), ('bowl', 2), ('dining table', 1), ('chair', 1), ('couch', 1)]