I am training a neural network that takes an input (H, W, 3) and has the output of size (H', W', C). Now, to augment my dataset, since I only have 45k images, I am using the following in my custom data generator

def Generator():



Since I am rotating my input images and the output_masks are generated from information about the input (specifically, heat maps around the joint locations) do I also need to rotate the masks as well?


Yes! This is crucial.

If you rotate your input images for segmentation, you need to rotate the output masks as well. Otherwise the loss of your network will not be correctly calculated and your network will not learn how to generalize to rotated input images.

If you use keras, you can use two ImageDataGenerator classes, one for the images and one for the masks, with the same random seed and augmentation parameters. It looks something like this:

data_gen_args = dict(rotation_range=45)

image_datagen = ImageDataGenerator(**data_gen_args)
mask_datagen = ImageDataGenerator(**data_gen_args)

image_generator = image_datagen.flow(train,seed=SEED)
mask_generator = mask_datagen.flow(y_train,seed=SEED)
train_generator = zip(image_generator, mask_generator)

model.fit_generator(train_generator, ...)
  • $\begingroup$ This is great thank you! So what I ended up doing was resizing the masks and concatenating the images with the masks and then tf.keras.preprocessing.images.random_shear((images,masks),*args) and then resizing the masks after. Would setting the same seed be sufficient to make sure the rotations or shearing agreed on both. Lastly, is there anyway to use ImageDataGenerator with a custom generator also? $\endgroup$
    – ADA
    Jul 10 '20 at 18:29

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.