将2D（类数，像素数）输出的图像分割模型重塑为3D通道最后一张图像的正确方法是什么

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