Good performance with Accuracy but not with Dice loss in Image Segmentation

Viewed 5637

I'm doing image segmentation with U-Net like architecture on Tensorflow w/Keras but I'm new in Deep Learning.
I've got this dataset with the following set shapes:

  • Train : X : (175250, 30, 30, 6) Y: (175250, 30, 30)
  • Validation: X: (29200, 30, 30, 6) Y: (29200, 30, 30)
  • Test : X: (29200, 30, 30, 6) Y: (29200, 30, 30)

Got some examples of these images and each channel, further down.

--> with 20% positive examples and 80% negative examples equally in each set

I run some series but for the best filter combination it made the plot for BCE with good accuracy: enter image description here

The plot for custom functions, Dice_Loss by Dice_Coeff:

enter image description here

And some images generated from the best model trained with test images:

enter image description here

The problem is when I change to dice loss and coefficient, there aren´t good predictions as we seen in the image plot and now it isn´t in the image prediction as we may see. enter image description here

Why it performs so badly in the dice loss? What other function do you recommend?

My dice loss and coefficient function:

def dice_coeff(y_true, y_pred, smooth=1):
    intersection = K.sum(K.abs(y_true * y_pred), axis=-1)
    return (2. * intersection + smooth) / (K.sum(K.square(y_true),-1) + K.sum(K.square(y_pred),-1) + smooth)


def dice_loss(y_true, y_pred):
    return 1-dice_coeff(y_true, y_pred)
1 Answers

Have you tried using soft-dice?
Recent works in computer vision have proposed soft surrogates to alleviate discrepancies and directly optimize the desired metric, either through relaxations (soft-Dice, soft-Jaccard).

There are other loss function that is popular for image segmentation:

The most commonly used loss function for the task of image segmentation is a pixel-wise cross entropy loss. This loss examines each pixel individually, comparing the class predictions (depth-wise pixel vector) to our one-hot encoded target vector.

Another popular loss function for image segmentation tasks is based on the Dice coefficient, (which you have tried already) which is essentially a measure of overlap between two samples. This measure ranges from 0 to 1 where a Dice coefficient of 1 denotes perfect and complete overlap.

In case you were wondering, there's a 2 in the numerator in calculating the Dice coefficient because our denominator "double counts" the common elements between the two sets. In order to formulate a loss function which can be minimized, we'll simply use 1−Dice.

This loss function is known as the soft Dice loss because we directly use the predicted probabilities instead of doing threshold and converting them into a binary mask.

A soft Dice loss is calculated for each class separately and then averaged to yield a final score. An example implementation is provided below.

def soft_dice_loss(y_true, y_pred, epsilon=1e-6): 
    """Soft dice loss calculation for arbitrary batch size, number of classes, and number of spatial dimensions.
    Assumes the `channels_last` format.
  
    # Arguments
        y_true: b x X x Y( x Z...) x c One hot encoding of ground truth
        y_pred: b x X x Y( x Z...) x c Network output, must sum to 1 over c channel (such as after softmax) 
        epsilon: Used for numerical stability to avoid divide by zero errors
    """
   
    # skip the batch and class axis for calculating Dice score
    axes = tuple(range(1, len(y_pred.shape)-1)) 
    numerator = 2. * np.sum(y_pred * y_true, axes)
    denominator = np.sum(np.square(y_pred) + np.square(y_true), axes)
    
    return 1 - np.mean(numerator / (denominator + epsilon)) # average over classes and batch

You may try this functions and find out what works best for your model.

Related