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?
If you use
learn.get_preds(), it will give you a tuple where the first element is the predicted probabilities
Thanks for the quick reply! I’ll try that out and see if I can get what I’m aiming for.
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:
Thanks for your help! Let me know if you have any thoughts on how to display an overlay like
img.show(y = mask).
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.