UNET Segmentation Model Learns only 1-Class [Colab Shared]

vision
#1

Hello all,

I am having a strange issue where my model is only learning a specific class in a data-set (Liver class ~ blue). I experimented more and realized if I remove the Liver class, the model can identify the kidneys very well (yellow and red). Interestingly the kidneys are much brighter in these images!

Does this essentially mean my model is short on parameters for predicting all organs? I would think adding classes would enhance the original model.

image
Image: Source, Ground, Prediction

Link to Original Colab Notebook. All data downloads into the notebook automatically! (You can just click ‘Open in playground’ to run the code).

https://colab.research.google.com/drive/1cDXzCnTaVDEe3Dnhafp-Rdu6rKs3BuXI

Link to Colab Notebook. Only predicts kidneys.
https://colab.research.google.com/drive/1_49DjRyQhnlgB5z_lIyXkV5pYydEYvd3

#2

Hi Akshay,

I’d probably look at two diagnostic (ha) questions:

  • If you permute the labels (labels = labels.flip(1) or so), does it still happen this way or can you then predict another organ?
  • If you stick with a single batch, can you overfit?

The former is a check whether there is something wrong in the tensor shapes or somesuch. I tend to print all shapes to debug that type of thing as a first step.

The second would go in the direction of capacity.

Best regards

Thomas

#3

Thanks for the feedback @tom!

I reduced the resolution of the input and that fixed the issue suggesting that it was a capacity problem. I had checked permuting the labels and that still had the same behavior.