2

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]
)
Mr K.
  • 1,064
  • 3
  • 19
  • 22
pkj
  • 559
  • 1
  • 9
  • 21

3 Answers3

2

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)
Gerry P
  • 7,662
  • 3
  • 10
  • 20
1

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)
buddemat
  • 4,552
  • 14
  • 29
  • 49
SvenWarnke
  • 11
  • 1
0

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))
Innat
  • 16,113
  • 6
  • 53
  • 101