Extracting class-probabilities when detecting in YOLOv8

[dbjorklund]

Hi! I am trying to implement an Active Learning implementation connected with YOLOv8's detection. To do this, I would need access to the class probabilities and not only the confidence of the prediction/bounding box which is given as an output now. Image embeddings would also be beneficial as this could be used as an diversifying measure in the learning step too,.

Does anyone have any ideas on how to solve this, or if any other detector include this which is comparable to the YOLO framework?

[pderrenger]

Hello! 👋 To access class probabilities in YOLOv8, you can use the probs attribute of the Results object after running predictions. Here's a quick example:

from ultralytics import YOLO

# Load your model
model = YOLO('yolov8n.pt')

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

# Access class probabilities
class_probs = results[0].probs

For image embeddings, you can extract features from specific layers using the embed argument during prediction:

# Extract embeddings from layer 24
results = model('path/to/image.jpg', embed=[24])

These embeddings can then be used in your Active Learning pipeline. If you need further assistance or have more questions, feel free to ask! Happy coding! 😊