Can't Select Activation / Output for DeepNetV3Plus, broken example for 1 class

Current Behavior:

The DeepNetV3Plus Task appears to have issues when there's only 1 class. It should allow setting the activation function and/or default to sigmoid etc. It appears to be a very subtle failure and only prints a warning.

UserWarning: You are using a softmax over axis -1 of a tensor of shape (1, 96, 96, 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?

This is broken in the example code shown on its webpage where it also only uses one class:

https://keras.io/api/keras_cv/models/tasks/deeplab_v3_segmentation/

images = np.ones(shape=(1, 96, 96, 3))
labels = np.zeros(shape=(1, 96, 96, 1))
backbone = keras_cv.models.ResNet50V2Backbone(input_shape=[96, 96, 3])
model = keras_cv.models.segmentation.DeepLabV3Plus(
    num_classes=1, backbone=backbone,
)

Note that other tasks allow changing the activation function for what I believe to be similar situations: https://keras.io/api/keras_cv/models/tasks/image_classifier/

I am interested in an application where I only have one output class, so having a single output of "confidence values" is appealing to me because it would be the most efficient. I did attempt training with 2 classes, but my gut says the internals of DeepLabV3Plus don't quite work well with only two classes, there may be some attempts to optimize this situation, and/or the training masks become too difficult to prepare.

Expected Behavior:

When classes == 1, the output should switch to a sigmoid or other function where a confidence threshold (typically 0.5) can be set to determine if class is of interest (1) or background/not of interest (0).

Steps To Reproduce:

https://colab.research.google.com/gist/chadrockey/d79f9e0aba4a316dd879d6ae4b8dd594/untitled0.ipynb

Version:

keras 3.3.3 keras_cv 0.9.0

Hi ! Is there a solution to this problem yet? I am trying to implement the DeepLabv3+ model and facing the same error

UserWarning: You are using a softmax over axis -1 of a tensor of shape (None, 32, 32, 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?

I am interested in binary segmented output. Following is the workaround I tried but it is still giving the same warning and the output is all 1 (which is my minority class, probably because of the loss function I used).

def deeplabModel(input_shape=(32, 32, 12), num_classes=1):

    backbone = keras_cv.models.ResNet50V2Backbone(input_shape=input_shape)
    base_model = keras_cv.models.DeepLabV3Plus(
        backbone=backbone,
        num_classes=num_classes, 
        projection_filters=48,
        # activation='sigmoid'  # Use 'sigmoid' for binary classification
    )

    inputs = tf.keras.Input(shape=input_shape)
    x = base_model(inputs)  
    outputs = tf.keras.layers.Activation('sigmoid')(x)  

    model = tf.keras.Model(inputs=inputs, outputs=outputs) 

    return model

keras-team / keras-cv