My semantic segmentation model classifies everything as background

Question

So, I am working on a semantic segmentation task using U-Net. The dataset is very unbalanced, with the background being by far the most common, and the last class being very scarce. First I trained it using Categorical Cross Entropy as the loss function, and in the end it simply classified everything as background (I used IoU as a measurement of success, and the confusion matrix had non-null values only on the first column, which can only mean that).

I also tried precomputing balanced weights using sklearn, however Keras does not accept precomputed weights for data that has more than three dimensions. I then tried using Focal Loss instead of CCE, because I saw it is very robust with imbalanced data, but I had the exact same results (to be fair, I only ran it for 25 epochs and the loss value was still very high and far from converging, so I will let it run again with 50 while I sleep, but I'm exploring other options in the meantime). Now I am running it with a DICE-based loss and hoping to see better results.

In case none of those work, is there anything you would recommend to improve the issue with imbalance?

Answer 1

Are you using some data augmentation with random crops / rotations / zooms ? If you do, you might have some images with only background labels and if so I would suggest you to add a condition to only retain the patches with a ratio of non-background pixels above a certain threshold value.

Answer 2

This answer might be a little late, but I believe that what you need is the Focal Tversky Loss (https://arxiv.org/abs/1810.07842).

Neither the vanilla Focal Loss nor the Dice Loss generalize well for this kind of problems. And it gave me good results, some time ago, in identifying the same small object across many pictures.

According to the authors, one must avoid missing any of the pixels of the imbalanced class (ideally, 0% FN) in detriment of misclassifying background class pixels (FPs) to improve recall:

One of the limitations of the Dice loss function is that it equally weighs false positive (FP) and false negative (FN) detections. In practice, this results in segmentation maps with high precision but low recall. With highly imbalanced data and small ROIs such as skin lesions, FN detections need to be weighted higher than FPs to improve recall rate.

With a quick google search, I could find you a Keras implementation of this Focal Tversky Loss function: https://www.kaggle.com/code/bigironsphere/loss-function-library-keras-pytorch#Focal-Tversky-Loss

Answer 3

Class weights are probably the best first choice when fighting foreground/background imbalance. But another similar approach is just not computing the loss/metrics over the background. If you do that, easily the model learns to segment all as background and still get low dice/iou losses. If your model does segment the foreground well, it also means a good background segmentation.