Skip to main content

7-9 Image Classification

PANDORA DEVELOPER GUIDE · CHAPTER 07

Learning Objectives​

Image Classification is one of the most fundamental and common tasks in computer vision. Its goal is to enable a computer to determine "which category an image belongs to."

Given an input image, the output includes:

- The category the image belongs to

- Confidence scores for each category

We will use Python and a pre-trained model to perform image classification.​

Perform image classification on a picture and print the Top-5 results.

import torch
from ultralytics import YOLO
model = YOLO("yolo11n-cls.pt")
results = model.predict("https://ultralytics.com/images/bus.jpg")
for result in results:
    print(f'top 5: {[result.names[cls.item()] for cls in torch.topk(result.probs.data, 5)[1].detach().cpu().numpy()]}')

You will then obtain the results, which show the top 5 most likely categories.

7-9 Image Classification

Perform real-time image classification using a webcam and visualize the results​

import cv2
from ultralytics import YOLO
model = YOLO("yolo11n-cls.pt")
video_path = 0
cap = cv2.VideoCapture(video_path)
while cap.isOpened():
    success, frame = cap.read()
    if success:
        results = model.predict(frame)
        annotated_frame = results[0].plot()
        cv2.imshow("YOLO Inference", annotated_frame)
        if cv2.waitKey(1) & 0xFF == ord("q"):
            break
    else:
        break
cap.release()
cv2.destroyAllWindows()

You will then obtain the results. The image displays the top 5 most likely categories, such as "stethoscope," along with their corresponding scores, e.g., 0.09.

7-9 Image Classification

Adapted from the original YUAN Developer Hub article. Source ID: 263.

On this page