Train YOLOv8 on a custom dataset

Khoa Le, Ph.D.
3 min readNov 12, 2023

--

Welcome to an in-depth tutorial on leveraging YOLOv8, a cutting-edge object detection algorithm, to train and validate detection models specifically tailored for HumanCrowd and MOT20 datasets. This tutorial assumes that you have already followed our previous part on data preparation, ensuring that your datasets are organized and annotated appropriately for the training process.

Why YOLOv8?

You Only Look Once version 8 (YOLOv8) has gained prominence for its real-time object detection capabilities, offering a compelling balance between speed and accuracy. Its modular design allows for customization, making it an ideal choice for diverse datasets with varying object types, such as HumanCrowd and MOT20.

Overview

In this tutorial, we will guide you through the following essential steps in training and validating your object detection model:

  1. Configuring YOLOv8:
  • Understand the YOLOv8 architecture and configuration files.
  • Choose an appropriate backbone network and set hyperparameters based on the characteristics of the HumanCrowd and MOT20 datasets.

2. Training the Model:

  • Initialize your YOLOv8 model with pre-trained weights or from scratch.
  • Train the model on the prepared datasets using the YOLOv8 training script.
  • Monitor and optimize training parameters for improved performance.

3. Validation Process:

  • Evaluate the trained model on the validation dataset to gauge its accuracy and generalization capabilities.
  • Explore metrics such as precision, recall, and mean average precision (mAP) to assess detection quality.
  1. Inference and Deployment:
  • Use the trained YOLOv8 model for real-time or batch inference on new images or videos.
  • Understand the post-processing techniques to refine and filter detection results.

By the end of this tutorial, you’ll have a comprehensive understanding of how to harness the power of YOLOv8 for object detection tasks on challenging datasets like HumanCrowd and MOT20. Let’s dive in and elevate your object detection capabilities with YOLOv8!

Training

In order to start training, we need to prepare a config yaml file yolov8.yaml as below. This config file defines training parameters such as image size, number of epochs, batch size, and the path to the dataset.yaml file which contains YOLO data format as in the previous tutorial. You can download a pre-trained model from the Ultralytics page, and add it to the model path, the architecture and weights will be initialized automatically.

#Verbose during prediction
verbose: False

img_size: 720
epochs: 200
batch_size: 8

#for training
mode: 'train'
data: "datasets/all_data/dataset.yaml"
output: "yolov8n_mot_ch"
model: weights/yolov8n.pt #download pretrained model on ultralytics
resume: False

I define a class TrainerYoloV8 as below to use the YOLO class to train and validate the detection model. If you want to add more parameters to the training, please check the instructions in this link.

from ultralytics import YOLO

class TrainerYoloV8():
def __init__(self, cfg):
self.cfg = cfg
self.data = cfg['data']
self.imgsz = cfg['img_size']
self.batch = cfg['batch_size']
self.epochs = cfg['epochs']
self.output = cfg['output']
self.mode = cfg['mode']
self.resume = cfg['resume']

# loading a pretrained YOLO model
self.model = YOLO(cfg['model'])

def train(self):
results = self.model.train(
mode=self.mode,
data=self.data,
imgsz=self.imgsz,
epochs=self.epochs,
batch=self.batch,
name=self.output,
resume=self.resume,
)

def validate(self):
results = self.model.val(
data=self.data,
imgsz=self.imgsz,
name=self.output,
)

with open(args.config) as file:
cfg = yaml.load(file, Loader=yaml.FullLoader)

trainer = TrainerYoloV8(cfg)
trainer.train()
trainer.validate()

There are two ways to run inference with YOLO library, the first way is in the code below, where you initiate the YOLO class by using the path of a trained model. then you create an iterator of detected results and loop through it. The Result class contains everything you need, from images to detected boxes, classes, or segments if you use segmentation models. An example of running the detector with camera input is presented in this code:

conf = 0.5
model_path = 'weights/yolov8n.pt'
model = YOLO(model_path)

results = model(source=0,
conf=conf,
show=True,
stream=True,
classes=[0],)

for frame_idx, r in enumerate(results):
continue

In summary, this tutorial has guided you through the YOLOv8 library for training, validation, and inference on a custom dataset. From meticulous data preparation to model fine-tuning, you’ve gained the skills to deploy an accurate and efficient object detection system.

Looking ahead, consider extending your capabilities by exploring the fusion of YOLOv8 with advanced tracking algorithms. Our next tutorial will delve into seamlessly integrating YOLOv8 detection with state-of-the-art tracking techniques, enabling you to elevate your computer vision applications through precise and dynamic human tracking. Stay tuned for an exciting continuation of your journey into the world of computer vision!

--

--

Khoa Le, Ph.D.

I do Data Science on Medical Imaging and Finance, and love them both.