How can I display a test image and display the mask for it based on my trained model?

[ptrblck]

The code plots the probability maps for both classes, not the mask.
If you are only interested on a single class, you can of course just plot the corresponding prediction.

You could use the IoU implementation or a dice score.