Train a YOLOv8 object detection model in Python

Object detection is a computer vision task that aims to locate objects in digital images. As such, it is an instance of artificial intelligence that consists of training computers to see as humans do, specifically by recognizing and classifying objects according to semantic categories.

Object localization is a technique for determining the location-specific objects in an image by demarcating the object through a bounding box. Object classification is another technique that determines to which category a detected object belongs. The object detection task combines subtasks of object localization and classification to simultaneously estimate the location and type of object instances in one or more images. Objects are identified by way of bounding boxes. A bounding box is the quadrate output demarcating a detected object as predicted by the model.

Object detection overlaps with other computer vision techniques, but developers nevertheless treat it as a discrete endeavor. Image classification aims to classify entire images according to defined categories. Object detection, by comparison, delineates individual objects in an image using bounding boxes according to specified categories. Image segmentation is similar to object detection, albeit more precise. Like object detection, segmentation delineates objects in an image according to object classes. But rather than mark objects using boxes, segmentation demarcates objects at the pixel level.

This tutorial provides an introduction to preparing images for object detection and fine-tuning a pretrained model for a specific task. Given its ease of implementation and quick image processing speed, we will focus on the YOLO (You Only Look Once) architecture.

YOLO

Machine learning algorithms for object detection vary depending on the detection model architecture with which one chooses to work. Moreover, object detection algorithms change depending on whether one works with video, which is a task for object tracking, or images. Other factors, such as whether one uses a pretrained model or trains a custom object detection model from scratch, can affect workflows, required dependencies, and so forth. It is therefore difficult, if not impossible, to provide a one-size-fits-all end-to-end template for object detection in general.

YOLO (You Only Look Once) is a family of single-stage network architectures based in Darknet, an open-source convolutional neural network (CNN) framework. The YOLO architecture prioritizes speed. Indeed, YOLO’s speed makes it preferable for real-time object detection and has earned it the common descriptor of state-of-the-art object detector.

YOLO differs from R-CNN (another popular architecture) in several ways. While R-CNN passes extracted image regions through multiple networks that separately extract features and classify images, YOLO condenses these actions into a single network. Secondly, compared to R-CNN’s approximately 2000 region proposals, YOLO makes less than 100 bounding box predictions per image.

In addition to being faster than R-CNN, YOLO also produces less background false positives, although it has a higher localization error. There have been many updates to YOLO since its inception, generally focusing on speed and accuracy. The latest iteration of YOLO, YOLOv8, further demonstrates increased performance in detecting small objects. Nevertheless, its single-stage YOLO architecture remains largely the same.

For our architecture in this tutorial, we adopt the latest YOLO version, YOLOv8. We will explore how to fine tune a pretrained object detector for a marine litter data set using Python code. In this way, you will explore a real-world application of object detection while becoming familiar with a YOLO algorithm and the fundamental terminology and concepts for object detection.

Prerequisites

You need an IBM Cloud account to create a watsonx.ai project.

You need a Kaggle account to load the data set.

Steps

Step 1. Set up your environment

While you can choose from several tools, this tutorial walks you through how to set up an IBM account to use a Jupyter Notebook. Jupyter Notebooks are widely used within data science to combine code, text, images, and data visualizations to formulate a well-formed analysis.

1. Log in to watsonx.ai using your IBM Cloud account.

2. Create a watsonx.ai project.

3. Create a Jupyter Notebook.

A notebook environment opens for you to load your data set and copy code from this tutorial to tackle a simple object detection model training task. To view how each block of code works, each step’s code block is best inserted as a separate cell of code in your watsonx.ai project notebook.

Step 2. Install and import relevant libraries

You need a few libraries for this tutorial. Make sure to import the necessary Python libraries. If they're not installed, you can resolve this with a quick pip install (included at the beginning of the following code).

Previous iterations of YOLO (for example, YOLOv5) require cloning the architecture's Github repository. YOLOv8, however, is the first iteration to have its own official Python package. Thus, we can install YOLOv8 via pip. Additional libraries, such as shutil and glob, are necessary for our data preprocessing.

An advantage of using YOLO to train and test our model is that it does not necessarily require other Python libraries, such as opencv, NumPy, or matplotlib, which are commonly used when training object detection models with platforms such as TensorFlow. Installing a litany of libraries and packages can become unwieldy, even with the TensorFlow object detection API. YOLOv8's simplified installation is a welcome advancement in increasing machine learning accessibility.

# install these to avoid potential dependency errors
%pip install torch torchvision torchaudio
%pip install opencv-contrib-python-headless

# install and import Ultralytics YOLOv8
%pip install ultralytics==8.0.196
import ultralytics
ultralytics.checks()
from ultralytics import YOLO
import yaml

# import packages to retrieve and display image files
import glob
import os
import shutil
from PIL import Image
from IPython.display import display

Step 3. Load the data

There are a wide array of benchmark data sets for testing object detector performance, such as the COCO data set or ImageNet, among many others. While these data sets' massive sizes make them ideal for training generalizable object detection models, they are computationally expensive, requiring advanced CPUs and GPUs to handle the massive number of input images.

For this tutorial, we will use a marine litter data set available on Kaggle through an MIT license and hosted by the University of Minnesota Libraries. With this data set, we will fine tune a pretrained YOLOv8 model to detect and classify litter and biological objects in underwater images. We will thereby explore computer vision's potential applications in ecological conservation efforts.

We will load the data set directly into our notebook using the Kaggle API. This requires creating a free Kaggle account. Once you have created an account, you need to generate an API key. Directions for generating your key are found in the Kaggle API documentation under the section "API credentials."

Copy and paste your Kaggle username and API key into the following code. Then run the code to install the API and load the data set into watsonx.

# install API
% pip install kaggle

# replace "username" string with your username
os.environ["KAGGLE_USERNAME"] = "username"
# replace "apiKey" string with your key
os.environ["KAGGLE_KEY"] = "apiKey"

# load dataset
!kaggle datasets download atiqishrak/trash-dataset-icra19 --unzip

Now we will print and save our working directory, as we will call this directory throughout the tutorial. We will also print the contents of our working directory to confirm the "trash_ICRA19" data set was loaded properly.

# store working directory path as work_dir
work_dir = os.getcwd()

# print work_dir path
print(os.getcwd())

# print work_dir contents
print(os.listdir(f"{work_dir}"))

If you see "trash_ICRA19" among the directory's contents, then it has loaded successfully. We can further confirm the data set loaded properly by viewing its contents with the following command:

# print trash_ICRA19 subdirectory contents
print(os.listdir(f"{work_dir}/trash_ICRA19"))

You will see three files/folders: a config.yaml file, a videos_for_testing directory, and a dataset directory.

We will ignore the videos_for_testing directory, so feel free to delete it with the following code:

# run this to delete videos_for_training folder
rm -r {work_dir}/trash_ICRA19/videos_for_testing

We will use the config.yaml file and the contents of the dataset directory to train our object detection model. Before doing so, however, we need to modify the dataset directory structure to ease processing. But first, let's discuss YOLO label formats.

Each object detection architecture requires a different annotation format and file type for processing bounding box labels. For example, Mask R-CNN, DETR, and Meta's Detectron2 all use COCO format labels stored in a central .json file. Faster R-CNN and MobileNet SSD v2 use Tensorflow's binary TFRecord format. YOLO models typically use Pytorch's .txt format. Because we are training a YOLO model in this tutorial, we will focus on this last format.

What is the YOLO text format? How does it differ from other formats? First, YOLO image labels are stored as .txt files, while VOC PASCAL stores iamges labels as .xml files, and COCO labels are compiled in .json. Second, in YOLO, each image possesses its own corresponding .txt file, which contains all of the labeled objects in that image. Each row in the file provides information on a different labeled object. Similarly, VOC PASCAL also creates a seperate .xml file for each image. By contrast, however, COCO generally saves all labeled objects in a single .json file, which points to all of the images in a given directory.

YOLO further differs from PASCAL VOC and COCO in how it records object labels. PASCAL VOC, for example, records bounding box coordinates using the box's top-left (x,y) coordinate point and its bottom-right coordinate point. COCO records the top-left (x,y) coordinate point with the box's width and height. YOLO, however, records a bounding box's central coordinate value alongside its width and height. A typical YOLO bounding box record adheres to this structure:

class_id x_center y_center width height

What does this look like in practice? Here is a sample image from our marine litter data set followed by its .txt label file:

0 0.58125 0.5375 0.3125 0.313888888889
1 0.827083333333 0.459722222222 0.345833333333 0.180555555556

Each row corresponds to a different labeled object. The image corresponding to our sample .txt file contains two labeled objects. The first object belongs to class 0 ('plastic'). The center of its bounding box is (.58125, .5375), and it has a width of .3125 and height of .313888888889. The second object belongs to class 1 ('rov'). This bounding box's center is (.827083333333, .459722222222) with a width of .345833333333 and height of .180555555556. We can visualize the bounding boxes and coordinates over the image:

Note the following stipulations on label records from Ultralytics' YOLO documentation:

• One .txt file per image.
• If an image contains no objects, it requires no .txt file.
• One object per row in a .txt file. Each row's information must follow the format: class_id x_center y_center width height.
• Bounding box coordinates are normalized, being values between 0 and 1. If bounding box coordinates are in pixels, you must divide x_center and width by the image's width, and y_center and height by the image's height.

Step 4. Preprocess the data

Fortunately, all labels in the marine litter data set are already formatted as YOLO .txt files. But we need to rearrange the structure of the image and label directories in order to help our model more readily process the image and labels.

As it stands, our loaded data set directory follows this structure:

But, YOLO models by default require seperate image and label to subdirectories within the train/val/test split. Additionally, the YOLOv5 architecture does not process annotations stored as .xml files. It is therefore better to delete these files in spirit of best practices in order to create a clean workspace and prevent any later confusion. We need to reorganize the directory into the following stucture:

To reorganize the data set directory, we can run the following script (be certain to replace directory path with your own):

# function to reorganize dir
def organize_files(directory):
    for subdir in ['train', 'test', 'val']:
        subdir_path = os.path.join(directory, subdir)
        if not os.path.exists(subdir_path):
            continue

        images_dir = os.path.join(subdir_path, 'images')
        labels_dir = os.path.join(subdir_path, 'labels')

        # create image and label subdirs if non-existent
        os.makedirs(images_dir, exist_ok=True)
        os.makedirs(labels_dir, exist_ok=True)

        # move images and labels to respective subdirs
        for filename in os.listdir(subdir_path):
            if filename.endswith('.txt'):
                shutil.move(os.path.join(subdir_path, filename), os.path.join(labels_dir, filename))
            elif filename.endswith('.jpg') or filename.endswith('.png') or filename.endswith('.jpeg'):
                shutil.move(os.path.join(subdir_path, filename), os.path.join(images_dir, filename))
            # delete .xml files
            elif filename.endswith('.xml'):
                os.remove(os.path.join(subdir_path, filename))

if __name__ == "__main__":
    directory = f"{work_dir}/trash_ICRA19/dataset"
    organize_files(directory)

Next, we need to modify the .yaml file for the data set. This is a master file, so to speak, that provides paths to the training, validation, and test data, as well as ids for our object classifications. As explained in Ultralytics' YOLO documentation, this .yaml file configures the data set. It specifically defines the root directory for the data set, relative subdirectory paths to the training, validation, and test image subsets, and finally a dictionary of label class IDs with names. This is the setup we will use in our .yaml file (you do not need to copy this into any file or code block, as we will run code to do so automatically):

path: /path/to/dataset/directory # root directory for dataset
train: train/images  # train images subdirectory
val: train/images  # validation images subdirectory
test: test/images # test images subdirectory

# Classes
names:
  0: plastic
  1: bio
  2: rov

Class ID numbers start from 0. This is necessary if ever creating your own object labels for a custom data set.

Run the following script to delete the current contents of config.yaml and replace it with the above contents that reflect our new data set directory structure. Be certain to replace the work_dir portion of the root directory path in line 4 with your own working directory path we retrieved earlier. Leave the train, val, and test subdirectory definitions. Also do not change {work_dir} in line 23 of the code.

# contents of new confg.yaml file
def update_yaml_file(file_path):
    data = {
        'path': 'work_dir/trash_ICRA19/dataset',
        'train': 'train/images',
        'val': 'train/images',
        'test': 'test/images',
        'names': {
            0: 'plastic',
            1: 'bio',
            2: 'rov'
        }
    }

    # ensures the "names" list appears after the sub/directories
    names_data = data.pop('names')
    with open(file_path, 'w') as yaml_file:
        yaml.dump(data, yaml_file)
        yaml_file.write('\n')
        yaml.dump({'names': names_data}, yaml_file) 

if __name__ == "__main__":
    file_path = f"{work_dir}/trash_ICRA19/config.yaml" #.yaml file path
    update_yaml_file(file_path)
    print(f"{file_path} updated successfully.")

Step 5. Train the YOLOv8 model

For this tutorial, we will fine tune a pretrained YOLO model for our underwater trash detection task. It is, of course, possible to train a model from scratch. Given our data set is relatively small, however, it may be better to leverage a pretrained model's weights. Another approach is to modify or only load select layers of a pretrained model. Doing so requires a more advanced knowledge of model architectures, however, and so beyond the scope of the present tutorial.

We will run the following command-line code to fine tune a pretrained the default YOLOv8 model. Depsite the relatively small size of our image data set, the training process might take up to around four hours to complete.

!yolo task=detect mode=train data={work_dir}/trash_ICRA19/config.yaml model=yolov8s.pt epochs=2 batch=32 lr0=.04 plots=True

While waiting for our model to finish training, we can review some of the parameters in the model training command:

• task: This parameter specifies the computer vision task for which you are using the specified YOLO model and data set. The available options for this parameter are detect, segment, and classify, which respectively denote object detection, image segmentation, and image classification. We could eschew this parameter, leaving YOLOv8 to guess our desired task from the data.

• mode: This denotes the purpose for which you are loading the specified model and data. Since we are training a model, it is set to "train." Later, when we test our model's performance, we will set it to "predict."

• epochs: This delimits the number of times YOLOv8 will pass through our entire data set. Admittedly, two epochs is incredibly low. But computer vision model training is computationally expensive, even with our small data set. In order to reduce training time, we are finetuning our model for only two epochs.

• batch: This numerical value stipulates the training batch sizes. Batches are the number of images a model processes before it updates its parameters. A smaller batch size allows a model to learn from more individualized examples, and in turn generalize better. But this also requires more training time. A larger batch size decreases training time, but is generally believed to produce a model that generalizes with less accuracy.

• lr0: This sepcifies the model's initial learning rate. In YOLOv8, the default value is .01. The learning rate controls how much a model adjusts its weights in response to estimated error. In other words, learning rate is how fast a model learns.

• plots: This parametere directs YOLO to generate and save plots of our models training and evaluation metrics. We will call these later in order to evaluate our model's performance.

Literature often designates learning rate and batch size as two of the most influential training parameters with respect to object detection models. Some research studies suggest these parameters should increase proportionally, although such studies deal with larger data sets and can show different recommended proportions. Ideal model training parameters can largely depend on one's data set and computational resources, such as GPUs.

Step 6. Test the model

We can now run inference to test the performance of our fine-tuned model:

!yolo task=detect mode=predict source={work_dir}/trash_ICRA19/dataset/test/images model={work_dir}/runs/detect/train/weights/best.pt conf=0.5 iou=.5 save=True save_txt=True

This brief script generates predicted labels for each image in our test set, as well as new output image files that overlaying the predicted bounding box atop the original image. Predicted .txt labels for each image are saved via the save_txt=True argument and the output images with bounding box overlays are generated through the save=True argument. The parameter conf=0.5 informs are model to ignore all predictions with a confidence level less than 50%. Lastly, iou=.5 directs the model to ignore boxes in the same class with an overlap of 50% or greater. This helps to reduce potential duplicate boxes generated for the same object.

If we want to view how our model performs on a handful of images, we can load the images with predicted bounding box overlays. To do so, we run the script below. This will display ten images from the test set with their predicted bounding boxes, accompanied by class name labels and confidence levels (or class probability).

# print first ten images from preceding prediction task
for pred_dir in glob.glob(f'{work_dir}/runs/detect/predict/*.jpg')[:10]:
    img = Image.open(pred_dir)
    display(img)

Step 7. Evaluate the model

Previous interations of YOLO used the tensorboard to visualize model performance over its training cycle. YOLOv8 offers an advancement from this. Whenever you train a model in YOLOv8, the system evaluates the model's performance and generates a series of data visualizations for these evaluations.

Two common evaluation metrics for object detection (as well as machine learning models and classifiers writ large) are precision and recall.

Precision, also called positive predicted value (PPV), signifies the proportion of predicted positive samples that actually belong to the class in question. Recall, also known as sensitivity or true positive rate (TPR), is the percentage of actual class instances correctly classified by the model. They are respectively represented by these equations:

YOLOv8 produces visualizations of a final generated model's precision and recall for each class. These are saved in the home directory, under the train folder. The precision score is displayed in P_curve.png:

img = Image.open(f'{work_dir}/runs/detect/train/P_curve.png')
display(img)

This should return an image similar (if not identical) to this:

The graph visualizes our model's precision for each individual class and every class together. The graph shows an exponential increase in precision as the model's confidence level for predictions increase. But model precision has not yet leveled out at a certain confidence level after two epochs. The recall graph (R_curve.png) displays an inverse trend:

img = Image.open(f'{work_dir}/runs/detect/train/R_curve.png')
display(img)

This code should return an image similar (if not identical) to this:

This graph shows a similarly exponential movement in model recall for predictions. But unlike precision, recall moves in the opposite direction, showing greater recall with lower confidence instances and lower recall with higher confidence instances. This is an apt example of the trade-off in precision and recall for classification models.

YOLOv8 provides a wide array of additional data visualizations, such as confusion matrices and loss function graphs. Although it is relatively simple to create one's own confusion matrix, these additional outputs save time in model evaluation. Let's call the automatically generated confusion matrix:

img = Image.open(f'{work_dir}/runs/detect/train/confusion_matrix_normalized.png')
display(img)

This code should return an image similar (if not identical) to this:

Here, rows are predicted values, and columns are ground truth values. We can see that objects classified as "plastic" are predicted with 88% accuracy. Similarly, while our model correctly classifies 88% of bio-class instances, 2% of bio-class instances are misclassifed as plastic, and 10% are ignored altogether, assumed to be background matter. A similar situation occurs with the "rov" class.

The final column suggests that most of our errors result from false positives. The top three squares of the final column list the percentage of background (that is, unlabled) information the model erroneously classifies as plastic, bio, and rov. One solution to reduce false positives is simply to acquire more training data. But, this is not possible for us. We could decrease the batch size to the YOLOv8 default: 16. While this may increase training time, it will allow our model to learn from more individualized data instances. If you choose to retrain the model, be sure to lower the learning rate (the YOLOv8 default is .01).

Step 8. Intersection over union

Intersection over union (IoU) is a common evaluation metric used in object detection models. IoU calculates the ratio of two bounding boxes’ area of intersection (that is, the area of boxes’ overlapping sections) over their area of union (that is, the total area of both boxes combined):

We can visualize this equation as:

Models use IoU to measure prediction accuracy by calculating the IoU between a predicted bounding box and ground truth bounding box for the same object. Model architectures also use IoU to generate final bounding box predictions. Because models often initially generate several hundred bounding box predictions for a single detected object, models use IoU to weigh and consolidate bounding box predictions into a single box per detected object.

Because we used the save_txt=True argument when running our YOLO model predictions, our model saved .txt file labels for all of its predictions.

Run the following code to save a function that calculate IoU between two bounding boxes:

# function to compute IoU for predicted versus ground truth labels
def compute_iou(gt_box, pred_box):
    x1, y1, w1, h1 = gt_box
    x2, y2, w2, h2 = pred_box

    # produce BBox coordinates from YOLO data
    xmin_gt_box = x1 - w1 / 2
    ymin_gt_box = y1 - h1 / 2
    xmax_gt_box = x1 + w1 / 2
    ymax_gt_box = y1 + h1 / 2

    xmin_pred_box = x2 - w2 / 2
    ymin_pred_box = y2 - h2 / 2
    xmax_pred_box = x2 + w2 / 2
    ymax_pred_box = y2 + h2 / 2

    # calculate intersection of two boxes
    xmin_inter = max(xmin_gt_box, xmin_pred_box)
    ymin_inter = max(ymin_gt_box, ymin_pred_box)
    xmax_inter = min(xmax_gt_box, xmax_pred_box)
    ymax_inter = min(ymax_gt_box, ymax_pred_box)
    inter_area = max(0, xmax_inter - xmin_inter) * max(0, ymax_inter - ymin_inter)

    # calculate union of two boxes
    gt_box_area = w1 * h1
    pred_box_area = w2 * h2
    union_area = gt_box_area + pred_box_area - inter_area

    # IoU equation
    iou = inter_area / union_area if union_area != 0 else 0
    return iou

This function takes two lists of YOLO formatted bounding box data and calculates the IoU between them. To calculate IoU, we need the ground turth and predicted bounding box coordinates for a single image. Run this script to print bounding box coordinates for the image bio0001_frame0000208.jpg:

# print ground truth cooridnates
gt_labels = f'{work_dir}/trash_ICRA19/dataset/test/labels/bio0001_frame0000208.txt'
with open(gt_labels, 'r') as file:
    print("Ground truth:\n" + file.read())

# print prediction cooridnates
pred_labels = f'{work_dir}/runs/detect/predict8/labels/bio0001_frame0000208.txt'
with open(pred_labels, 'r') as file:
    print("Predictions:\n" + file.read())

Your output should look like this:

The printed output reveals a potential problem. We see that the model's predicted labels are organized according to ascending class ID values, while the manually-created ground truth labels are organized according to descending class ID values. If we wanted to calcualte IoU for all of our predicted objects, we need to assure that each line in respective ground truth and predicted label files correspond to the same objects and classes.

Nevertheless, for now, we can calculate the IoU for one object. The following script takes the label coordinates for the class 1 object in bio0001_frame0000208.png:

gt_box = [0.9, 0.35, 0.2, 0.327777777778]
pred_box = [0.895526, 0.328866, 0.208602, 0.357924]

iou = compute_iou(gt_box, pred_box)
print("IoU:", iou)

The IoU is approximately .8467. This is not terrible, but neither is it ideal. Thus, it may be worth retraining the model with a lower batch size and measure its improvement.

Note that IoU can be difficult to compute en masse. An object detector might fail to identify all of the user-labeled objects in an image, might assign objects to the wrong class, and might record object labels in output .txt files in a different order than the user-generated label record (as with our data set). Nevertheless, IoU remains one common method for interpreting object detection performance.

Summary and next steps

This beginner tutorial provides an overview for how to use Python to train a YOLOv8 object detection model and compute common evaluation metrics for its predictions. Note that there are a myriad other object detection algorithms and architectures, such as Fast R-CNN or Detectron 2. Each of these require different code for training, fine tuning, and testing models. Object detection has a wide array of applications, from robotics to security systems. Here, we've explored its potential role in environmental conservation.

Try watsonx for free

Build an AI strategy for your business on one collaborative AI and data platform called IBM watsonx, which brings together new generative AI capabilities, powered by foundation models, and traditional machine learning into a powerful platform spanning the AI lifecycle. With watsonx.ai, you can train, validate, tune, and deploy models with ease and build AI applications in a fraction of the time with a fraction of the data.

Try watsonx.ai, the next-generation studio for AI builders.

Next steps

Explore more articles and tutorials about watsonx on IBM Developer.