Restoring mask to image
My PyTorch model outputs a segmented image with values (0,1,2) for each one of the three classes. During the preparation of the set, I mapped black to 0, red to 1 and white to 2. I have two questions:

How can I show what each class represents? for example take a look at the image:

I am currently using the following method to show each class:
     output = net(input)

     input = input.cpu().squeeze()
     input = transforms.ToPILImage()(input)

     probs = F.softmax(output, dim=1)
     probs = probs.squeeze(0)

     full_mask = probs.squeeze().cpu().numpy()

     fig, (ax0, ax1, ax2, ax3, ax4) = plt.subplots(1, 5, figsize=(20,10), sharey=True)

     ax0.set_title('Input Image')
     ax1.set_title('Background Class')
     ax2.set_title('Neuron Class')
     ax3.set_title('Dendrite Class')
     ax4.set_title('Predicted Mask')

     ax0.imshow(input)
     ax1.imshow(full_mask[0, :, :].squeeze())
     ax2.imshow(full_mask[1, :, :].squeeze())
     ax3.imshow(full_mask[2, :, :].squeeze())

     full_mask = np.argmax(full_mask, 0)
     img = mask_to_image(full_mask)



But there appears to be shared pixels between the classes, is there a better way to show this (I want the first image to only of the background class, the the second only of the neuron class and the third only of the dendrite class)?
2.My second question is about generating a black, red and white image from the mask, currently the mask is of shape (512,512) and has the following values:
[[0 0 0 ... 0 0 0]
 [0 0 0 ... 2 0 0]
 [0 0 0 ... 2 2 0]
 ...
 [2 1 2 ... 2 2 2]
 [2 1 2 ... 2 2 2]
 [0 2 0 ... 2 2 2]]

And the results look like this:

Since I am using this code to convert to image:
def mask_to_image(mask):
   return Image.fromarray((mask).astype(np.uint8))



But there appears to be shared pixels between the classes, is there a
better way to show this (I want the first image to only of the
background class, the the second only of the neuron class and the
third only of the dendrite class)?

Yes, you can take argmax along 0th dimension so the one with highest logit (unnormalized probability) will be 1, rest will be zero:
output = net(input)

binary_mask = torch.argmax(output, dim=0).cpu().numpy()
ax.set_title('Neuron Class')
ax.imshow(binary_mask == 0)


My second question is about generating a black, red and white image
from the mask, currently the mask is of shape (512,512) and has the
following values

You can spread [0, 1, 2] values into the zero-th axis making it channel-wise. Now [0, 0, 0] values across all channels for single pixel will be black, [255, 255, 255] would be white and [255, 0, 0] would be red (as PIL is RGB format):
import torch

tensor = torch.randint(high=3, size=(512, 512))

red = tensor == 0
white = tensor == 2

zero_channel = red &amp; white

image = torch.stack([zero_channel, white, white]).int().numpy() * 255
Image.fromarray((image).astype(np.uint8))

