Albumentations

Efficient Image Augmentation for Machine Learning in Python

Albumentations is a fast, flexible image augmentation library designed for machine learning practitioners working on computer vision tasks. Our aim is simple: provide a comprehensive set of tools that can transform any image to augment your datasets, thereby improving model accuracy and robustness.

Features:

• Wide Range of Augmentations: Supports geometric transforms, color augmentations, flips, rotations, and more, tailored for classification, segmentation, object detection, and working with key points.
• Easy Integration: Designed to easily fit into any machine learning pipeline.
• Performance Optimized: Minimizes CPU/GPU load with efficient implementation.
• Community Driven: Open to contributions and feedback. We evolve with your needs.

from albumentations import (
    HorizontalFlip, Affine, CLAHE, RandomRotate90,
    Transpose, ShiftScaleRotate, Blur, OpticalDistortion, GridDistortion, HueSaturationValue,
    GaussNoise, MotionBlur, MedianBlur,
    RandomBrightnessContrast, Flip, OneOf, Compose
)
import numpy as np

def strong_aug(p=0.5):
    return Compose([
        RandomRotate90(),
        Flip(),
        Transpose(),
        GaussNoise(),
        OneOf([
            MotionBlur(p=0.2),
            MedianBlur(blur_limit=3, p=0.1),
            Blur(blur_limit=3, p=0.1),
        ], p=0.2),
        Affine(translate_percent=0.0625, scale=(0.8, 1.2), rotate_limit=(-45, 45), p=0.2),
        OneOf([
            OpticalDistortion(p=0.3),
            GridDistortion(p=0.1)
        ], p=0.2),
        OneOf([
            CLAHE(clip_limit=2),
            RandomBrightnessContrast(),
        ], p=0.3),
        HueSaturationValue(p=0.3),
    ], p=p)

image = np.ones((300, 300, 3), dtype=np.uint8)
mask = np.ones((300, 300), dtype=np.uint8)
whatever_data = "my name"
augmentation = strong_aug(p=0.9)
data = {"image": image, "mask": mask, "whatever_data": whatever_data, "additional": "hello"}
augmented = augmentation(**data)
image, mask, whatever_data, additional = augmented["image"], augmented["mask"], augmented["whatever_data"], augmented["additional"]