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?
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])