Proper mask shape for a multiclass task

Hello,

I’m trying to use custom_unet for a multiclass task. In my case all masks inputs have a following shape (BATCH_SIZE, NUMBER_OF_CLASSES, IMAGE_WIDTH, IMAGE_HEIGHT, 1). Since I have 4 different colours on the black background and my images’re 256x256 pixels and I use batch size of 4, I end up with the following shape: (4, 4, 256, 256, 1). Unfortunately, custom_unet doesn’t like this shape giving me this error:

<...>
  File "C:\Users\E-soft\Anaconda3\envs\Explorium\lib\site-packages\keras_unet\losses.py", line 44, in jaccard_distance
    intersection = K.sum(K.abs(y_true * y_pred), axis=-1)
  File "C:\Users\E-soft\Anaconda3\envs\Explorium\lib\site-packages\tensorflow_core\python\ops\math_ops.py", line 899, in binary_op_wrapper
    return func(x, y, name=name)
  File "C:\Users\E-soft\Anaconda3\envs\Explorium\lib\site-packages\tensorflow_core\python\ops\math_ops.py", line 1206, in _mul_dispatch
    return gen_math_ops.mul(x, y, name=name)
  File "C:\Users\E-soft\Anaconda3\envs\Explorium\lib\site-packages\tensorflow_core\python\ops\gen_math_ops.py", line 6698, in mul
    _six.raise_from(_core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.InvalidArgumentError: Incompatible shapes: [4,4,256,256] vs. [4,256,256,4] [Op:Mul] name: loss/conv2d_18_loss/mul/

It seems that I should reshape my input for masks, Should it be (BATCH_SIZE IMAGE_WIDTH, IMAGE_HEIGHT, NUMBER_OF_CLASSES, 1)?

[karolzak]

Hi @EtagiBI !

For multiclass semantic segmentation tasks you need to reshape your masks into (BATCH_SIZE IMAGE_WIDTH, IMAGE_HEIGHT, NUMBER_OF_CLASSES).
And to explain why behind it: with this shape you should treat NUMBER_OF_CLASSES as a number of layers in your output masks where each layer represents each class you are trying to predict. In other words your output mask will be a volume of binary masks where each mask represents predictions for specific class.

To give you some visuals on a dummy example:

print(x.shape, y.shape, y_pred.shape)
(1, 256, 256, 3) (1, 256, 256, 3) (1, 256, 256, 3)

utils.plot_imgs(
    org_imgs=np.vstack(
        (x, x, x)),
    mask_imgs=np.stack(
        (y[0][:,:,0],  # first layer of mask is the first class gt
         y[0][:,:,1],   # second layer of mask is the second class gt
         y[0][:,:,2])),   # third layer of mask is the third class gt
    pred_imgs=np.stack(
        (y_pred[0][:,:,0],  # first layer of mask is the first class predictions
         y_pred[0][:,:,1],   # second layer of mask is the second class predictions
         y_pred[0][:,:,2]))   # third layer of mask is the third class predictions
)

[karolzak]

@EtagiBI the code you shared looks fine although I prefer array slicing to extract specific layers. Also GRAY2RGB conversion is unnecessary unless you really need it. Grayscale image is perfectly fine especially for segmentation mask.

If your prediction results look messy your model probably failed to learn anything.
In my dummy example (training on a single image…) I used code below:

from keras_unet.models import custom_unet

model = custom_unet(
    (256,256,3), num_classes=3,
    use_batch_norm=False,
    dropout=0.0, dropout_type='standard',
    output_activation='sigmoid')

from tensorflow.keras.optimizers import Adam
from keras_unet.metrics import iou

model.compile(
    optimizer=Adam(),
    loss='binary_crossentropy',
    metrics=[iou]
)

model.fit(x, y, epochs=200)

When I tried building the network with softmax it couldn’t learn anything.