Search code examples
yoloyolov8

How to limit default "yolo" model to recognize less objects?

i want to use default YOLOv8 model (yolov8m.pt) for object detection. I know that default YOLO models uses COCO dataset and can detect 80+ objects. I just want to detect like 5 of them, how can i achieve this?

Solution

  • To filter the classes in the prediction mode use the classes parameter of the predict function, where you need to specify the relevant classes IDs.

    classes, list[int]. Filters predictions to a set of class IDs. Only detections belonging to the specified classes will be returned. Useful for focusing on relevant objects in multi-class detection tasks. https://docs.ultralytics.com/modes/predict/#inference-arguments

    The class list for the COCO dataset can be found here: https://github.com/ultralytics/ultralytics/blob/main/ultralytics/cfg/datasets/coco.yaml.

    from ultralytics import YOLO

model = YOLO('yolov8n.pt')
results = model.predict(source='img.jpg', classes=[0, 3, 5])