Tensorflow data augmentation

Question

I want to convert this keras data augmentation workflow:

datagen = ImageDataGenerator( 
    rescale=1./255,
    rotation_range = 10,
    horizontal_flip = True,
    width_shift_range=0.1,
    height_shift_range=0.1,
    fill_mode = 'nearest')

here is a code snippet but both functions does not work because It does not support batch dimensions!

import numpy as np
def augment(x, y):
    x = tf.keras.preprocessing.image.random_shift(x, 0.1, 0.1)
    x = tf.keras.preprocessing.image.random_rotation(
    x, 10, row_axis=1, col_axis=2, channel_axis=0, fill_mode='nearest', cval=0.0,
    interpolation_order=1)
    return x, y

X = np.random.random(size=(256, 48, 48, 1))
y = np.random.randint(0, 7, size=(256,))
dataset = tf.data.Dataset.from_tensor_slices((X, y))
dataset = dataset.map(augment)
dataset = dataset.batch(16, drop_remainder=False)
dataset = dataset.prefetch(buffer_size=1)

akensert akensert · Accepted Answer · 2020-08-07T20:06:35

I get the following error from running your code: AttributeError: 'Tensor' object has no attribute 'ndim'. It doesn't seem possible to run the augment function with tf.data.Dataset because it can't deal with the Tensor(s). A workaround is to wrap your augment function in tf.py_function:

import tensorflow as tf
import numpy as np

def augment(x, y):
    x = x.numpy()
    x = tf.keras.preprocessing.image.random_shift(x, 0.1, 0.1)
    x = tf.keras.preprocessing.image.random_rotation(
    x, 10, row_axis=1, col_axis=2, channel_axis=0, fill_mode='nearest', cval=0.0,
    interpolation_order=1)
    return x, y

X = np.random.random(size=(256, 48, 48, 1))
y = np.random.randint(0, 7, size=(256,))

dataset = tf.data.Dataset.from_tensor_slices((X, y))
dataset = dataset.map(
    lambda x, y: tf.py_function(
        func=augment,
        inp=[x, y],
        Tout=[tf.float32, tf.int64]))
dataset = dataset.batch(16, drop_remainder=False)
dataset = dataset.prefetch(buffer_size=1)

The code above should run without any errors. If you often need to wrap your functions with tf.py_function, it can be convenient (and also clean) to write a decorator instead. Something like this:

import tensorflow as tf
import numpy as np

def map_decorator(func):
    def wrapper(*args):
        return tf.py_function(
            func=func,
            inp=[*args],
            Tout=[a.dtype for a in args])
    return wrapper

@map_decorator
def augment(x, y):
    x = x.numpy()
    x = tf.keras.preprocessing.image.random_shift(x, 0.1, 0.1)
    x = tf.keras.preprocessing.image.random_rotation(
    x, 10, row_axis=1, col_axis=2, channel_axis=0, fill_mode='nearest', cval=0.0,
    interpolation_order=1)
    return x, y

X = np.random.random(size=(256, 48, 48, 1))
y = np.random.randint(0, 7, size=(256,))

dataset = tf.data.Dataset.from_tensor_slices((X, y))
dataset = dataset.map(augment)
dataset = dataset.batch(16, drop_remainder=False)
dataset = dataset.prefetch(buffer_size=1)

Hope it helps!

Tensorflow data augmentation

1 Answers