FCN output into binary image


This might be a noobie question, but I hope to learn. I have a FCN operating on 3D images for segmentation purposes. I have only two classes, so the output is of the same size as the input. My question is; how do I turn the output into a binary image so that I can compare it with ground truth for evaluation? I have no softmax at the end of the network. The data is of type [batchsize, channels, x, y , z]

I assume you are able to train your model and are probably using the CrossEntropyLoss, since your model seems to return logits.
Now you would like to compare your model output visually with the target?
Could you give me a hint on your output size for your 3D images? I will try to create a sample for 2D images.

You could try the following:

# Your model output (assuming you have 2 channels)
output = Variable(torch.randn(1, 2, 24, 24))
pred_prob = F.softmax(output, dim=1)[:, 1]
pred = pred_prob > 0.5

plt.imshow(pred.data.numpy()[0, ...])

Let me know, if this helps you or if I misunderstood your question!

Thank you for your reply.

Maybe I was not clear enough. The problem I have is, how do I turn the output from the network into a binary prediction image? I think you indirectly answered that with the “pred = pred_prob > 0.5” line of the code you wrote. Is this the usual way; to use softmax and then prob > 0.5 outside of the network?

Yes. Else it would be hard to differentiate I think…