The number of labels during training and testing are different

[HripsimeS]

Hi @yiheng-wang-nv I would like to mention one important mark concerning the visualization of labels. I have totally 5 labels (including a background) and after training process I can see all 4 labels in true_masks and predicted_masks. I used following Compose during training:
post_pred = Compose([EnsureType(), AsDiscrete(argmax=True, to_onehot=5)])
post_label = Compose([EnsureType(), AsDiscrete(to_onehot=5)])

But during evaluation of model on original images/masks visualization shows only 1 label for both true_masks and predicted_masks.
I assume the issue comes from post_transform, although I mentioned to_onehot=5

post_transforms = Compose([
EnsureTyped(keys="pred"),
Invertd(keys="pred", transform=test_transforms, orig_keys="image", meta_keys="pred_meta_dict",
orig_meta_keys="image_meta_dict", meta_key_postfix="meta_dict", nearest_interp=False,to_tensor=True,),
AsDiscreted(keys="pred", argmax=True, to_onehot=5),
AsDiscreted(keys="mask", to_onehot=5),
])

Do you know by chance what could cause these changes with number of labels. Thanks in advance for your help!

[yiheng-wang-nv]

Hi @HripsimeS ,

But during evaluation of model on original images/masks visualization shows only 1 label for both true_masks and predicted_masks.

Here what do you mean visualization shows only 1 label for both true_masks and predicted_masks?
can you check the shape of pred and mask? I think their shape[0] should be 5.

[HripsimeS]

Hi @yiheng-wang-nv Thanks for your quick reply! I checked the shapes of test_labels[0] and test_outputs[0]. For both of them the shapes are [5, 512, 512, 350]. But when I try to plot the results, I don't see all labels even when I check on different sizes. Here is my notebook:

https://github.com/HripsimeS/Computer-Vision/blob/main/MONAI_3DUnet_Github.ipynb

Let me know if you need additional information.

[yiheng-wang-nv]

I checked your code and found that only channel 1 is plotted:

plt.imshow(test_labels[0].detach().cpu()[1, :, :, 250])
plt.subplot(1, 3, 3)
plt.title(f"predicted mask")
plt.imshow(test_outputs[0].detach().cpu()[1, :, :, 250])

If you need to visualize all labels, you need to do sth like:

for cls in range(1, 5):
    plt.imshow(test_labels[0].detach().cpu()[cls, :, :, 250])
    ...

[HripsimeS]

@yiheng-wang-nv thank you, it worked but for each label it plots separately (with the same yellow color).
Would be nice to have all labels on the same plot with different colors :)

[yiheng-wang-nv]

HI @HripsimeS , a simple way is to convert the one hot data into the single channel format first and then draw it, like:

data = data_one_hot.argmax(0)
plt.imshow(data[:, :, 250])
...

you may need to learn more about the matplotlib library, or other powerful tools to see how to visualize it better. I am also not familiar with the visualization thing.