Issue with Binary Crossentropy Loss during Training with Custom Data Generator in Keras

Question

I am encountering a problem while training my neural network using a custom data generator in Keras. I have images and corresponding ground truth masks for a binary segmentation task. However, I am consistently facing the following error:

ValueError: in user code:
  
File "/usr/local/lib/python3.10/dist-packages/keras/src/losses.py", line 2532, in binary_crossentropy
    backend.binary_crossentropy(y_true, y_pred, from_logits=from_logits),
File "/usr/local/lib/python3.10/dist-packages/keras/src/backend.py", line 5822, in binary_crossentropy
    return tf.nn.sigmoid_cross_entropy_with_logits(

ValueError: `logits` and `labels` must have the same shape, received ((None, 1) vs (None, None, None, None, None)).

Here is my data generator code:

def custom_data_generator(directory, batch_size=32, input_size=(256, 256), is_training=True):
    image_folder = os.path.join(directory, 'image')
    mask_folder = os.path.join(directory, 'gt')

    filenames = os.listdir(image_folder)
    num_samples = len(filenames)
    image_size = (256, 256)
    while True:
        np.random.shuffle(filenames)

        for i in range(0, num_samples, batch_size):
            batch_filenames = filenames[i:i+batch_size]
            batch_images = []
            batch_masks = []

            for filename in batch_filenames:
                image_path = os.path.join(image_folder, filename)
                mask_path = os.path.join(mask_folder, filename.replace('.jpg', '.png'))

                # Load and preprocess image
                image = cv2.imread(image_path)
                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                image = cv2.resize(image, (input_size[1], input_size[0]))
                image = image / 255.
                batch_images.append(image)

                # Load and preprocess mask
                mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
                mask = np.reshape(mask, (*image_size, 1))  # Resize width and height
                mask = mask < 20
                batch_masks.append(mask)

            x_batch = np.array(batch_images)
            y_batch = np.expand_dims(np.array(batch_masks), axis=-1)  # Add channel dimension

            yield x_batch, y_batch

train_generator = custom_data_generator(train_directory, batch_size=8, input_size=(256, 256, 3), is_training=True)
validation_generator = custom_data_generator(validation_directory, batch_size=8, input_size=(256, 256, 3), is_training=False)

And the model creation and training code:

def create_binary_deeplabv3_model(input_shape):

    base_model = tf.keras.applications.MobileNetV2(input_shape=input_shape, include_top=False, weights=None)

    x = base_model.output
    x = layers.Conv2D(256, (1, 1), activation='relu', padding='same', name='aspp0')(x)
    x = layers.Conv2D(256, (3, 3), activation='relu', padding='same', dilation_rate=6, name='aspp1')(x)
    x = layers.Conv2D(256, (3, 3), activation='relu', padding='same', dilation_rate=12, name='aspp2')(x)
    x = layers.Conv2D(256, (3, 3), activation='relu', padding='same', dilation_rate=18, name='aspp3')(x)
    x = layers.GlobalAveragePooling2D()(x)

    # Flatten and add dense layers for binary classification
    x = layers.Flatten()(x)
    x = layers.Dense(256, activation='relu')(x)
    x = layers.Dropout(0.5)(x)

    x = layers.Flatten()(x)
    
    output = layers.Dense(1, activation='sigmoid')(x)

    model = Model(inputs=base_model.input, outputs=output)
    
    return model


binary_deeplab_model = create_binary_deeplabv3_model(input_shape=(256, 256, 3))
binary_deeplab_model.summary()

binary_deeplab_model.compile(
    optimizer='adam',
    loss='binary_crossentropy',
    metrics=['accuracy']
)

binary_deeplab_model.fit(
    train_generator,
    epochs=10,
    validation_data=validation_generator
)