Understand how to use Albumentations with Keras

user_7455917 0
Viewed 437

I am trying to understand how to build a data augmentation pipeline with Albumentations to feed a Keras model. I am following this example-> https://albumentations.ai/docs/examples/tensorflow-example/ where they create a dataset object PrefetchDataset and pass it to model.fit(). See code below:

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import tensorflow_datasets as tfds
from functools import partial
from albumentations import (
    Compose, RandomBrightness, JpegCompression, HueSaturationValue, RandomContrast, HorizontalFlip,
    Rotate
)
AUTOTUNE = tf.data.experimental.AUTOTUNE

#load data
data, info= tfds.load(name="tf_flowers", split="train", as_supervised=True, with_info=True)
data

# augmentations
transforms = Compose([
            Rotate(limit=40),
            RandomBrightness(limit=0.1),
            JpegCompression(quality_lower=85, quality_upper=100, p=0.5),
            HueSaturationValue(hue_shift_limit=20, sat_shift_limit=30, val_shift_limit=20, p=0.5),
            RandomContrast(limit=0.2, p=0.5),
            HorizontalFlip(),
        ])


def aug_fn(image, img_size):
    data = {"image":image}
    aug_data = transforms(**data)
    aug_img = aug_data["image"]
    aug_img = tf.cast(aug_img/255.0, tf.float32)
    aug_img = tf.image.resize(aug_img, size=[img_size, img_size])
    return aug_img

def process_data(image, label, img_size):
    aug_img = tf.numpy_function(func=aug_fn, inp=[image, img_size], Tout=tf.float32)
    return aug_img, label

# create dataset
ds_alb = data.map(partial(process_data, img_size=120),
                  num_parallel_calls=AUTOTUNE).prefetch(AUTOTUNE)


# ...And then just pass this dataset object to the model

def create_model(input_shape):
    pass
    # define model layers...

model = create_model(input_shape)

model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics='accuracy', run_eagerly=True)
model.fit(ds_alb, epochs=2)

My question is: Does this object ds_alb returns different images (as per the random parameters set) for each batch during training? I have read the code and it seems like all the augmentations in Compose are being performed only once, here:

ds_alb = data.map(partial(process_data, img_size=120), num_parallel_calls=AUTOTUNE).prefetch(AUTOTUNE)

But I believe the purpose of building this pipeline is to feed different augmented images on each pass, not to augment only once... There is something I am not understanding. How can I check if different images are being generated?

albumentations keras tensorflow python
1 Answers

I'm not sure whether the function will be called once or many times by tf.numpy_function per batch. But there is a simple way to test it. Put a print inside aug_fn and take two batches:

def aug_fn(...):
    print("Called")
    ....

batch = next(ds_alb) #depending on the type of generator, you might need ds_alb[0]   
another_batch = next(ds_alb) #or ds_alb[1]

#also check the shape of the batch to make sure there are many images

If "Called" appears many times, it's making one transformation per image. If it appears only twice, it is making one transformation per batch. If it appears only once, then there is a problem in the code.

Honestly, I prefer "one transformation per batch" because this would represent a big gain in performance, which is important. Sometimes preprocessing can be the operation limiting your performance.

But will this not beat the purpose of augmentation?

No! You're feeding many batches to the model, and presumably, although you have the same preprocessing for the entire batch, you have different source images! And your next batch will have a different preprocessing, of course.

Since you will be training for many epochs, of course, the same images will appear again in a later epoch with a different preprocessing. In the end, the number of passes will be so big that using the same preprocessing in a single batch will not be a problem for the model.

user_20972400
Related