[Python] Problem in fine tuning DeepLabV3Plus using keras_cv for Semantic Segmentation

I'm using open-images-v7 dataset (accessing via fiftyone lib) and keras_cv lib to fine tune DeepLabV3Plus with mobilenet_v3_small backbone, but the accuracy doesn't improve with epochs at all, and I'm getting shape warning.

Dataset Preparation Code:

import tensorflow as tf
import numpy as np

def preprocess_sample(sample):
img = tf.io.read_file(sample["filepath"])
img = tf.image.decode_jpeg(img, channels=3)
img.set_shape([None, None, 3])
img = tf.image.resize(img, (512, 512))
# img = img / 255.0 # Normalize

for detection in sample.ground_truth.detections:
if detection.label == 'Vehicle registration plate':
mask = detection.mask.astype(np.float32)
break
mask = tf.expand_dims(mask, axis=-1)
mask.set_shape([None, None, 1])
mask = tf.image.resize(mask, (512, 512), method="nearest")

return img, mask

# Convert FiftyOne samples to TF Dataset
tf_train_dataset = tf.data.Dataset.from_generator(
lambda: (preprocess_sample(s) for s in train_dataset),
output_signature=(
tf.TensorSpec(shape=(512, 512, 3), dtype=tf.float32), # Image
tf.TensorSpec(shape=(512, 512, 1), dtype=tf.float32), # Mask
)
)
tf_val_dataset = tf.data.Dataset.from_generator(
lambda: (preprocess_sample(s) for s in val_dataset),
output_signature=(
tf.TensorSpec(shape=(512, 512, 3), dtype=tf.float32), # Image
tf.TensorSpec(shape=(512, 512, 1), dtype=tf.float32), # Mask
)
)

# Batch and shuffle
tf_train_dataset = tf_train_dataset.batch(8).prefetch(tf.data.AUTOTUNE)
tf_val_dataset = tf_val_dataset.batch(8).prefetch(tf.data.AUTOTUNE)


Fine Tuning Code

model = keras_cv.models.DeepLabV3Plus.from_preset(
"mobilenet_v3_small",
input_shape=(512, 512, 3),
num_classes=1,
# activation=None # Remove final activation
)
# outputs = tf.keras.layers.Activation("sigmoid")(model.output)
# model = tf.keras.Model(inputs=model.input, outputs=outputs)

model.compile(
optimizer=keras.optimizers.Adam(learning_rate=1e-4), # 'adam'
loss="binary_crossentropy",
metrics=["binary_accuracy"] # accuracy
)

# Train
model.fit(
tf_train_dataset,
validation_data=tf_val_dataset,
epochs=5,
callbacks=[
keras.callbacks.EarlyStopping(patience=3, restore_best_weights=True),
],
)


Error/Output

Epoch 1/5

/usr/local/lib/python3.11/dist-packages/keras/src/models/functional.py:237: UserWarning: The structure of `inputs` doesn't match the expected structure.
Expected: ['keras_tensor']
Received: inputs=Tensor(shape=(None, 512, 512, 3))
warnings.warn(msg)
/usr/local/lib/python3.11/dist-packages/keras/src/ops/nn.py:907: UserWarning: You are using a softmax over axis -1 of a tensor of shape (None, 512, 512, 1). This axis has size 1. The softmax operation will always return the value 1, which is likely not what you intended. Did you mean to use a sigmoid instead?
warnings.warn(

13/Unknown 200s 12s/step - binary_accuracy: 0.7506 - loss: 0.6739

/usr/local/lib/python3.11/dist-packages/keras/src/trainers/epoch_iterator.py:151: UserWarning: Your input ran out of data; interrupting training. Make sure that your dataset or generator can generate at least `steps_per_epoch * epochs` batches. You may need to use the `.repeat()` function when building your dataset.
self._interrupted_warning()

13/13 ━━━━━━━━━━━━━━━━━━━━ 211s 13s/step - binary_accuracy: 0.7506 - loss: 0.6725 - val_binary_accuracy: 0.6972 - val_loss: 0.6890
Epoch 2/5
13/13 ━━━━━━━━━━━━━━━━━━━━ 202s 13s/step - binary_accuracy: 0.7506 - loss: 0.5131 - val_binary_accuracy: 0.6972 - val_loss: 0.6835
Epoch 3/5
13/13 ━━━━━━━━━━━━━━━━━━━━ 182s 14s/step - binary_accuracy: 0.7506 - loss: 0.4288 - val_binary_accuracy: 0.6972 - val_loss: 0.6792
Epoch 4/5
13/13 ━━━━━━━━━━━━━━━━━━━━ 166s 13s/step - binary_accuracy: 0.7506 - loss: 0.3606 - val_binary_accuracy: 0.6972 - val_loss: 0.6756
Epoch 5/5
13/13 ━━━━━━━━━━━━━━━━━━━━ 166s 13s/step - binary_accuracy: 0.7506 - loss: 0.3141 - val_binary_accuracy: 0.6972 - val_loss: 0.6723

Continue reading...