Keras ImageDataGenerator for segmentation with images and masks in separate directories

Question

I am trying to build a semantic segmentation model using tensorflow.keras . The dataset that I am using has the images and masks stored in separate directories and each filename has is an id for mapping an image file with its respective mask.

Following is the structure of my dataset directory:

new
   - rendered_imges
      - render
         - image_1.tif
         - image_2.tif
         - image_3.tif
   - ground_truths
      - masks
         - mask_1.tif
         - mask_2.tif
         - mask_3.tif

In the above directory structure, image_{i}.tif corresponds to mask_{i}.tif .

I tried writing an ImageDataGenerator for augmenting both the images and their respective masks in exactly the same way. My approach was the following:

SEED = 100

image_data_generator = ImageDataGenerator(
    width_shift_range = 0.1,
    height_shift_range = 0.1,
    rotation_range = 10,
    zoom_range = 0.1
).flow_from_directory('./new/rendered_images', batch_size = 16, target_size = (150, 150), seed = SEED)

mask_data_generator = ImageDataGenerator(
    width_shift_range = 0.1,
    height_shift_range = 0.1,
    rotation_range = 10,
    zoom_range = 0.1
).flow_from_directory('./new/ground_truths', batch_size = 16, target_size = (150, 150), seed = SEED)

With the above approach although I am getting the same augmentations applied to both image and mask, the images are not getting paired with their respective masks according to their filenames. It would be great if someone can suggest a way to do this properly using Keras or Tensorflow.

Answer 1

You will need to make a new function that will generate both the training image and corresponding mask that you will use to feed into the fit_generator method. Specifically, the way fit_generator works is that yields a sequence of tuples such that the first element of the tuple is the image and the second element of the tuple is the expected output. By simply using the data generator on its own, the sub-directories will implicitly encode the expected label of the image. This is of course no longer the case when you're trying to do semantic segmentation.

Therefore, create a new function that will output a collection of tuples that give you the image and the mask. In summary, you will simply take the two ImageDataGenerators that you created, zip them together, then have a loop that will yield each batch of training images and expected output labels.

One final thing I need to mention is that the filenames for both directories need to match if you want to successfully pair the image and corresponding mask. For example, if you have a training image called 1.tif in your rendered_imges/render subdirectory, you will need to name your corresponding mask the same way in ground_truths/mask . The reason is that even though you match the seeds, it will randomly choose the filenames before loading the images in memory, so to make sure that the same order of selection is maintained among the training images and corresponding masks, their filenames also need to match. Make sure you do that before we proceed here.

Therefore, do something like this:

def my_image_mask_generator(image_data_generator, mask_data_generator):
    train_generator = zip(image_data_generator, mask_data_generator)
    for (img, mask) in train_generator:
        yield (img, mask)

Next, create your data generators as you did normally:

SEED = 100

image_data_generator = ImageDataGenerator(
    width_shift_range = 0.1,
    height_shift_range = 0.1,
    rotation_range = 10,
    zoom_range = 0.1
).flow_from_directory('./new/rendered_images', batch_size = 16, target_size = (150, 150), seed = SEED)

mask_data_generator = ImageDataGenerator(
    width_shift_range = 0.1,
    height_shift_range = 0.1,
    rotation_range = 10,
    zoom_range = 0.1
).flow_from_directory('./new/ground_truths', batch_size = 16, target_size = (150, 150), seed = SEED)

Finally, call the fit_generator method on your model. Assuming you've already constructed the model correctly:

from keras.optimizers import Adam
# Your other related imports here...

# Create custom generator for training images and masks
my_generator = my_image_mask_generator(image_data_generator, mask_data_generator)

model = ... # Define your model here
# Compile your model here
model.compile(optimizer = Adam(lr = 1e-4), loss = 'binary_crossentropy', metrics = ['accuracy'])

# Train your model here
model.fit_generator(my_generator,...)

Take note that given your directory structure, it looks like you are performing a binary segmentation per image, so that's why I chose binary cross-entropy as the loss function.