SuperGradients How To Use a model to Predict.ipynb does not work for yolo_nas models

[mikel-brostrom]

🐛 Describe the bug

This:

import torchvision.transforms as transforms
from PIL import Image
from super_gradients.training import models

from super_gradients.training.models.detection_models.yolo_base import YoloPostPredictionCallback
import torch

model = models.get("yolo_nas_s", pretrained_weights="coco", num_classes=80)

preprocess = transforms.Compose([
    transforms.Resize([640, 640]),
    transforms.PILToTensor()
])

im = Image.open('/path/to/im.jpg')
transformed_image = preprocess(im).float().unsqueeze(0)

with torch.no_grad():
    raw_predictions = model(transformed_image)

# Callback uses NMS, confidence threshold
predictions = YoloPostPredictionCallback(conf=0.1, iou=0.4)(raw_predictions)
print(predictions)

Gives None

Versions

3.1.3

[BloodAxe]

You should be using PPYoloEPostPredictionCallback and not YoloPostPredictionCallback for yolo nas model. Or as an alternative - use the predict() method. See https://github.com/Deci-AI/super-gradients/blob/ac21534a5505da1e2758acddd7dead0e7ac19dad/YOLONAS.md

[mikel-brostrom]

Sorry, I checked https://github.com/Deci-AI/super-gradients/blob/ac21534a5505da1e2758acddd7dead0e7ac19dad/YOLONAS.md but I couldn't find this information. Thanks anyways :smile: