Train YOLOv8 on a Custom Dataset: A Step-by-Step Guide

Object detection is in high demand, and YOLOv8 is leading the way, improving upon previous YOLO versions. This guide will walk you through fine-tuning YOLOv8 on your own data, leveraging the Ultralytics API and DigitalOcean GPU Droplets. By the end, you'll be able to adapt YOLOv8 for your specific needs.

Prerequisites to Get Started

Before diving in, make sure you have these basics covered:

• Python: A foundational understanding of Python.
• Deep Learning: Familiarity with neural networks and object detection.
• PyTorch/TensorFlow: Knowledge of at least one of these frameworks.
• OpenCV: Understanding of image processing.
• CUDA: Experience with GPU acceleration.
• RoboFlow: An account and basic familiarity.
• Git: For code management.

Understanding How YOLO Works

YOLO, or "You Only Look Once," revolutionized object detection by predicting both location (bounding boxes) and class probabilities in a single pass.

• Single-Stage Detection: A convolutional network predicts multiple bounding boxes and class probabilities simultaneously.
• Grid Division: An image is divided into N grids (s*s) that the model parses for object detection and localization.
• Bounding Box Prediction: The model predicts bounding box coordinates (B) with a label and confidence score for each identified object within a grid.

YOLO excels at object classification, detection, and image segmentation. This core functionality remains in YOLOv8.

What's New in YOLOv8? Major Improvements

While the official YOLOv8 paper is pending, here's what we know about the updates in YOLOv8 so far:

Key Architectural Changes

• C2f Module: Replaces the C3 module, concatenating outputs from all Bottleneck layers.
• Convolutional Layers: Replaces the first 6x6 Conv with a 3x3 Conv block in the Backbone.
• Streamlined Backbone: Deletes two Conv layers.
• Switched to using a decoupled head, and deleted the 'objectness' branch.

Enhanced Accessibility

The Ultralytics API provides a simplified approach to training and inference, no longer requiring manual repository cloning.

Anchor-Free Bounding Boxes

YOLOv8 eliminates the need for manual anchor box definition. It automatically predicts the center of objects that are detected.

Smarter Augmentation Strategies

Mosaic augmentation, combining four images, now stops before the final training epochs for optimized accuracy.

Performance Boost

YOLOv8 hits accuracy and efficiency improvements during training and inference. Data shows YOLOv8 outperforms YOLOv7, YOLOv6-2.0, and YOLOv5-7.0 in mean Average Precision, size, and latency.

Step-by-Step: Fine-Tuning YOLOv8 on Your Data

Fine-tuning a YOLOv8 model can be broken down into three steps:

• Creating and labeling the dataset.
• Training the model.
• Deploying it to production.

Following This Demo

A GPU-powered machine is essential. Consider using cloud-based DigitalOcean GPU Droplets for optimal performance. This tutorial assumes you're using a Jupyter Notebook

Follow these steps to set up your environment:

git clone https://github.com/gradient-ai/YOLOv8-Ballhandler
cd YOLOv8-Ballhandler
jupyter lab

Prepare your Dataset

Use RoboFlow to set up your dataset and label data and pull it into our Notebook.

• Replace the workspace and project values with your own to access your dataset.
• Change the API key to your own.

!pip install roboflow

from roboflow import Roboflow
rf = Roboflow(api_key="")
project = rf.workspace("james-skelton").project("ballhandler-basketball")
dataset = project.version(11).download("yolov8")
!mkdir datasets
!mv ballhandler-basketball-11/ datasets/

Training Your Model

The ultralytics library simplifies the training process. The model.train() method lets you fine-tune using your loaded dataset.

from ultralytics import YOLO

# Load a model
model = YOLO("yolov8n.yaml") # build a new model from scratch
model = YOLO("yolov8n.pt") # load a pretrained model (recommended for training)

# Use the model
results = model.train(data="datasets/ballhandler-basketball-11/data.yaml", epochs=10) # train the model

Testing The Model

To evaluate our new model, execute the code below.

results = model.val() # evaluate model performance on the validation set

Prediction

From there, it’s simple to submit any photo. It will output the predicted values for the bounding boxes, overlay those boxes to the image, and upload to the ‘runs/detect/predict’ folder.

from ultralytics import YOLO
from PIL import Image
import cv2

# from PIL
im1 = Image.open("assets/samp.jpeg")
results = model.predict(source=im1, save=True) # save plotted images
print(results)
display(Image.open('runs/detect/predict/image0.jpg'))