Detect particular class

[Cutinho007]

I have a yolov5 custom trained model, trained with 32 different classes. I tried using SAHI+ yolov5 and it works perfectly.
Well I am looking for a way to detect only 2or 3 particular classes. Is there a way I can do so?

[nakuda1]

Hi, did you run your own model and get .pt? Or just using the pt document sahi provide.

[Cutinho007]

I trained my own model and have my own best.pt.

[nakuda1]

Thanks a lot!!! I also got my best.pt, but when I use my .pt, the program have the error. Can you tell me how to figure it out?

---

KeyError Traceback (most recent call last)
File ~/miniconda3/lib/python3.8/site-packages/sahi/model.py:370, in Yolov5DetectionModel.load_model(self)
TypeError: ('model_path is not a valid yolov5 model path: ', KeyError('model'))

[nakuda1]

About your question, I also make two class objections. And I modify the classes before my training, maybe it works.

[fcakyon]

@Cutinho007 this should work:

import yolov5
from sahi import AutoDetectionModel

model = yolov5.load('custom_model.pt")
model.classes = [0, 15, 16]

sahi_model = AutoDetectionModel(
    model_type='yolov5',
    model=model,
)

# use `sahi_model` with `get_prediction` or `get_sliced_prediction`