Visualization of model evaluation - Plot all labels at same time

[HripsimeS]

@Nic-Ma @yiheng-wang-nv Hello. I need to do a visualization of all labels on the same picture after evaluation of model on true_mask and pred_mask. I would like to have the same kind of visualization you had after training in cell 11 of this notebook:

https://github.com/Project-MONAI/tutorials/blob/master/3d_segmentation/spleen_segmentation_3d.ipynb

But in my case I need to do after evaluation of model on original image spacings. I tried to convert the one hot data into the single channel format using argmax function this way plt.imshow(torch.argmax(test_labels[0], dim=1).detach().cpu()[0, :, :, 250])
But got this error "IndexError: too many indices for tensor of dimension 3". Then I removed 0 to make it 3 dimension, it did plotting but only a small part of the mask. Please can you take a look at my code, you might guess quickly what is wrong and how to fix it.

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

[Nic-Ma]

Hi @HripsimeS ,

Thanks for your interest and experiments here.
The key feature here to convert the spacing is the Invertd transform:

Invertd(
    keys="pred",
    transform=val_org_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,
)

Could you please help check whether your logic is correct about inverting?
@yiheng-wang-nv May help analyze more details if you still have issues.

Thanks in advance.

[HripsimeS]

@Nic-Ma thanks for the advice. I modified the Invertd transform to have a correct coding structure and still did not help.
https://github.com/HripsimeS/Computer-Vision/blob/main/3DUnet_GitHub..ipynb

@yiheng-wang-nv can you please take a look to this tricky issue. Look forward to hearing from you.

[yiheng-wang-nv]

Hi @HripsimeS , sorry that I was busy on other tasks, I will take a look at this issue later.

[HripsimeS]

@yiheng-wang-nv no problem :) Let me know if you need more details from my side. Thanks in advance!

[yiheng-wang-nv]

Hi @HripsimeS , I checked your code and the issue is from these lines:

plt.imshow(torch.argmax(test_outputs[0], dim=1).detach().cpu()[:, :, 250])

Let me explain it step by step:

1. test_outputs is from from_engine, which will return a list of tensors, and the length of the list equals to the batch size of the data loader, thus in your case len(test_outputs) = 1.

2. Therefore test_outputs[0] is the prediction you wanted. Its shape should be in the form of (channels, H, W, D). In your case, channels = 5.

3. Since you'd like to draw a single image to show all labels, it's needed to convert the one-hot form (channels, H, W, D) into a single channel form (H, W, D). What you need is to do torch.argmax(test_outputs[0], dim=0). It is the place that has the issue. You set dim = 1 which will convert (channels, H, W, D) into the shape (channels, W, D).

[HripsimeS]

@yiheng-wang-nv thank you so much for all these detailed explanation, it was really helpful to understand the issue. I used plt.imshow(torch.argmax(test_outputs[0], dim=0).detach().cpu()[:, :, 250]) and could receive nice visualization with all labels!
My warm greetings to MONAI team :) 👍 🥇