MRI Segmentation and Tumor Detection Pipeline

Hi guys, currently I’m working on a pipeline for head tissues and brain tumor segmentation.

If you are excited about some brain stuff, keep reading…..else keep reading.... :D

The pipeline consist of two major parts:
1- The tissues segmentation model
2- the tumor segmentation model


In this pipeline, U-net architecture is used for semantic segmentation of the tissues and the tumor

The datasets used are:
1- first stage: self-made dataset using SPM12 for tissue masks
2- second stage: Brain Tumor Segmentation (BraTS) Challenge 2017

The first stage is used to extract binary masks for the brain regions, the ROI for the tumor model, the current result achieved for brain extraction model is at mIOU = 95.6%:


For the second stage The BraTS 2017 dataset contains 3 stages for brain tumor:

  1. Edema
  2. Non-enhancing tumor
  3. Enhancing tumor


In the beginning, due to the small size of the tumor regions in comparison to the background, the model generalized all predictions as background.

For a fix, I have changed the loss function form Dice loss + categorical cross-entropy to:

  1. Generalised dice loss: which uses classes weights inversely proportional to the label area
  2. and, Weighted cross-entropy: which allows to manually set the classes weights
def weighted_log_loss(y_true, y_pred):
    y_pred /= K.sum(y_pred, axis=-1, keepdims=True)
    y_pred = K.clip(y_pred, K.epsilon(), 1 - K.epsilon())
    loss = y_true * (K.log(y_pred) * weights)
    loss = K.mean(-K.sum(loss, -1))
    return loss

def gen_dice_loss(y_true, y_pred):
    y_true_f = K.reshape(y_true,shape=(-1,4))
    y_pred_f = K.reshape(y_pred,shape=(-1,4))
    sum_pr=K.sum(y_true_f * y_pred_f,axis=-2)
    generalised_dice_numerator =2*K.sum(weights*sum_pr)
    generalised_dice_denominator =K.sum(weights*(sum_r+sum_p))
    generalised_dice_score =generalised_dice_numerator /(generalised_dice_denominator)
    del sum_p,sum_r,sum_pr,weights

    return GDL+weighted_log_loss(y_true,y_pred)

Here some of the obtained results:


These results can be considered promising, yet further fine-tuning is required to segment more tissues classes in the brain tissues model, and to increase the accuracy for the different tumor classes in the tumor model.

Hope this sharing excites you, give me some comments below if you have any question, suggestion or idea on how to further improve this work.