what is the correct way to reshape image segmentation model output from 2D (num of classes, num of pixels) to 3D channel last images

Question

I am using keras and python for satellite image segmentation. It is my understanding that to get (pixel level)predictions for image segmentation, model reshapes layer of dimension(-1,num_classes,height,width) to shape (-1,num_classes,height*width).This is then followed by applying activation function like softmax or sigmoid. My question is how to recover images after this step back in the format either channel first or channel last? example code

o = (Reshape((  num_classes , outputHeight*outputWidth)))(o)
o = (Permute((2, 1)))(o)
o = (Activation('softmax'))(o)

I have tried adding following layer to the model at the end

o = (Reshape((outputHeight, outputWidth, num_classes)))(o)

Is this correct? will this reorient the image pixels in the same order as original or not? Another alternative may be to use following code on individual images.

array.reshape(height, width, num_classes)

Which method should i use to get pixel level segmentation result?

Answer 1

No, if you are interested in an image segmentation, you should not flatten and then reshape your tensors. Instead, use a fully convolutional model, like the U-Net . You find a lot of example implementations of it on github, eg here