YOLOv9 -Object Detection On Custom Dataset | Computer Vision Project

Asad iqbal
6 min readMar 1, 2024

A Comprehensive Guide to build Object Detection Model using YOLOv9 on custom dataset!

In the rapidly evolving field of computer vision, object detection stands as a cornerstone application. Among the myriad of object detection algorithms, YOLOv9 has emerged as a powerful and versatile solution, offering real-time detection capabilities with impressive accuracy. In this comprehensive guide, we will explore YOLOv9 training process on custom dataset and make inferences on the test data.

By the end of this tutorial, you’ll have a solid understanding of how YOLOv9 works and how to implement it in your own projects using custom dataset. So, let’s dive in and explore the power of YOLOv9 together!

Video Tutorial:

We will be using Construction Site Safety Image Dataset Roboflow for this tutorial. To obtain the Construction Site Safety Image Dataset. You can access it from Kaggle, a popular platform for data science competitions, datasets, and machine learning resources.

After downloading the dataset, you may need to extract the files from the compressed format (e.g., ZIP or TAR file) if the dataset is packaged.

Dataset: Link

Clone the YOLOv9 Repository

!git clone https://github.com/SkalskiP/yolov9.git

Let’s clone the YOLOv9 repository from GitHub using the git clone command and start exploring the real-time object detection with YOLOv9. Git is a version control system that allows developers to track changes to their codebase, collaborate with others, and manage project histories.

Once the cloning process is complete, we’ll have a local copy of the YOLOv9 repository on our machine. This will enable us to explore the codebase, make modifications, and use the YOLOv9 algorithm for our own object detection tasks.

dataDir = '/content/css-data/' # css-data is the unzip path of the dataset
workingDir = '/content/' # Working Dir in google colab

Here, we’re setting up directory paths for our project. The variable dataDir holds the path to the directory where our CSS (Construction Site Safety) data is stored. We've set it to '/content/css-data/', indicating that our CSS dataset files are located in a directory named 'css-data' within the '/content/' directory.

Next, we have the variable workingDir, which defines the path to our working directory.

num_classes = 10
classes = ['Hardhat', 'Mask', 'NO-Hardhat', 'NO-Mask', 'NO-Safety Vest', 'Person', 'Safety Cone', 'Safety Vest', 'machinery', 'vehicle']

Two variables are defined:

  • num_classes is set to the integer value 10, indicating the number of classes in a classification task.
  • classes is a list containing 10 string elements representing different classes: 'Hardhat', 'Mask', 'NO-Hardhat', 'NO-Mask', 'NO-Safety Vest', 'Person', 'Safety Cone', 'Safety Vest', 'machinery', and 'vehicle'.
import yaml
import os

file_dict = {
'train': os.path.join(dataDir, 'train'),
'val': os.path.join(dataDir, 'valid'),
'test': os.path.join(dataDir, 'test'),
'nc': num_classes,
'names': classes
}

with open(os.path.join(workingDir,'yolov9', 'data.yaml'), 'w+') as f:
yaml.dump(file_dict, f)

Defines a dictionary named file_dict containing the following key-value pairs:

  • ‘train’: the path to the training data directory, formed by joining the ‘dataDir’ variable with ‘train’.
  • ‘val’: the path to the validation data directory, formed by joining the ‘dataDir’ variable with ‘valid’.
  • ‘test’: the path to the testing data directory, formed by joining the ‘dataDir’ variable with ‘test’.
  • ‘nc’: the number of classes, derived from the ‘num_classes’ variable.
  • ‘names’: a list of class names, derived from the ‘classes’ variable.

Opens a file named ‘data.yaml’ in write mode within the ‘workingDir/yolov9’ directory. Writes the content of the file_dict dictionary into the 'data.yaml' file using the YAML format. If the file doesn't exist, it creates a new one.

Download YOLOv9 Weights

!wget  https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-e.pt

We're now fetching the pre-trained weights file for YOLOv9. Using !wget, we can retrieve files directly from the web within our environment.

In this command, we're downloading the file named yolov9-c.pt from the release version v0.1 of the YOLOv9 repository hosted on GitHub. These pre-trained weights are essential for initializing our YOLOv9 model with learned parameters, allowing it to detect objects accurately.

Once downloaded, we'll be able to use these weights to fine-tune our custom dataset for object detection tasks, saving us the time and computational resources required for training the model from scratch.

cd yolov9

Now, let’s navigate to the ‘yolov9’ directory. ‘yolov9’ is the repository we’ve just cloned from GitHub. By doing this, we gain access to the code of the ‘yolov9’ model, which we will use for the training the yolov9 model on our custom dataset.

!pip install -r requirements.txt -q

Next up, we’re installing the necessary Python packages and dependencies for our project using the pip package manager. The !pip install command allows us to install Python packages from the Python Package Index (PyPI) or from a specified requirements file.

Here, -r requirements.txt tells pip to install all the packages listed in the 'requirements.txt' file. This file contains a list of Python libraries that are essential for the training of yolov9 model.

The -q flag stands for 'quiet' mode, which means that pip will execute the installation process silently without displaying detailed output unless there's an error. It's helpful for keeping the installation process clean and concise, especially in scripts or automation tasks.

!python train_dual.py --workers 8 --batch 4  --img 640 --epochs 50 --data /content/yolov9/data.yaml --weights /content/yolov9-e.pt --device 0 --cfg /content/yolov9/models/detect/yolov9.yaml --hyp /content/yolov9/data/hyps/hyp.scratch-high.yaml

This command initiates the training process for the YOLOv9 model. But before executing the cell we need to modify two arguments, first argument is — data , we will pass the data.ymal file path in — data argument that we have created above, and the 2nd argument is — cfg, we need to make little changes in the yolov9.yaml file because This original yolov9.yaml file is configured for the COCO dataset, on which the yolov9 model has trained, which contain 80 classes. As our custom dataset we have 10 classes, so we need to adjust the nc parameter in yolov9.yaml and pass nc=10 because we have 10 number of classes, instead of 80, i hope you got my point what i am saying and — hyp contain the hyperparameter for the model and we don't need to modify it, so will leave this as is it.

Now lets fine-tune the yolov9 model for our custom dataset using !python train_dual.py

  • --workers 8: This argument specifies the number of data loading workers for data preprocessing during training. A higher number can speed up data loading and preprocessing.
  • --batch 4: This argument sets the batch size for training. The batch size determines the number of samples used in each iteration of training.
  • --img 640: This argument specifies the input image size for training. In this case, images will be resized to a height and width of 640 pixels during training.
  • --epochs 5: This argument sets the number of epochs for training. An epoch is one complete pass through the entire training dataset.
  • --data /content/yolov9/data.yaml: This argument specifies the path to the YAML file containing dataset configuration information.
  • --weights /content/yolov9-c.pt: This argument specifies the path to the pre-trained weights file (yolov9-c.pt) used to initialize the model before training.
  • --device 0: This argument specifies the device to use for training. Here, 0 likely indicates the first GPU device available for training.
  • --cfg /content/yolov9/models/detect/yolov9.yaml: This argument specifies the path to the model configuration file (yolov9.yaml) defining the architecture and parameters of the YOLOv9 model.
  • --hyp /content/yolov9/data/hyps/hyp.scratch-high.yaml: This argument specifies the path to the hyperparameters file (hyp.scratch-high.yaml) containing tuning parameters for the model during training.

Inferences

!python detect.py --img 640 --conf 0.1 --device 0 --weights /content/yolov9/runs/train/exp2/weights/best.pt --source /content/css-data/test/images/004763_jpg.rf.46484e6ca73caeaa9de45822cf1085a9.jpg
from IPython.display import Image
Image(filename="/content/yolov9/runs/detect/exp2/004763_jpg.rf.46484e6ca73caeaa9de45822cf1085a9.jpg", width=600)

“Embrace each learning moment with joy, for happiness fuels the journey of knowledge.” 🌟

Thanks for reading; if you liked my content and want to support me, the best way to supporting me on Patreon —

  • Subscribe my YouTube channel
  • Connect With Me On LinkedIn and Github where I keep sharing such free amazing content to become more productive and effective at what you do using Technology and AI.
  • Need help with ML & DL? Check out my Fiverr services!.

--

--