Load and Inference local YOLOv8.pt with PyTorch

The YOLOv8 model, distributed under the GNU GPL3 license, is a popular object detection model known for its runtime efficiency and detection accuracy.

YOLOv8 offers unparalleled performance in terms of speed and accuracy. Its streamlined design makes it suitable for various applications and easily adaptable to different hardware platforms, from edge devices to cloud APIs.

Many readers were also interested in learning “how to load the YOLOv8 model using PyTorch”.To address this issue and with the recent release of the YOLOv8 model from Ultralytics, we present this post on how to Load and Predict a YOLOv8 PyTorch Model.

Install

Install YOLOv8 via the ultralytics pip package for the latest stable release.

!pip install -U ultralytics

Note that pip automatically installs all required dependencies.

Download YOLO model

Ultralytics has multiple YOLOv8 models with different capabilities. They are subdivided into Object Detection, Segmentation, and Classification.

Model size
(pixels)
mAPval
50-95
Speed
CPU ONNX
(ms)
Speed
A100 TensorRT
(ms)
params
(M)
FLOPs
(B)
YOLOv8n 640 37.3 80.4 0.99 3.2 8.7
YOLOv8s 640 44.9 128.4 1.20 11.2 28.6
YOLOv8m 640 50.2 234.7 1.83 25.9 78.9
YOLOv8l 640 52.9 375.2 2.39 43.7 165.2
YOLOv8x 640 53.9 479.1 3.53 68.2 257.8

In this blog, we focus on object detection using yolov8l.pt PyTorch model and load YOLOv8 model and inference. The model weights yolov8l.pt file must be in local directory and the main inference python script contains the functions needed for loading the model, parsing the input, running the inference, and post-processing the output.

Use with PyTorch

YOLOv8’s Python interface allows for seamless integration into your Python projects, making it easy to load, run, and process the model’s output. Python interface enables users to quickly implement object detection, segmentation, and classification in their projects. This makes YOLOv8’s Python interface an invaluable tool for anyone looking to incorporate these functionalities into their Python projects.

import torch
from PIL import Image
from google.colab.patches import cv2_imshow
from ultralytics import YOLO

model = YOLO("/content/yolov8l.pt")

Predict YOLOv8

YOLOv8 predict mode can generate predictions for various tasks, returning either a list of Results objects or a memory-efficient generator of Results objects when using the streaming mode. 

im1 = Image.open('/content/office_desk.jpg')
im2 = Image.open('/content/bus_station.jpg')

YOLOv8 can accept various input sources like,images, URLs, PIL images, OpenCV, numpy arrays, torch tensors, CSV files, videos, directories, globs, YouTube videos, and streams. 

# Inference
results = model([im1, im2]) # batch of images

The Results object contains the following components:

  • Results.boxes: Boxes object with properties and methods for manipulating bounding boxes
  • Results.masks: Masks object for indexing masks or getting segment coordinates
  • Results.probs: torch.Tensor containing class probabilities or logits
  • Results.orig_img: Original image loaded in memory
  • Results.path: Path containing the path to the input image

Each result is composed of a torch.Tensor by default, which allows for easy manipulation.

Plotting results

You can use plot() function of Result object to plot results on in image object. It plots all components(boxes, masks, classification logits, etc.) found in the results object.

res_plotted = results[1].plot()
cv2_imshow(res_plotted)
PyTorch YOLOv8 Predict

Related Post

Convert PASCAL dataset to TFRecord for object detection in TensorFlow

Read more here: Source link