Limit the class of my prediction?

[clodoaldocodes]

Search before asking

• I have searched the YOLOv8 issues and discussions and found no similar questions.

Question

If I want to use the YOLOv8 pre-trained network and I know that my image has oranges and my intention is to detect only these, is there a way for me to set the return to only return oranges? And any other classes that return will be automatically deleted.

Additional

No response

[glenn-jocher]

Absolutely, you can restrict YOLOv8 to detect only specific classes, such as oranges, by filtering the detection outputs based on class IDs associated with your target class(es).

Here's a simple example using the Python API:

from ultralytics import YOLO

# Load your model
model = YOLO('yolov8n.pt')  # Assuming 'yolov8n' is the model you are using

# Define the class ID for oranges (just an example ID)
orange_class_id = 32  # You'll need to replace this with the actual class ID for oranges in your model

# Run prediction
results = model('path/to/your/image.jpg')

# Filter results to only keep detections of oranges
filtered_results = [d for d in results if d['class'] == orange_class_id]

# Now 'filtered_results' will only contain detections for oranges

This snippet loads the model, predicts on an image, and then filters the results to include only detections of the class representing oranges. Make sure to replace 'path/to/your/image.jpg' with the actual path to your image and orange_class_id with the correct class ID for oranges from the model you are using.