Brain Tumor MRI segmentation using Deep Learning.
Introduction
Biomedical Imaging like X-rays, CT scans, MRI assists the doctor in the evaluation of the patient’s current condition. Tumor is a mass or growth of abnormal cells. Early tumor detection in the brain will save many lives. Due to the varying shape and sizes of tumors, it’s hard to locate the exact location. A considerable amount of time is required to detect and locate the tumor in MRI.
Automating this process will add efficiency to the system and free medical staff time which can be utilized to serve more patients. It eliminates the human errors which are more common in developing nations due to the shortage of trained medical staff.
About the dataset:
The dataset source Kaggle .
Dataset consists of 110 patients MRI(Magnetic resonance imaging) together with manual FLAIR (Fluid-attenuated inversion recovery) abnormality segmentation masks obtained from The Cancer Imaging Archive (TCIA).
data.csv file contains patients data and Tumor genomic clusters. MRI scans are in 110 folders named after the CASE_ID. Each folder contains MR images with the following naming convention:
TCGA_<institution-code>_<patient-id>_<slice-number>.tif
Corresponding masks have a ” _mask” suffix. Images are in .tif format having three channels per image.
What is image segmentation?
It is the process of classifying each pixel belonging to a particular class.
Here for this task, we will have two classes :
- tumor part
- non-tumor part
Business Objective:
To create segmentation masks for tumors in the human brain using MRI scans.
Using the segmentation masks we can determine whether the tumor is present or not.
Performance metric:
Focal Tversky Loss (FTL): is a generalisation of the tversky loss. The non-linear nature of the loss gives you control over how the loss behaves at different values of the Tversky index obtained.
Default parameter that works as a good starting point.
α = 0.7, 𝜷 = 0.3, γ = 3/4
Focal Tversky loss is used to tackle the class imbalance.
γ controls the non-linearity of the loss.When γ < 1 it penalising more when TI > 0.5 which changes model to improve more on such examples.
As for many scans, only a small portion of the scan is the tumor.
Exploratory Data Analysis:
We have a total of 3929 MRI scans and their respective segmentation masks. Each MRI scan is of dimensions (256, 256, 3).
We will first classify the MRI scan into tumor and non-tumor using the following function:
Sample of Brain MRI and segmentation mask with no tumor.
Sample of Brain MRI and segmentation mask with tumor.
Distribution of MRI scans before data augmentations
Percentage of patients with no tumor 65.05472130313056
Percentage of patients with tumor 34.9452786968694365% of MRI scans contain no tumor.
Applying Data Augmentation
Three types of image augmentation are applied.
- Flip
- Rotation by a random angle
- Blur operation.
We will include more augmentation images from the minority class (tumor class)
Distribution after data augmentation:
Percentage of patients with no tumor 50.0
Percentage of patients with tumor 50.0I have balanced the data here by adding more images of tumor class after augmentation.
After data augmentation, we have 5312 MRI scans as data points.
Modeling
A) Separate Models
Initially, I tried making creating different models for both classification and segmentation tasks.
Used flow_from_dataframe of ImageDataGenerator class to create a train, test, and validation data generator.
A.1) Simple classifier:
Used to two convolution layers, followed by a flatten layer and at last a dense layer with sigmoid as activation.
The model started to over-fitting as the gap between training and validation accuracy increased.
accuracy: 0.9604
val_accuracy: 0.8540A.2 Classification using Xception Model
As Xception model has more trainable parameter, it performed well than simple classifier but here also the model has over-fitted.
accuracy: 0.9875
val_accuracy: 0.8728A.3) Segmentation:
I created a custom data generator for loading the images in batches.
Built an Unet model with focal Tversky as loss function. Passed (256, 256, 3) images to the model which resulted in (256, 256, 1) segmentation masks.
Note: This model was built on only the MRI which contains tumor.
loss: 0.2907 - val_loss: 0.2974
tversky: 0.8060 — val_tversky: 0.7994
Unet model did not over-fit and the stabilised after 10 epochs.
Below is sample original segmentation and its respective predicted mask.
The resulted mask did a well job here, it misses to capture the sharp edges.
B) Segmentation Model on both classes data
Both the input and output from the model are of dimension (256, 256, 3).
Note: training is slow and it is initially over-fitting to the train data if the learning rate is high.
Tip: Keep learning rate low for better training.
B.1) Unet model with Resblocks:
Used a four-stage encoder-decoder Unet model.
This is the simplest Unet model for this task.
loss: 0.1888 - tversky: 0.8880
val_loss: 0.1840 - val_tversky: 0.8925B.2)Unet with EfficientNet as the backbone
Instead of using the Resnet as the backbone, used EfficientNet trained on Image net dataset.
loss: 0.1480 - tversky: 0.9206
val_loss: 0.1513 - val_tversky: 0.9183
Sample result:
Failed example:
The model failed to find the tumor when the size is very small.
B.3) Double- Unet
It was a work published in 2020. It aims to create a model that generalise well on various segmentation datasets.
Here there two masks are predicted instead of a single mask.
loss: 0.1539 — tversky: 0.9168
val_loss: 0.1551 — val_tversky: 0.9158
Comparing the performance of different models:
Efficient Unet has 10x more trainable parameters than Unet. Double Unet has almost 3x more trainable parameters.
Efficient Unet has turned out to be the best performer.
Deployment:
Saved the best model and using it for deployment.
The model was deployed by a web app using streamlit.
Future Work:
- Different Deep learning architecture from academic research can be explored.
- Networks utilising the 3D spatial nature of MRI like 3d tiramisu net,Deep Medic.
- Collection of diverse MRI scans can be combined.
- Using general segmentation models like SegNet on medical images.
Codebase and Deployment.
Github repository: containing the full code of this project. [Github repo]
Below is the video demonstrating the model at deployment.
