Getting Started with Albumentation: Deep Learning Image Augmentation Technique in PyTorch | by Kevin Bok | Towards Data Science

Either you are quietly participating Kaggle Competitions, trying to learn a new cool Python technique, a newbie in data science / deep learning, or just here to grab a piece of codeset you want to copy-paste and try right away, I guarantee this post would be very helpful.
In this post, I will make a thorough introduction of one of the most widely used (2021 and ongoing) image augmentation libraries, Albumentation, with the example Python code and output. You should be able to establish a basic understanding of Albumentation from this post and eventually trying it out in your own workspace. Anyone who does not (or minimal) have previous experience in Pytorch or Python or exposure to the Image Augmentation technique itself is also welcome.
I personally found the library very useful in my day-to-day usage for any computer-vision-related tasks, so I thought, ‘ Why don’t I just make a post about it?
Disclaimer: This writeup is not sponsored by or affiliated with Albumentation.
So, what actually is Albumentation? What are the features and advantages of using it? And What are the real-world applications?
Before you get to know what Albumentation helps, it’s critical to understand what Image Augmentation means in Computer Vision.
Ryan Allred’s post has a great explanation of it.
Deep Neural Networks, particularly Convolutional Neural Networks (CNNs), are particularly proficient at image classification tasks. State-of-the-art CNNs have even been shown to exceed human performance in image recognition (...) Image Augmentation is the process of taking images that are already in a training dataset and manipulating them to create many altered versions of the same image. This both provides more images to train on, but can also help expose our classifier to a wider variety of lighting and coloring situations so as to make our classifier more robust
Basically, it custom-adds all sorts of different varieties to the given images/pictures, in order to increase the size of the training dataset, eventually help to improve the accuracy of the deep learning model.
Albumentation is a tool that can customize [elastic, grid, motion blur, shift, scale, rotate, transpose, contrast, brightness, etc] to the images/pictures before you slot those into the model.
Official Albumentation website describes itself as
Albumentations is a Python library for fast and flexible image augmentations. Albumentations efficiently implements a rich variety of image transform operations that are optimized for performance, and does so while providing a concise, yet powerful image augmentation interface for different computer vision tasks, including object classification, segmentation, and detection.
Albumentations has been officially published with its title Albumentations: Fast and Flexible Image Augmentations in 2020 to the Infomation Journal, and at this moment it is maintained by 5 core team members from Russia, with consistent feature updates.
What Albumentation differentiates itself from other image augmentation related packages is that,
IT IS FAST
The package has been optimized by several OpenCV-based libraries (including NumPy, OpenCV, imgaug).
Now you get a chance to face the final boss without beating up all mid-bosses
My go-to python framework for deep learning has been Pytorch, so I have been initially exposed to the usage of torchvision.transforms that are natively offered by TorchVision. Torchvision.transforms already gives pretty solid custom augmentation methods and documentation, so I have been stick to its offerings.
While searching for a better augmentation open-source library, I discovered this Albumentation package. And I found it very fast, highly customizable, and most importantly, it only takes <5 min to refactor the Torchvision code to Albumentation.
The above chart is the test result with Intel Xeon Platinum 8168 CPU by 2000 validation set images in ImageNet. Values in each cell indicate the number of images processed in a single core. You can see that Albumentation beats all other libraries at least 2X for many transforms.
You can find more about the benchmark with the latest 0.5.0 release in Albumentation official GitHub.
For the scope of the project, I will introduce key components and usages in Albumentation. The codeset is mostly based on tutorial notebooks from the Albumentations team. I have referred to the following notebook:
migrating_from_torchvision_to_albumentations.ipynb
I persistently use Google Colab for easy/sharable notebook prototypes. Their Jupyter is free and you can use free GPU!
You can upload your photos to your Google Drive and mount the Google Drive to Colab with the following codes.
from google.colab import drivedrive.mount("/content/gdrive")
This cell will return the following Go to this URL in a browser. After you click the URL you can retrieve the authorization code. Copy-paste this code and hit enter and you are good to go!
Also, I will import all other Python libraries that are needed for this tutorial
from PIL import Imageimport timeimport torchimport torchvisionfrom torch.utils.data import Datasetfrom torchvision import transformsimport albumentationsimport albumentations.pytorchfrom matplotlib import pyplot as pltimport cv2import numpy as np
For the demonstration purpose, I have used a street photo in Venice, Italy beautiful that I took during my trip to Europe a while ago.
I normally create a Dataloader to process image data pipelines using PyTorch and Torchvision. In the below code, it
Creates a simple Pytorch Dataset class
Calls an image and do a transformation
Measure the whole processing time with 100 loops
First, get Dataset abstract class from torch.utils.data, and crates a TorchVision Dataset Class. Then I slot in the image and do transformation using the __getitem__ method. Also, I used total_time = (time.time() - start_t for measuring how long it takes
class TorchvisionDataset(Dataset):    def __init__(self, file_paths, labels, transform=None):        self.file_paths = file_paths        self.labels = labels        self.transform = transform            def __len__(self):        return len(self.file_paths)    def __getitem__(self, idx):        label = self.labels[idx]        file_path = self.file_paths[idx]                # Read an image with PIL        image = Image.open(file_path)                start_t = time.time()        if self.transform:            image = self.transform(image)        total_time = (time.time() - start_t)        return image, label, total_time
Then we resize the image to 256x256 (height * weight) and do a random crop to the size 224x224. Then apply Horizontal flip with 50% probability and convert it to Tensor. The input file path should be the path of Google Drive where your images are in.
torchvision_transform = transforms.Compose([    transforms.Resize((256, 256)),     transforms.RandomCrop(224),    transforms.RandomHorizontalFlip(),    transforms.ToTensor(),])torchvision_dataset = TorchvisionDataset(    file_paths=["/content/drive/MyDrive/img5.jpg"],    labels=[1],    transform=torchvision_transform,)
Then we calculate the time taken to pull a sample image from torchvision_dataset and transform it, followed by running 100 loops to check the average millisecond it takes.
total_time = 0for i in range(100):  sample, _, transform_time = torchvision_dataset[0]  total_time += transform_timeprint("torchvision time/sample: {} ms".format(total_time*10))plt.figure(figsize=(10, 10))plt.imshow(transforms.ToPILImage()(sample))plt.show()
A hundred loops of Resize + RandomCrop + RandomHorizontalFlip took about 40 ms in my Colab environment, and the last image takes 224x224 as you can see. You can also see that the Flip happened for the 100th image!
Now I will refactor the data pipeline from TorchVision to Albumentation
Similar to TorchVision, we create an Albumentations Dataset class.
class AlbumentationsDataset(Dataset):    """__init__ and __len__ functions are the same as in TorchvisionDataset"""    def __init__(self, file_paths, labels, transform=None):        self.file_paths = file_paths        self.labels = labels        self.transform = transform            def __len__(self):        return len(self.file_paths)    def __getitem__(self, idx):        label = self.labels[idx]        file_path = self.file_paths[idx]                # Read an image with OpenCV        image = cv2.imread(file_path)                # By default OpenCV uses BGR color space for color images,        # so we need to convert the image to RGB color space.        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)        start_t = time.time()        if self.transform:            augmented = self.transform(image=image)             image = augmented['image']	    total_time = (time.time() - start_t)        return image, label, total_time
Now create a transform in Albumentation. A small syntax difference you can spot in this example is that Torchvision’s RandomHorizontalFlip() produces the same result by HorizontalFlip() in Albumentation
"""torchvision_transform = transforms.Compose([    transforms.Resize((256, 256)),     transforms.RandomCrop(224),    transforms.RandomHorizontalFlip(),    transforms.ToTensor(),])"""# Same transform with torchvision_transformalbumentations_transform = albumentations.Compose([    albumentations.Resize(256, 256),     albumentations.RandomCrop(224, 224),    albumentations.HorizontalFlip(), # Same with transforms.RandomHorizontalFlip()    albumentations.pytorch.transforms.ToTensor()])
Perform the same image transformation, get the average time taken, and visualize the result.
# Same dataset with torchvision_datasetalbumentations_dataset = AlbumentationsDataset(    file_paths=["/content/drive/MyDrive/img5.jpg"],    labels=[1],    transform=albumentations_transform,)total_time = 0for i in range(100):  sample, _, transform_time = albumentations_dataset[0]  total_time += transform_timeprint("albumentations time/sample: {} ms".format(total_time*10))plt.figure(figsize=(10, 10))plt.imshow(transforms.ToPILImage()(sample))plt.show()
Surprisingly, it only took 1.77ms, which is about 23x faster than with Torchvision! It even has a wider computation gap than what we saw from the official benchmark documentation, and also the cropped area seems to have a slight difference.
As you saw, it’s super-fast, and speed matters.
In code competitions, if you can reduce the computation bottleneck in image processing, more resources can be used in other parts (model fitting, hyper-parameter tuning, etc).
In real-world industries, the influx of new images in the production database might come in high volume (i.e. 1000 images per second). Imagine developing a real-time deep learning model. A lightning-fast method to process the image is critical in the model, which might affect the user experience and eventually revenues/profits.
For your learning purpose, it is helpful to be on the verge of the latest and best-performing tech skills, which could be useful when building your own project, and maybe in job-searching (i.e. resume)
Lastly, I will show how to perform augmentation with OneOf function, which I personally felt most complex but useful in Albumentation
albumentations_transform_oneof = albumentations.Compose([    albumentations.Resize(256, 256),     albumentations.RandomCrop(224, 224),    albumentations.OneOf([                          albumentations.HorizontalFlip(p=1),                          albumentations.RandomRotate90(p=1),                          albumentations.VerticalFlip(p=1)                ], p=1),    albumentations.OneOf([                          albumentations.MotionBlur(p=1),                          albumentations.OpticalDistortion(p=1),                          albumentations.GaussNoise(p=1)                     ], p=1),    albumentations.pytorch.ToTensor()])
It has the same code in Resize and Random Crop in the previous codeset. OneOf randomly takes one of the listed transforms inside the bracket. We can even put a probability of happening in the function itself. For example, if one of([...], p=0.5), it skips the whole transforms in 50% chance, and randomly takes one of three transforms with the 1/6 chance.
I let it randomly choose from Horizontal Flip, Rotation, Vertical Flip, and again let it randomly choose from Blur, Distortion, Noise. So in this case, we allow 3x3 = 9 combinations.
albumentations_dataset = AlbumentationsDataset(    file_paths=["/content/gdrive/My Drive/img5.png"],    labels=[1],    transform=albumentations_transform_oneof,)num_samples = 5fig, ax = plt.subplots(1, num_samples, figsize=(25, 5))for i in range(num_samples):  ax[i].imshow(transforms.ToPILImage()(albumentations_dataset[0][0]))  ax[i].axis('off')
Five samples from the defined image transformation are as follows:
We see that rotation 360, horizontal flip, vertical flip are applied, and different distortion is applied to all images under the hood.
In addition to what I have explained, there are numerous examples available. Please refer to the official Albumentation website or Albumentation Github to apply the best suitable augmentation to your own needs!
*Note* You can also find the Jupyter Notebooks that contains the real-world usage of the Albumentation: Kaggle Competition
Bengali.AI Handwritten Grapheme Classification was a code competition to classify three constituent elements in the image, given the image of a handwritten Bengali grapheme. It was an image classification problem, and Albumentation is used to transform the tens of thousands of images that originally exist in the Parquet file, with GBs of size.
GitHub: Bengali.AI Handwritten Grapheme Classification Competition
To summarize, I gave an introduction to the Image Augmentation technique, the Albumentation library in Python, with example codes in the tutorial. The experiment was only performed with a single image for the sake of the project, but a big speed improvement could be spotted. Albumentation offers a variety of transforms, so I highly recommend my readers to utilize it from today.
For more details on learning other data science topics, below Github repositories will also be helpful
AWS SageMaker in Production End-to-End examples that show how to solve business problems using Amazon SageMaker and its ML/DL algorithm.
PySparkFunctions and utilities with Real-world Data examples. Can be used to build a complete ETL process of data modeling
Recommendation SystemProduction-level Implementations of Recommender System in Pytorch. Clone repo and start training by running ‘main.py’
Natural Language Processing (NLP)Full implementation examples of several Natural Language Processing methods in Python. Ordered in a level of complexity in learning
HyunJoon is a Data Geek with a degree in Statistics. He enjoys sharing data science/analytics knowledge. Message him on LinkedIn.
Reference:(1) Ryan Allred, https://towardsdatascience.com/aws-glue-and-you-e2e4322f0805(2) Hoya012, https://hoya012.github.io/blog/albumentation_tutorial/(3) Albumentation Official Docs, https://albumentations.ai/docs/