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

Asked By: TAN

||

Source

Answer 1

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)
for i, (result) in enumerate(results):
    print('Do something with class 0')

Answered By: Mike B

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

Question:

Answers: