# loss function for multi-label segmentation with class inbalance

In order to use a binary segmentation loss function in a multi label problem, I would like to permute the batch axis with the channel axis in the loss computation in order to compute the loss by channel rather than batches to fight a class inbalance problem. (for example, if a class A is very predominant on the others, I could obtain a good loss result by attributing class A to each pixels / voxels.)

Here is an example using the dice loss on 3d images:


def dice_coef(y_true, y_pred, smooth=1e-10):

# permute channels with batches
# eg:  [num_batches, height, width, depth, num_channels] --becomes--> [num_channels, height, width, depth, num_batches]
# in order to perform the computation per channel rather than batches
y_true = K.permute_dimensions(y_true, (4,1,2,3,0))
y_pred = K.permute_dimensions(y_pred, (4,1,2,3,0))

y_true_f = K.batch_flatten(y_true)
y_pred_f = K.batch_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_coef_loss(y_true, y_pred):
return 1 - dice_coef(y_true, y_pred)



is my raisonning correct ?