My semantic segmentation model classifies everything as background

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?

Topic semantic-segmentation imbalanced-data keras machine-learning

Category Data Science


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.

About

Geeks Mental is a community that publishes articles and tutorials about Web, Android, Data Science, new techniques and Linux security.