An example of Keras preprocessing layers in TF2

keras_preprocessing_layers.py

def augment_images(x, config):
    if config['use_contrast'] == "True":
        x = tf.keras.layers.experimental.preprocessing.RandomContrast(
            config['contrast_factor']
        )(x)

    if config['use_rotation'] == "True":
        x = tf.keras.layers.experimental.preprocessing.RandomRotation(
            config['rotation_factor']
        )(x)
    
    if config['use_flip'] == "True":
        x = tf.keras.layers.experimental.preprocessing.RandomFlip(
            config['flip_mode']
        )(x)

    return x