Imshow differs drastically from applying matplolib.cm to a segmented image

Question

Hi and thanks for reading.

What I am trying to do is to make a web app that would take an image, run it through the model and return a segmented version. I can not use imshow in the webapp though. So I tried adding colormap through matplolib.cm.viridis however it returns a much darker image.

Here are some code and images for refernce:

pred = new_model.predict(np.expand_dims(img, 0))
pred_mask = np.argmax(pred, axis=-1)
pred_mask = pred_mask[0]

This returns me a 2D grayscale image, which when put into matplolib imshow looks like this.(last picture on the right is the output of the model). Code and image below.

axs[0].imshow(m1)
axs[0].set_title('Image')
axs[1].imshow(test_label1)
axs[1].set_title('Ground Truth')
axs[2].imshow(new_pred)
axs[2].set_title('Prediction')

However, when applying colormap to an image using matplolib.cm (something I have to do for app to function) I get this image. Code and image presented below.

Adding colormap. (Viridis, as far as I know is default one from matplolib 3.5)

from matplotlib import cm
pred_mask = cm.viridis(pred_mask / 255)*255
pred_mask = np.asarray(pred_mask, dtype='uint8')

Plotting Image

fig, axs = plt.subplots(1, 3, figsize=(20, 10))
axs[0].imshow(m1)
axs[0].set_title('Image')
axs[1].imshow(test_label1)
axs[1].set_title('Ground Truth')
axs[2].imshow(pred_mask)
axs[2].set_title('Prediction')

But as you can see image is much darker, without even a hint of lighter blue or yellow, ie worse. How can I make it closer to imshow output?

PS. Thank you very much for reading and hope that someone has an answer to that. Any suggestions would be much appreciated though.

Answer 1

This is most likely related to the number range of the image or colormap, respectively. As the prediction mask can be faintly seen my money would be on either multiplying the prediction data with 255 or to set the vmax of imshow to a smaller value. In any case, it would be useful to know the min/max value of pred_mask and additionally show a colorbar for the right plot.

I hope that gets you on the right track.