Specify a class to detect using YOLOv8 on pre-trained model

Question

I'm new to YOLOv8, I just want the model to detect only some classes, not all the 80 classes the model trained on. How can I specify YOLOv8 model to detect only one class? For example only person.

from ultralytics import YOLO
model = YOLO('YOLOv8m.pt')

I remember we can do this with YOLOv5, but I couldn't do same with YOLOv8:

model = torch.hub.load("ultralytics/yolov5", 'custom', path='yolov5s.pt')
model.classes = [0]  # Only person
model.conf = 0.6

why not just filter the resulting detections, selecting only those classes you care about? — Christoph Rackwitz, Feb 15 '23 at 08:31
That's one method. But I think it's better to restrict the model to detect only certain class. — TAN, Feb 15 '23 at 09:55

Mike B · Accepted Answer · 2023-05-31T19:19:38.627

3

Just specify classes in predict with the class IDs you want to predict

from ultralytics.yolo.engine.model import YOLO

model = YOLO("yolov8n.pt")
model.predict(source="0", show=True, stream=True, classes=0)  # [0, 3, 5] for multiple classes
for i, (result) in enumerate(results):
    print('Do something with class 0')

edited May 31 '23 at 19:19

answered Feb 15 '23 at 14:29

Mike B

2,136
2
12
31

1

if you need to predict various classes: classes = [2,3,5,7] – SachaDee May 29 '23 at 21:52

Specify a class to detect using YOLOv8 on pre-trained model

1 Answers1