CODEWITHSUNDEEP
yoloyolov8
anbvy
anbvy

Reputation: 39

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?

Upvotes: 0

Views: 641

Answers (1)

hanna_liavoshka
hanna_liavoshka

Reputation: 1323

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

Upvotes: 1

Related Questions