YOLO Object Classification with Python

Object classification is a computer vision task that involves classifying an entire image into one of a set of predefined classes. The output of an image classifier is a single class label or multi-label along with confidence score. Image classification is useful when you need to know only what class an image belongs to and don't need to know where objects of that class are located or what their exact shape is.

YOLOv8 Segmentation models have -cls suffix and are pretrained on ImageNet dataset containing 1000 classes

YOLOv8 pretrained Classification models (nano, small, medium, large and extra large based on number of parameters) are shown in the table below:

Setup UltraLytics for YOLOv8

%pip install ultralytics
import ultralytics
ultralytics.checks()

Load YOLOv8 for Classification

from ultralytics import YOLO

# Load a model
# You can use different YOLOv8 variants (yolov8n-cls, yolov8s-cls, yolov8m-cls, yolov8l-cls, yolov8nx-cls)
model = YOLO('yolov8n-cls.pt')  # load a pretrained model

# Use the model
results = model('https://ultralytics.com/images/bus.jpg')  # predict on an image
# Save the output image after classification in Google Colab 
results[0].save('/content/output_cls.jpg')

# Print the COCO dataset classes on which model is trained.
print(model.names.values())

Input Image

Display the Output Image after Classification

from IPython.display import Image

# Display the image in Google Colab 
Image(filename='/content/output_cls.jpg')

Output Predicted Image

Print Objects Classes Probabilities

for r in results:
    print(r.probs)