YOLOv9 Tutorial: YOLOv9 Segmentation Training On Custom Data | YOLOv9 Training | YOLOv9 Python

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

In the rapidly evolving field of computer vision, object segmentation plays a pivotal role in extracting meaningful information from images. Among the array of segmentation algorithms, YOLOv9 has emerged as a robust and adaptable solution, offering efficient segmentation capabilities with remarkable accuracy.

In this comprehensive guide, we will delve into the training process of YOLOv9 for object segmentation on custom datasets and perform inference on test data.

By the end of this tutorial, you’ll gain a firm grasp of YOLOv9’s segmentation mechanisms and learn how to apply it to your projects using custom datasets using ultralytics.

Video Tutorial:

📋 Table of Contents

• Step 1 | Download Dataset
• Step 2 | Install Ultralytics
• Step 3 | Loading YOLOv9-seg Pre-trained Model & Inference
• Step 4 | Fine-Tuning YOLOv9-seg On Custom Dataset
• Step 5 | Load Custom Model
• Step 6 | Inference on Test Images

Step 1 | Download Dataset

We will be using Furniture BBox To Segmentation (SAM) for this tutorial. To obtain the Furniture BBox To Segmentation (SAM) 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

Step 2 | Install Ultralytics

!pip install ultralytics -q

Import Packages

from ultralytics import YOLO
import matplotlib.pyplot as plt
import cv2
import pandas as pd
import seaborn as sns

Step 3 | Inference Using Pre-Trained YoloV9 Weights

model = YOLO('yolov9c-seg.pt')
model.predict("image.jpg", save=True)

1. model = YOLO('yolov9c-seg.pt'):

• This line initializes a YOLOv9 (You Only Look Once) model for object segmentation.
• The model is loaded from the file named ‘yolov9c-seg.pt’, which contains pre-trained weights and configurations for the YOLOv9 architecture specifically designed for segmentation tasks.

1. model.predict("image.jpg", save=True):

• This line executes a prediction on an input image named “image.jpg” using the initialized YOLOv9 model.
• The predict function takes the input image and performs segmentation, identifying and delineating objects within the image.
• The save=True parameter indicates that the results of the segmentation will be saved.

Step 4 | Fine-Tuning YOLOv9-seg On Custom Dataset

Configs for yolov9:

dataDir = '/content/Furniture/sam_preds_training_set/'
workingDir = '/content/'

The variable dataDir represents the directory path where the training data for the object segmentation model is located. The training data is stored in a directory named "sam_preds_training_set" under the "Furniture" directory, which is located in the "/content" directory.

Similarly, the variable workingDir represents the directory path where the main working files are stored.

num_classes = 2
classes = ['Chair', 'Sofa']

1. num_classes = 2: This variable specifies the total number of classes or categories that the model will be trained to segment. In this case, num_classes is set to 2, indicating that there are two distinct object categories that the model will learn to recognize.
2. classes = ['Chair', 'Sofa']: This list contains the names of the classes or objects that the model will be trained to identify. Each element in the list corresponds to a specific class label. The classes are defined as 'Chair' and 'Sofa', on which model will be trained to segment objects belonging to these categories.

import yaml
import os

file_dict = {
    'train': os.path.join(dataDir, 'train'),
    'val': os.path.join(dataDir, 'val'),
    'test': os.path.join(dataDir, 'test'),
    'nc': num_classes,
    'names': classes
}
with open(os.path.join(workingDir, 'data.yaml'), 'w+') as f:
  yaml.dump(file_dict, f)

1. file_dict: Creates a dictionary containing information about the dataset:

• 'train', 'val', and 'test': Paths to the training, validation, and test data directories, respectively. These paths are obtained by joining the dataDir (the directory containing the dataset) with the corresponding directory names.
• 'nc': The number of classes in the dataset, represented by the variable num_classes.
• 'names': A list of class names, represented by the variable classes.
• with open(...) as f:: Opens a file named 'data.yaml' in write mode ('w+'). If the file does not exist, it will be created. The with statement ensures that the file is properly closed after writing.
• yaml.dump(file_dict, f): Writes the contents of the file_dict dictionary to the YAML file f. The yaml.dump() function serializes Python objects into YAML format and writes them to the specified file object.

model = YOLO('yolov9c-seg.pt')
model.train(data='/content/data.yaml' , epochs=30 , imgsz=640)

Initializes a YOLOv9 model for object segmentation using the specified pre-trained weights file 'yolov9c-seg.pt'. The model is then trained on a custom dataset specified by the data parameter, which points to the YAML file 'data.yaml' containing dataset configuration details such as the path to the training and validation images, number of classes, and name of the classes.

Step 5 | Load Custom Model

best_model_path = '/content/runs/segment/train/weights/best.pt'
best_model = YOLO(best_model_path)

we’re defining the path to the best-performing model obtained during training. The best_model_path variable holds the file path where the best model weights are stored. These weights represent the learned parameters of the YOLOv9 model that achieved the highest performance on the training data.

Next, we instantiate the YOLO object using the best_model_path as the argument. This creates an instance of the YOLO model initialized with the weights from the best model obtained during training. This instantiated YOLO model, referred to as best_model, is now ready to be used for making predictions on new data.

Step 6 | Inference on Test Images

# Define the path to the validation images
valid_images_path = os.path.join(dataDir, 'test', 'images')

# List all jpg images in the directory
image_files = [file for file in os.listdir(valid_images_path) if file.endswith('.jpg')]
# Select images at equal intervals
num_images = len(image_files)
selected_images = [image_files[i] for i in range(0, num_images, num_images // 4)]
# Initialize the subplot
fig, axes = plt.subplots(2, 2, figsize=(10, 10))
fig.suptitle('Test Set Inferences', fontsize=24)
# Perform inference on each selected image and display it
for i, ax in enumerate(axes.flatten()):
    image_path = os.path.join(valid_images_path, selected_images[i])
    results = best_model.predict(source=image_path, imgsz=640)
    annotated_image = results[0].plot()
    annotated_image_rgb = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
    ax.imshow(annotated_image_rgb)
    ax.axis('off')
plt.tight_layout()
plt.show()

1. Define the path to the validation images: This line constructs the path to the directory containing the test images within the dataDir directory.
2. List all jpg images in the directory: It creates a list of all the JPEG image files present in the specified directory.
3. Select images at equal intervals: It selects a subset of images from the list to visualize. In this case, it selects one-fourth of the total images.
4. Initialize the subplot: This line creates a 2x2 grid of subplots to display the selected images and their corresponding predictions.
5. Perform inference on each selected image and display it: It iterates through each subplot, performs inference on the corresponding selected image using the best_model.predict() function, and displays the annotated image with bounding boxes or segmentation masks.
6. Finally, the subplots are arranged neatly using plt.tight_layout() and displayed using plt.show().

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

• 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!.