Apply Data Augmentation on YOLOv5/YOLOv8 Dataset

How to apply the augmentation on YOLOv5 or YOLOv8 dataset using albumentations library in Python?

Variations of Augmented Images — An Example. (Image: Pranjal Ostwal on Medium)

Data Augmentation

Data augmentation is a technique in machine learning or deep learning where new data is created from the existing data by applying various transformations such as flipping, rotating, and changing the brightness/contrast.

Data augmentation is commonly used in computer vision tasks but can also be used in other domains such as natural language processing and speech recognition.

Why Data Augmentation?

It can help prevent overfitting by increasing the diversity of the training data and helping the model learn to generalize to new and unseen data. It can also improve the accuracy and robustness of the model by exposing it to a wider range of data variations, and can be useful when the amount of available training data is limited, as it helps increase the size of the dataset and improve the performance of the model.

Data Augmentation

Dataset Format of YOLOv5 and YOLOv8

Both YOLOv8 and YOLOv5 have same dataset format which mainly contain two directories.

• Images directory contains the images
• labels directory contains the .txt files. Each .txt file contains the normalized bounding boxes in a following format.

[class_label, x_center, y_center, width, height]

How to Use Albumentations Library?

Albumentations is a Python library for image augmentation that offers a simple and flexible way to perform a variety of image transformations. The steps to use this library are followed.

• Install the albumentations library.

pip install -U albumentations

• Minimum augmentation pipeline.

import albumentations as A

transform = A.Compose([
    A.RandomCrop(width=450, height=450),
    A.HorizontalFlip(p=0.5),
    A.RandomBrightnessContrast(p=0.2),
], bbox_params=A.BboxParams(format='yolo'))
transformed = transform(image=image, bboxes=bboxes)
transformed_image = transformed['image']
transformed_bboxes = transformed['bboxes']

Code description

• Firstly, we need to import the albumentations library.
• Defined the augmented operation in a A.Compose object (as in the code) which you would like to apply.
• We need to convert the YOLO input bounding box label into following albumentation format.

[x_center, y_center, width, height, class_name]

Example input and output data for bounding boxes augmentation

• Passed the YOLO inputs image and bounding box list in albumentation format to transform object which will return the augmented results in transformed (dictionary).

Save Augmentation

• If you want to save the transformed images and labels.
  - Use cv2.imwrite function to save transformed images.
  - Before saving the augmented label need to convert into yolo format.

Results

Input

Input Image

Input Labels

Output

Augmented Image

Augmented Text File

Visualize the augmented output by drawing label on it

GitHub Code

You can get the full code from my GitHub repo.

Summary

Here I have just discussed how to get the augmented dataset of YOLOv5 and YOLO8 dataset for object detection. We can also get the augmented dataset of other format of dataset using same library in Python. Not only for object detection but also used for classification as well segmentation purpose.

Reference

https://albumentations.ai/docs/

Muhammad Faizan is a Machine Learning Engineer at Red Buffer. You may reach out to him at muhammad.faizan@redbuffer.net or through his LinkedIn.