How to use a cv2 image augmentation function with tensorflow tf.data.Dataset?

Question

I am using tf.data.Dataset to create my dataset and training a CNN with keras. I need to apply masks on the images, and the mask depends on the shape of the image, there are no predefined pixel coordinates.

When looking for an answer on the internet, I found that there are 2 ways of accessing shapes of images in TensorFlow (in training time):

1. Using eager execution (which is not enabled by default in my case, I'm using tf v 12.0)

2. Using a session

I do not want to use eager execution because it slows down training, and cannot use a session because I train and test the CNN using Keras (I feed the data to model.train() using iterators of tf.data.Dataset).

As a consequence, I have no way of knowing the shapes of images, and thus cannot access specific pixels for data augmentation.

I wrote a function using OpenCV (cv2) that applies the masks. Is there a way to integrate it with the TensorFlow data pipeline?

EDIT : I found a solution. I used tf.py_func to wrap the python functions

Answer 1

You can use map to transform elements of your dataset. You can then use tf.py_function to wrap your cv2 function into a tf op that executes eagerly. In tensorflow 1.x, you may use tf.py_func but the behavior is a bit different. See tf.py_function documentation for more info.

So, in TF-2.x it will look something like:

def cv2_func(image, label):
    # your code goes here

def tf_cv2_func(image, label):
    [image, label] = tf.py_function(cv2_func, [image, label], [tf.float32, tf.float64])
    return image, label

train_ds = train_ds.shuffle(BUFFER_SIZE).map(tf_cv2_func).batch(BATCH_SIZE)