Augment a tf.data.Dataset with has image and mask. Both Needs to be augmented

Question

I have a tf.data.Dataset of image paths of image and masks

# Creating list of image and mask path
all_val_img = np.array(sorted([os.path.join(VAL_DIR,i) for i in os.listdir(VAL_DIR)]))
all_val_mask = np.array(sorted([os.path.join(VAL_MASK_DIR,i) for i in os.listdir(VAL_MASK_DIR)]))

# doing tf.data.Dataset
val_data = tf.data.Dataset.from_tensor_slices((all_val_img, all_val_mask))

mapping a function to turn paths into images

#mapping the function
val_data = val_data.map(make_image)

Now How do i Augment it . I do not want to use ImageDataGenerator as tf says Deprecated: tf.keras.preprocessing.image.ImageDataGenerator is not recommended for new code.

Note: Both Image and Mask needs to be augmented the same way.

Answer 1

You can define your agumented layers for both images and masks as follows, notice for each training paris (image and mask), you create two instance of same processing layers with same seed value.

class Augment(tf.keras.layers.Layer):
  def __init__(self, seed=42):
    super().__init__()
    # both use the same seed, so they'll make the same random changes.
    self.augment_inputs = keras.layers.RandomFlip(mode="horizontal", seed=seed)
    self.augment_labels = keras.layers.RandomFlip(mode="horizontal", seed=seed)

  def call(self, inputs, labels):
    inputs = self.augment_inputs(inputs)
    labels = self.augment_labels(labels)
    return inputs, labels

Then you can create your tf.data API as usual with above.

train_batches = (
    train_images
    .shuffle(BUFFER_SIZE)
    .batch(BATCH_SIZE)
    .map(Augment())
    .prefetch(-1)

Check this details tutorials. It explains nicely.