TDM 19000: Introduction to YOLO

YOLO stands for "You Only Look Once", as this algorithm detects using only a single regression problem. More traditional detectors are multi-stage, and take two steps to propose the regions and then classify them. YOLO does this in a single neural network to predict the bounding boxes AND to classify its certainty, all within one forward pass.

For this project, we will be using a pretrained model - YOLOv8 nano - along with a selected image from a mall camera footage video. To load these in, read in the image and load the model:

from PIL import Image
from ultralytics import YOLO

yolo_model = YOLO("yolov8n.pt")

image_path = "/anvil/projects/tdm/data/images/seq_000014.jpg"

Image.open(image_path)

It’s actually simple to use the model to predict the contents of the image. Use “yolo_results = yolo_model.predict(source=image_path)”, and show the results. There should be a count of how many people (and plants) there were found to be. Did the model correctly predict all relevant objects in the image?

We can use this model further to show us what parts of the image it thinks are people, and how certain of this it is. The confidence score will be calculated for each box in the YOLO results - the coordinates for each box will get stored based on the x-min, y-min, x-max, and y-max values. The given class ID that the model assigned will need to be converted to a human-readable format (e.g. "person"). To store the results, save all bounding boxes, get the total count of how many of each object type there are, and track the labels given to the objects.

input_boxes = []
counts = {}
detected_names = []

for result in yolo_results:
    for box in result.boxes:
        # Get coordinates in [xmin, ymin, xmax, ymax] format
        x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()

        # Prediction confidence score
        conf = box.conf[0]

        # Get class ID and assign readable name (e.g., 'person')
        cls = int(box.cls[0])
        name = result.names[cls]

        # use ':.2f' to shorten the shown decimals of the confidence score
        print(f"Found {name} with {conf:.2f} confidence at {x1, y1, x2, y2}")

        input_boxes.append([float(x1), float(y1), float(x2), float(y2)])
        counts[name] = counts.get(name, 0) + 1
        detected_names.append(name)

Print 'counts' to show how many occurrences there were of each object type in the image.

We’ve created these "boxes" to classify the image objects, but this doesn’t really have much of a meaning to us yet. To help make this make sense, we can loop through the results and overlay them on the original image.

yolo_results is where everything is stored. For each result in yolo_results, plot it to display the box and score.

import matplotlib.pyplot as plt

for r in yolo_results:
    im_array = r.plot()
    plt.imshow(im_array)
    plt.axis('off')
    plt.show()