Everything is classified as background by segmentation model

Question

I am training a U-NET model for medical image segmentation. Problem is that the binary masks that im using to train the model mostly consist of background pixels and a very small region of the whole image is the foreground, for example :\

1. Image
   
2. Corresponding Mask:
   

Consequently, my model is just predicting everything as background.

I have tried using three different loss functions - Dice Loss,Focal Loss and Focal Tvwersky Loss. The Dice Loss and Focal Tvwersky Loss gave me high accuracies and high losses - close to 1, whereas Focal loss gave me high accuracies and low losses(around 0.01-0.08). Here are the loss functions I have tried using in the tensorflow model.fit method.

#Dice loss

def dice_coef(y_true, y_pred, smooth=1):
    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)
    intersection = K.sum(y_true_f * y_pred_f)
    return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)

def dice_loss(y_true, y_pred):
    return 1 - dice_coef(y_true, y_pred)

#Focal tversky
ALPHA = 0.5
BETA = 0.5
GAMMA = 1

def FocalTverskyLoss(targets, inputs, alpha=ALPHA, beta=BETA, gamma=GAMMA, smooth=1e-6):

        #flatten label and prediction tensors
        inputs = K.flatten(inputs)
        targets = K.flatten(targets)

        #True Positives, False Positives & False Negatives
        TP = K.sum((inputs * targets))
        FP = K.sum(((1-targets) * inputs))
        FN = K.sum((targets * (1-inputs)))

        Tversky = (TP + smooth) / (TP + alpha*FP + beta*FN + smooth)
        FocalTversky = K.pow((1 - Tversky), gamma)

        return FocalTversky


ALPHA = 0.8
GAMMA = 2

def FocalLoss(targets, inputs, alpha=ALPHA, gamma=GAMMA):

    inputs = K.flatten(inputs)
    targets = K.flatten(targets)

    BCE = K.binary_crossentropy(targets, inputs)
    BCE_EXP = K.exp(-BCE)
    focal_loss = K.mean(alpha * K.pow((1-BCE_EXP), gamma) * BCE)

    return focal_loss

How should I approach the data or model training so it doesnt learn to always predict the background and thereby giving better predictions?

Answer 1

Your model is doing something, but not a lot. Because so many pixels in the training set are background pixels, it learns to predict mostly background pixels. There are many reasons why a model doesn't work, ranging from codebugs to problems with the data to more fundamental problems. I would not look at the loss function first, most likely you just don't have enough data or at least didn't train with enough data.

It can be a good idea to crop your images to only (or mostly) the part of the image of interest. This will speed up training as well as make it more effective.