Data augmentation in the data loading of tensorflow dataset, tfds, resulting in TypError or AttributeError

Question

I am trying to do some data augmentation but i am not so familiar with tensors. This is the code i started with:

def _random_apply(func, x, p):
  return tf.cond(tf.less(tf.random.uniform([], minval=0, maxval=1, dtype=tf.float32),
                  tf.cast(p, tf.float32)),
        lambda: func(x),
        lambda: x)

def _resize_with_pad(image):
  image = tf.image.resize_with_pad(image, target_height=IMG_S, target_width=IMG_S)   
  return image

def augment(image, label):
  img = _random_apply(tf.image.flip_left_right(image), image, p=0.2)
  img = _random_apply(_resize_with_pad(img), img, p=1)
  return img, label

train_dataset = (
    train_ds
    .shuffle(1000)
    .map(augment, num_parallel_calls=tf.data.AUTOTUNE)
    .prefetch(tf.data.AUTOTUNE)
)

which resulted in the following error.

----> 4     .map(augment, num_parallel_calls=tf.data.AUTOTUNE)

TypeError: 'Tensor' object is not callable

Then i thought maybe it would work if i converted it to numpy.

def _random_apply(func, x, p):
  return tf.cond(tf.less(tf.random.uniform([], minval=0, maxval=1, dtype=tf.float32),
                  tf.cast(p, tf.float32)),
        lambda: func(x),
        lambda: x)

def _resize_with_pad(image):
  image = image.numpy()
  image = tf.image.resize_with_pad(image, target_height=IMG_S, target_width=IMG_S).numpy()  
  return image

def augment(image, label):
  image = image.numpy()
  img = _random_apply(tf.image.flip_left_right(image).numpy(), image, p=0.2)
  img = _random_apply(_resize_with_pad(img), img, p=1)
  return img, label

train_dataset = (
    train_ds
    .shuffle(1000)
    .map(augment, num_parallel_calls=tf.data.AUTOTUNE)
    .prefetch(tf.data.AUTOTUNE)
)

But now i get this error.

----> 4     .map(augment, num_parallel_calls=tf.data.AUTOTUNE)

 AttributeError: 'Tensor' object has no attribute 'numpy'

I tried to do something like in this answer and now i get no error directly but rather in the next block of code:

for image, _ in train_dataset.take(9):
etc

InvalidArgumentError 
----> 1 for image, _ in train_dataset.take(9):

InvalidArgumentError: TypeError: 'tensorflow.python.framework.ops.EagerTensor' object is not callable

Anyone know what I am doing wrong?

Answer 1

In augment, you're passing tensors to _random_apply. tf.image.flip_left_right(image) returns a tensor. Then, in _random_apply, you're using that tensor like it's a function. You need to pass tf.flip_left_right as a callable:

def augment(image):
    img = _random_apply(tf.image.flip_left_right, image, p=0.2)
    img = _random_apply(_resize_with_pad, img, p=1)
    return img

Full working example:

import tensorflow as tf

train_ds = tf.data.Dataset.from_tensor_slices(tf.random.uniform((100, 224, 224, 3)))


def _random_apply(func, x, p):
    return tf.cond(tf.less(tf.random.uniform([], minval=0, maxval=1, dtype=tf.float32),
                           tf.cast(p, tf.float32)),
                   lambda: func(x),
                   lambda: x)


def _resize_with_pad(image):
    image = tf.image.resize_with_pad(image, target_height=200, target_width=200)
    return image


def augment(image):
    img = _random_apply(tf.image.flip_left_right, image, p=0.2)
    img = _random_apply(_resize_with_pad, img, p=1)
    return img


train_dataset = train_ds.map(augment)

batch = next(iter(train_dataset))