Agumenting data in keras using albumentations - tensorflow

Question

I am trying to train a keras ResNet50 model for image classification model using a tutorial. Instead of the inbuilt data generator, I want to use albumentations library for augmentation.
from albumentations import Compose
transforms = Compose([HorizontalFlip()])
I have read a few articles, but I could not figure out how to implement albumentations.
Which line of code should I modify to implement albumentations.
I am reproducing the code below after removing non necessary lines.
NUM_CLASSES = 2
CHANNELS = 3
IMAGE_RESIZE = 224
RESNET50_POOLING_AVERAGE = 'avg'
DENSE_LAYER_ACTIVATION = 'softmax'
OBJECTIVE_FUNCTION = 'categorical_crossentropy'
LOSS_METRICS = ['accuracy']
NUM_EPOCHS = 300
EARLY_STOP_PATIENCE = 20
STEPS_PER_EPOCH_TRAINING = 20
STEPS_PER_EPOCH_VALIDATION = 20
BATCH_SIZE_TRAINING = 10
BATCH_SIZE_VALIDATION = 10
# %% ---------------------------------------------------------------------
TrainingData_directory = 'C:/datafolder/Train'
ValidationData_directory = 'C:/datafolder/Validation'
ModelCheckpointPath = 'C:/datafolder/ResNet50_Weights.hdf5'
# %% ---------------------------------------------------------------------
from albumentations import Compose
import tensorflow as tf
from tensorflow.keras.applications import ResNet50
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
# %% ---------------------------------------------------------------------
model = Sequential()
model.add(ResNet50(include_top = False, pooling = RESNET50_POOLING_AVERAGE, weights = 'imagenet'))
model.add(Dense(NUM_CLASSES, activation = DENSE_LAYER_ACTIVATION))
model.layers[0].trainable = False
from tensorflow.keras import optimizers
sgd = optimizers.SGD(lr = 0.001, decay = 1e-6, momentum = 0.9, nesterov = True)
model.compile(optimizer = sgd, loss = OBJECTIVE_FUNCTION, metrics = LOSS_METRICS)
from keras.applications.resnet50 import preprocess_input
from keras.preprocessing.image import ImageDataGenerator
image_size = IMAGE_RESIZE
data_generator = ImageDataGenerator(preprocessing_function = preprocess_input)
train_generator = data_generator.flow_from_directory(TrainingData_directory,
target_size = (image_size, image_size),
batch_size = BATCH_SIZE_TRAINING,
class_mode = 'categorical')
validation_generator = data_generator.flow_from_directory(ValidationData_directory,
target_size = (image_size, image_size),
batch_size = BATCH_SIZE_VALIDATION,
class_mode = 'categorical')
from tensorflow.python.keras.callbacks import EarlyStopping, ModelCheckpoint
cb_early_stopper = EarlyStopping(monitor = 'val_loss', patience = EARLY_STOP_PATIENCE)
cb_checkpointer = ModelCheckpoint(filepath = ModelCheckpointPath,
monitor = 'val_loss', save_best_only = True, mode = 'auto')
fit_history = model.fit_generator(
train_generator,
steps_per_epoch=STEPS_PER_EPOCH_TRAINING,
epochs = NUM_EPOCHS,
validation_data=validation_generator,
validation_steps=STEPS_PER_EPOCH_VALIDATION,
callbacks=[cb_checkpointer, cb_early_stopper]
)

Answer 1

I think you can do it using the ImageDataGenerator preprocessing_function. The function should taken in a single image as input and return an image. So in your case.
def augmentor (img)
# place you code here do to the albumentations transforms
# your code should result in a single transformed image I called aug_img
return aug_img/127.5-1 #scales the pixels between -1 and +1 which it what preprocees_input does
data_generator = ImageDataGenerator(preprocessing_function = augmentor)

Answer 2

You can include it into the preprocessing function passed to ImageDataGenerator:
def preprocessing_function(x):
preprocessed_x = preprocess_input(x)
transformed_image = transforms(image=preprocessed_x)['image']
return transformed_image
ImageDataGenerator(preprocessing_function = preprocessing_function)

Answer 3

That's (IMO) the limitation or losing the flexibility that one might come across using a built-in data generator (ImageDataGenerator). You should implement your own custom data generator.
Check this kernel: [TF.Keras]: SOTA Augmentation in Sequence Generator, where we've shown how one can use albumentation, cutmix, mixup, and fmix type advance augmentation into the custom generator. Here is a basic approach of how to use albumentaiton in a custom data generator.
import albumentations as A
# For Training
def albu_transforms_train(data_resize):
return A.Compose([
A.ToFloat(),
A.Resize(data_resize, data_resize),
A. [.....what ever......]
], p=1.)
class Generator(tf.keras.utils.Sequence):
def __getitem__(self, index):
...........
Data = np.empty((self.batch_size, *self.dim))
Target = np.empty((self.batch_size, 5), dtype = np.float32)
for i, k in enumerate(idx):
# load the image file using cv2
image = cv2.imread(self.img_path + self.data['image_id'][k])
image = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)
# call augmentor / albumentation
res = self.augment(image=image)
image = res['image']
# assign
Data[i,:, :, :] = image
Target[i,:] = self.label.loc[k, :].values
return Data, Target
# call the generator
check_gens = Generator(...., transform = albu_transforms_train(128))