Why do we use a softmax activation function in Convolutional Autoencoders?

I have been working on an image segmentation project where I have created a convolutional autoencoder. I saw this image and implemented it using Keras.

At the output layer, the author has used the softmax activation function. Shouldn't it be ReLU?

According to me, this seems to be a regression problem, where we need to predict the continuous values for the segmented image pixels. If so, why are we using a softmax function, instead of ReLU or a linear function?

2 Answers

Looking at the images, it seems an image segmentation problem. As you see, the Autoencoder is not returning a more or less accurate representation of the original input image, but the segmentation of it. Each of the output colors are segments of the image that have been classified. Probably the classes are something such as: "road", "tree", "building", "sky", etc.

In other words, it's an Autoencoder performing a classification task. That explains the softmax layer at the end.

Since its a multi-class classification problem, each class will have its own probability value between 0 and 1 and sum of all probabilities of classes will equal to 1.