Search code examples
pythoncomputer-visiondetectionyolo

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

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
Solution

  • 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')