ImageDataGenerator preparing labels to semantic segmentation

[mgajewski73]

Hi,

I prepare data for semantic segmentation and have a problem. I use Keras Image Generator. I had no problem with binary segmentation. But now I have 7 classes, the images are fine, but the labels are returned in shape [no classes, w, h, 1 or 3] (grayscale or rgb).
To creating labels I used Pixel Annotation tool.
My labels have shape [w,h,1] range of values 0-7

My code:

import tensorflow as tf
batch = 8
datagen_params = dict(
    width_shift_range=0.05,
    height_shift_range=0.05,
    zoom_range=(0.95,1.05),
    rotation_range=3,
    validation_split=0.15
)

images_datagen = tf.keras.preprocessing.image.ImageDataGenerator(
    rescale = 1 / 255,
    **datagen_params
)

masks_datagen = tf.keras.preprocessing.image.ImageDataGenerator(
    # preprocessing_function=my_preprocessing_function,
    **datagen_params
)

train_images_generator = images_datagen.flow_from_directory(
    'path/Images/images',
    color_mode = 'rgb',   
    target_size=(480, 640),
    batch_size=batch,
    class_mode=None,
    shuffle=True,
    seed=42, 
    subset='training'
)

train_masks_generator = masks_datagen.flow_from_directory(
    'path/Images/watershed_mask',    
    # color_mode = 'grayscale',    
    target_size=(480, 640),
    batch_size=batch,
    class_mode='categorical',
    shuffle=True,
    seed=42,
    subset='training'
)

val_images_generator = images_datagen.flow_from_directory(
    'path/Images/images',  
    color_mode = 'rgb',    
    target_size=(480, 640),
    batch_size=batch,
    class_mode='categorical',
    shuffle=True,
    seed=42,
    subset='validation'
)

val_masks_generator = masks_datagen.flow_from_directory(
    'path/Images/watershed_mask',
    # color_mode = 'grayscale',    
    target_size=(480, 640),
    batch_size=batch, 
    class_mode=None,
    shuffle=True,
    seed=42,
    subset='validation'
)

train_combined_generator = (pair for pair in zip(train_images_generator, train_masks_generator))
val_combined_generator = (pair for pair in zip(val_images_generator, val_masks_generator))

Return:

Found 929 images belonging to 1 classes. 
Found 929 images belonging to 1 classes. 
Found 163 images belonging to 1 classes. 
Found 163 images belonging to 1 classes

I check shape of image, or one layer of labels using this code:

joined_generator = train_combined_generator
images, masks = next(joined_generator)

print(masks[0].shape, images[0].shape)

Return:
(8, 480, 640, 3) (480, 640, 3)

I expect a shape (480, 640, 8) (480, 640, 3).
What am I doing wrong, how can I get this label?

Further code:

model = segmentation_models.Unet(backbone_name='efficientnetb1', input_shape=(480, 640, 3), classes=8)
model.summary()

training_samples = train_images_generator.n
validation_samples = val_images_generator.n

model.compile(optimizer='adam',
              loss='categorical_crossentropy',
              metrics=['accuracy'])

model.fit(
  train_combined_generator,
  steps_per_epoch=training_samples // batch,
  epochs=50,
  validation_data=val_combined_generator, 
  validation_steps=validation_samples // batch
)

[mgajewski73]

I've loaded dataset via tf.data.