YOLOv8 Custom Object Detection

A Simple Guide for Parameter Tuning and Class-Based Detection with YOLOv8

YOLOv8 is the latest version of the YOLO object detection and image segmentation model developed by Ultralytics.

The model that builds upon the success of previous YOLO versions and introduces new features and improvements to further boost performance and flexibility.

In addition to its extensibility, YOLOv8 includes a number of other innovations that make it an appealing choice for a wide range of object detection and image segmentation tasks. These include a new backbone network, a new anchor-free detection head, and a new loss function.

In this article, we are going to use YOLOv8 to train our custom object detection model.

Prerequisites

It’s good to have a basic knowledge of deep learning computer vision and how to work in a Google Colab environment.

Steps Covered in this Tutorial

To train our own custom object detector these are the steps to follow

• Preparing the dataset
• Environment Setup: Install YOLOv8 dependencies
• Setup the data and the directories
• Setup the YAML files for training
• Training the model
• Evaluate the model
• Running inference on test images
• Export the weight files for later use

Preparing the dataset

I get the data from pictures i took from my own phone.
Instead, you can work with ready-made datasets via Roboflow or Kaggle.

My dataset consists of 62 images.
I label the data with roboflow and separate 2 classes as book and notebook.
Then duplicate the data. This is called the augmentation process.
You can duplicate your data from Roboflow according to the needs of your project. I reproduced it with blur, noise and greyscale.

Setup the YAML files for training

To train the model we need a yaml file like below.

train: ./train/images 
val: ./valid/images  
nc: 2
names: ['book', 'notebook']

The “train” and “val” sections in the YAML file determine the paths to the training and validation datasets, respectively. These sections specify where the image data used during training is located. In this example, the image data is stored in the directories “./train/images” and “val: ./valid/images”.

The “nc” section specifies the number of classes in the model. In this example, there are two classes, so the “nc” value is set to 2.

The “names” section determines the names of the classes. In this example, there are two classes (‘book’, ‘notebook’). The class names represent the object classes that will be shown in the model’s output.

This YAML file defines the parameters used in training the YOLO model and the paths to the dataset. This file can affect the accuracy and efficiency of the model.
After the dataset is prepared, we can start training our model in the Colab environment.

Environment Setup

You need a google account to use Google Colab.
We use Colab for tasks that require intensive computation, such as deep learning.
Since my computer's GPU is insufficient, we need to activate Colab's GPU support.

After doing this, we check the gpu activity.

The NVIDIA Tesla T4 enterprise graphics card offers breakthrough performance for mainstream servers and data centers. It supports up to 16GB of memory and 2560 CUDA cores to accelerate a wide range of modern applications.

With the code below, we download the necessary libraries and make their definitions.

from ultralytics import YOLO
from IPython.display import display,Image

Exporting Dataset from Roboflow

Once the dataset version is generated, we have a hosted dataset we can load directly into our notebook for easy training. Click Export and select the YOLOv8dataset format.

!pip install roboflow --quiet

from roboflow import Roboflow
rf = Roboflow(api_key="your-api-key")
project = rf.workspace("batuhansener").project("bookornotebook")
dataset = project.version(1).download("yolov8")

Custom Training

After the configuration is done we can begin our training.

There are multiple hyperparameters that we can specify which are:

• img: define input image size
• batch: determine batch size
• epochs: define the number of training epochs.
• data: set the path to our YAML file
• name: result names
• cache: cache images for faster training

We need to specify the path of our YAML files which we created above.

!yolo task=detect mode=train model=yolov8m.pt data={dataset.location}/data.yaml epochs=128 imgsz=800 plots=True

The trained model is saved as best.pt.
Google Colab notebooks have an idle timeout of 90 minutes and absolute timeout of 12 hours.
I want to run my trained model for this in my local.

We will examine our results under the runs folder.

You can assign the model to the variable and return the results you want from its predictions.

from ultralytics import YOLO

model = YOLO('best2.pt')
res=model.predict(
   source='book.jpg',
   conf=0.80,
   save=True
)
for result in res:
    
    result.boxes.xyxy   # box with xyxy format, (N, 4)
    result.boxes.xywh   # box with xywh format, (N, 4)
    result.boxes.xyxyn  # box with xyxy format but normalized, (N, 4)
    result.boxes.xywhn  # box with xywh format but normalized, (N, 4)
    result.boxes.conf   # confidence score, (N, 1)
    result.boxes.cls    # cls, (N, 1)

  

We can select the classes that we want to predict using the class parameter.

model = YOLO('best2.pt')
results=model.predict(source='b-nb.jpg',save=True,classes=0)

model = YOLO('best2.pt')
results=model.predict(source='b-nb.jpg',save=True,classes=1)

Conclusion

I hope you were able to follow along and was able to train successfully.

I explained YOLOv8 and taught how to do custom object detection. We trained the model and tried to show some useful parameters.

If you have any questions, recommendations, or critiques, I can be reached via LinkedIn or via my mail. Feel free to reach out to me.

References:

• Official Repo: https://github.com/ultralytics/ultralytics
• Roboflow
• Roboflow Blog