Calculate dice score for arbitrary number of classes (beyond 2)

gistfile1.txt

#to_categorical() acquired from Keras implementation
#dice_coef() inspired by Marco Jocic's implementation: https://github.com/jocicmarko/ultrasound-nerve-segmentation/blob/master/train.py


def to_categorical(y, nb_classes=None):
    """Converts a class vector (integers) to binary class matrix.
    E.g. for use with categorical_crossentropy.
    # Arguments
        y: class vector to be converted into a matrix
            (integers from 0 to nb_classes).
        nb_classes: total number of classes.
    # Returns
        A binary matrix representation of the input.
    """
    y = np.array(y, dtype='int').ravel()
    if not nb_classes:
        nb_classes = np.max(y) + 1
    n = y.shape[0]
    categorical = np.zeros((n, nb_classes))
    categorical[np.arange(n), y] = 1
    return categorical

def dice_coef(y_pred,y_true):
    smooth = 1.0
    y_true = to_categorical(y_true,nb_classes=2) #Number of classes=2 , to_categorical makes y_true in the following shape (,2)
    y_pred = np.transpose(y_pred,(0,2,3,4,1)) #Transpose prediction array so that classes dimension is at the end
    y_pred = np.reshape(y_pred,(y_pred.shape[0]*y_pred.shape[1]*y_pred.shape[2]*y_pred.shape[3],y_pred.shape[4])) # Reshape pred_y (,2)
    intersection = np.sum(y_true * y_pred_)
    return (2. * intersection + smooth) / (np.sum(y_true) + np.sum(y_pred) + smooth)