Exclude classes

[st43r]

Hey, could someone please help me

i want to try make object detection on pretrained YoloNasL model but i need to exclude some classes for example class person

///
modelName='yolo_nas_l'

model = models.get(modelName, pretrained_weights="coco").to(DEVICE)

image = cv2.imread(SOURCE_IMAGE_PATH)
result = list(model.predict(image, conf=0.35))[0]
///

how can i set some properties of model like class_names and as example give array with class names to that property?

i checked all properties but didnt find out where or how i can exclude some yolo classes

i would be extremely grateful if you try to help

thanks in advance!!

*environment is google colab

[SkalskiP]

Hi, @st43r 👋🏻! You can use supervision to do that. UNWANTED_CLASSES is the list of class ids that you want to filter out.

import cv2
import torch
import numpy as np
import supervision as sv
from super_gradients.training import models

DEVICE = 'cuda' if torch.cuda.is_available() else "cpu"
MODEL_ARCH = 'yolo_nas_l'
UNWANTED_CLASSES = [1, 4, 8]

model = models.get(MODEL_ARCH, pretrained_weights="coco").to(DEVICE)

image = cv2.imread(SOURCE_IMAGE_PATH)

result = list(model.predict(image, conf=0.35))[0]
detections = sv.Detections.from_yolo_nas(result)
selection_mask = ~np.isin(detections.class_id, UNWANTED_CLASSES)
detections = detections[selection_mask]

[st43r]

Piotr, you are the bestest man!

Thanks a lot!

Have a great one!

Sincerely,
Sergey Protasov