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

image

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  http://medicaldecathlon.com/

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%:

image

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

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


source

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())
    weights=np.array([1,5,3,8])
    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_p=K.sum(y_pred_f,axis=-2)
    sum_r=K.sum(y_true_f,axis=-2)
    sum_pr=K.sum(y_true_f * y_pred_f,axis=-2)
    weights=K.pow(K.square(sum_r)+K.epsilon(),-1)
    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)
    GDL=1-generalised_dice_score
    del sum_p,sum_r,sum_pr,weights

    return GDL+weighted_log_loss(y_true,y_pred)

Here some of the obtained results:

image
image

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.