U-net Output: Probability instead of Class?

Part 1 (2019) Advanced (Part 1 v3)
1

Hi everyone!

I’m working on a pretty basic U-net segmentation task with only two classes.

Instead of the output being binary 0’s and 1’s, is it possible to have the final output mask be a grayscale representation of the probability that each pixel belongs to class 1?

Thanks,
Jona

2

If you use learn.get_preds(), it will give you a tuple where the first element is the predicted probabilities

3

Thanks for the quick reply! I’ll try that out and see if I can get what I’m aiming for.

4

A bit of a hassle…
I wasn’t able to display the probability output of preds (floats between 0 and 1) as a segment mask, because it is a Tensor not an Image. If I try to force it into a FastAI Image type, it complains about not having a cpu

I hacked my way through, and just display them side by side instead of overlay (disappointing) using this code:

matplotlib.pyplot.imshow(learn.predict(img)[2][1], cmap="binary")

Thanks for your help! Let me know if you have any thoughts on how to display an overlay like img.show(y = mask).

5

You can try with ImageSegment(tensor) as long as the tensor has the right shape. That by itself displays a heatmap of your probabilities. If you want to use it as a mask overlay, you also need to resolve how to decide the class, such as ImageSegment(tensor > 0.5). If you have logits, might want to use torch.sigmoid(tensor) > 0.5 instead.