Python tips and tricks - 9: Performing additional tasks during data augmentation in keras

For example scaling inputs, performing preprocessing operations, converting masks to categorical, etc.

Code snippet from the video...

#Libraries
import segmentation_models as sm
from sklearn.preprocessing import MinMaxScaler
scaler = MinMaxScaler()
from keras.utils import to_categorical

#Some scaling operation to be applied to images
from sklearn.preprocessing import MinMaxScaler
scaler = MinMaxScaler()
from keras.utils import to_categorical

#Some preprocessing operation on images
BACKBONE = 'resnet34'
preprocess_input = sm.get_preprocessing(BACKBONE)

#Define a function to perform additional preprocessing after datagen.
#For example, scale images, convert masks to categorical, etc.

def preprocess_data(img, mask, num_class):
#Scale images
img = scaler.fit_transform(img.reshape(-1, img.shape[-1])).reshape(img.shape)
img = preprocess_input(img) #Preprocess based on the pretrained backbone...
#Convert mask to one-hot
mask = to_categorical(mask, num_class)

return (img,mask)

#Define the generator.
#We are not doing any rotation or zoom to make sure mask values are not interpolated.
#It is important to keep pixel values in mask as 0, 1, 2, 3, .....
from tensorflow.keras.preprocessing.image import ImageDataGenerator
def trainGenerator(train_img_path, train_mask_path, num_class):

img_data_gen_args = dict(horizontal_flip=True,
vertical_flip=True,
fill_mode='reflect')

image_datagen = ImageDataGenerator(**img_data_gen_args)
mask_datagen = ImageDataGenerator(**img_data_gen_args)

image_generator = image_datagen.flow_from_directory(
train_img_path,
class_mode = None,
batch_size = batch_size,
seed = seed)

mask_generator = mask_datagen.flow_from_directory(
train_mask_path,
class_mode = None,
color_mode = 'grayscale',
batch_size = batch_size,
seed = seed)

train_generator = zip(image_generator, mask_generator)

for (img, mask) in train_generator:
img, mask = preprocess_data(img, mask, num_class)
yield (img, mask)

train_img_path = "data/data_for_keras_aug/train_images/"
train_mask_path = "data/data_for_keras_aug/train_masks/"
train_img_gen = trainGenerator(train_img_path, train_mask_path, num_class=4)

#Make sure the generator is working and that images and masks are indeed lined up.
#Verify generator.... In python 3 next() is renamed as __next__()
x, y = train_img_gen.__next__()

for i in range(0,3):
image = x[i]
mask = np.argmax(y[i], axis=2)
plt.subplot(1,2,1)
plt.imshow(image)
plt.subplot(1,2,2)
plt.imshow(mask, cmap='gray')
plt.show()