How to augment images and OBB labels to expand a YOLO training set

thijs · September 2, 2026, 12:17am

Based on contributions by Luckyhipster10, Mathijs.

If you don’t have enough hand-annotated images for a good oriented bounding box (OBB) model, you can generate extra training images by applying random transformations (flips, rotation, brightness changes) to the ones you already have — and automatically update their OBB labels to match. This uses Albumentations, which supports keypoint-based transforms that this script reuses to move OBB corner points along with the image.

This works for OBB datasets in the YOLO label format class_id x1 y1 x2 y2 x3 y3 x4 y4 (normalized corner coordinates). For plain axis-aligned boxes, use Albumentations’ built-in bounding-box augmentation instead.

What you need

Python with pip install opencv-python albumentations
An existing YOLO OBB dataset: an images/ folder and a matching labels/ folder (one .txt per image), as exported from How to train a YOLO OBB model and get the center and angle of detections

Steps

1. Install and import the libraries

import os
import cv2
import albumentations as A

2. Set your input/output folders

image_dir = r"C:\datasets\data1\train\images"          # Folder with original images
label_dir = r"C:\datasets\data1\train\labels"           # Folder with YOLO OBB labels
output_image_dir = r"C:\datasets\data2\train\images"    # Folder to save augmented images
output_label_dir = r"C:\datasets\data2\train\labels"    # Folder to save augmented labels

num_aug (used below) controls how many augmented variants are generated per original image.

3. Load and save OBB labels

def load_obb_labels(label_file):
    """Load OBB labels in the format: class_id x1 y1 x2 y2 x3 y3 x4 y4 (normalized)."""
    boxes = []
    with open(label_file, 'r') as f:
        for line in f:
            parts = line.strip().split()
            if len(parts) != 9:  # class_id + 4 corners (x, y)
                continue
            class_id = int(parts[0])
            points = [(float(parts[i]), float(parts[i + 1])) for i in range(1, 9, 2)]
            boxes.append([class_id, points])
    return boxes


def save_obb_labels(label_file, boxes):
    """Save OBB labels back to: class_id x1 y1 x2 y2 x3 y3 x4 y4 (normalized)."""
    with open(label_file, 'w') as f:
        for class_id, points in boxes:
            coords = ' '.join(f"{x:.6f} {y:.6f}" for x, y in points)
            f.write(f"{class_id} {coords}\n")

4. Augment one image and its OBB labels

The OBB corners are treated as keypoints so Albumentations moves them along with the image transform. Points are denormalised to pixel coordinates before augmenting, then normalised back afterwards.

def augment_obb_image(image, bboxes, img_height, img_width):
    # Denormalize the corner points to pixel coordinates
    keypoints = []
    for class_id, points in bboxes:
        denormalized_points = [(x * img_width, y * img_height) for x, y in points]
        keypoints.extend(denormalized_points)

    # Define your augmentation pipeline — tune this to what actually happens in your setup
    transform = A.Compose([
        A.HorizontalFlip(p=0.5),
        A.VerticalFlip(p=0.2),
        A.RandomBrightnessContrast(p=0.2),
        A.Rotate(limit=20, p=0.5),
    ], keypoint_params=A.KeypointParams(format='xy', remove_invisible=False))

    augmented = transform(image=image, keypoints=keypoints)

    # Clip keypoints back into the image bounds after augmentation
    def clip_keypoint(keypoint):
        x, y = keypoint
        x = max(0, min(x, img_width - 1))
        y = max(0, min(y, img_height - 1))
        return (x, y)

    clipped_keypoints = [clip_keypoint(kp) for kp in augmented['keypoints']]

    # Regroup the augmented keypoints back into OBBs (4 points each)
    new_bboxes = []
    for i in range(0, len(clipped_keypoints), 4):
        class_id = bboxes[i // 4][0]
        points = clipped_keypoints[i:i + 4]
        normalized_points = [(x / img_width, y / img_height) for x, y in points]
        new_bboxes.append([class_id, normalized_points])

    return augmented['image'], new_bboxes

5. Run augmentation over the whole dataset

def augment_dataset(image_dir, label_dir, output_image_dir, output_label_dir, num_aug=5):
    """
    image_dir / label_dir: source images and YOLO OBB label files.
    output_image_dir / output_label_dir: where augmented copies are written.
    num_aug: number of augmented variants generated per source image.
    """
    os.makedirs(output_image_dir, exist_ok=True)
    os.makedirs(output_label_dir, exist_ok=True)

    for image_file in os.listdir(image_dir):
        if not (image_file.endswith(".jpg") or image_file.endswith(".png")):
            continue

        image_path = os.path.join(image_dir, image_file)
        label_path = os.path.join(label_dir, os.path.splitext(image_file)[0] + ".txt")

        image = cv2.imread(image_path)
        img_height, img_width = image.shape[:2]

        bboxes = load_obb_labels(label_path)
        if len(bboxes) == 0:  # Skip images with no annotations
            continue

        for i in range(num_aug):
            augmented_image, augmented_bboxes = augment_obb_image(image, bboxes, img_height, img_width)
            if len(augmented_bboxes) == 0:  # Skip if augmentation left no valid boxes
                continue

            base_name = os.path.splitext(image_file)[0]
            cv2.imwrite(os.path.join(output_image_dir, f"{base_name}_aug_{i}.jpg"), augmented_image)
            save_obb_labels(os.path.join(output_label_dir, f"{base_name}_aug_{i}.txt"), augmented_bboxes)

6. Run it

augment_dataset(image_dir, label_dir, output_image_dir, output_label_dir, num_aug=5)

Processing time depends on how many images and augmentations you generate. Once it finishes, output_image_dir/output_label_dir (data2/train/... in the example paths above) contain your expanded training set, ready to train on (see How to train a YOLO object detection model with Roboflow and Ultralytics).

Common mistakes

Augmenting transforms that don’t make sense for your object. Think about what can realistically happen in your setup before picking transforms. For example, if you need to tell whether an object is upside down or not, don’t flip images horizontally or vertically — that destroys the very distinction you’re trying to teach the model.
Duplicating effort with Ultralytics’ built-in augmentation. Ultralytics already applies a substantial set of augmentations automatically during training (see the Ultralytics augmentation settings docs). A separate pipeline like this one is most useful when you need an augmentation Ultralytics doesn’t support, or you’re training with a framework that has no built-in augmentation at all. If you use both, make sure to disable any built-in augmentation that would hurt your model (e.g. a flip you already disabled here for the reason above).