How calculate the dice coefficient for multi-class segmentation task using Python?

Question

I am wondering how can I calculate the dice coefficient for multi-class segmentation.

Here is the script that would calculate the dice coefficient for the binary segmentation task. How can I loop over each class and calculate the dice for each class?

Thank you in advance

import numpy 



def dice_coeff(im1, im2, empty_score=1.0):

    im1 = numpy.asarray(im1).astype(numpy.bool)
    im2 = numpy.asarray(im2).astype(numpy.bool)

    if im1.shape != im2.shape:
        raise ValueError("Shape mismatch: im1 and im2 must have the same shape.")

    im_sum = im1.sum() + im2.sum()
    if im_sum == 0:
        return empty_score

    # Compute Dice coefficient
    intersection = numpy.logical_and(im1, im2)

    return (2. * intersection.sum() / im_sum)

Answer 1

You can use dice_score for binary classes and then use binary maps for all the classes repeatedly to get a multiclass dice score.

I'm assuming your images/segmentation maps are in the format (batch/index of image, height, width, class_map).

import numpy as np
import matplotlib.pyplot as plt

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

def dice_coef_multilabel(y_true, y_pred, numLabels):
    dice=0
    for index in range(numLabels):
        dice += dice_coef(y_true[:,:,:,index], y_pred[:,:,:,index])
    return dice/numLabels # taking average

num_class = 5

imgA = np.random.randint(low=0, high= 2, size=(5, 64, 64, num_class) ) # 5 images in batch, 64 by 64, num_classes map
imgB = np.random.randint(low=0, high= 2, size=(5, 64, 64, num_class) )


plt.imshow(imgA[0,:,:,0]) # for 0th image, class 0 map
plt.show()

plt.imshow(imgB[0,:,:,0]) # for 0th image, class 0 map
plt.show()

dice_score = dice_coef_multilabel(imgA, imgB, num_class)
print(f'For A and B {dice_score}')

dice_score = dice_coef_multilabel(imgA, imgA, num_class)
print(f'For A and A {dice_score}')