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 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:
- Non-enhancing tumor
- 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:
- Generalised dice loss: which uses classes weights inversely proportional to the label area
- 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:
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.
Xpress AI Newsletter
Join the newsletter to receive the latest updates in your inbox.