BraTS(Brain Tumour Segmentation) using V-Net
This project is an approach to detect brain tumours using BraTS 2016,2017 dataset.
Description
BraTS is a dataset which provides multimodal 3D brain MRIs annotated by experts. Each Magnetic Resonance Imaging(MRI) scan consists of 4 different modalities(Flair,T1w,t1gd,T2w).
Expert annotations are provided in the form of segmentation masks to detect 3 classes of tumour – edema(ED),enhancing tumour(ET),necrotic and non-enhancing tumour(NET/NCR). The dataset is challenging in terms of the complex and heterogeneously-located targets.
We use Volumetric Network(V-Net) which is a 3D Fully Convolutional Network(FCN) for segmentation of 3D medical images. We use Dice Loss as the objective function for the present scenario. Future implementation will include Hausdorff Loss for better boundary segmentations.
Fig 1: Brain Tumour Segmentation
Getting Started
Dataset
4D Multimodal MRI dataset
The dataset contains 750 4D volumes of MRI scans(484 for training and 266 for testing). Since the test set is not publicly available we split the train set into train-val-split. We use 400 scans for training and validation and the rest 84 for evaluation. No data augmentations are applied to the data. The data is stored in NIfTI file format(.nii.gz). A 4D tensor of shape (4,150,240,240) is obtained after reading the data where the 1st dimension denotes the modality(Flair,T1w,t1gd,T2w), 2nd dimension denotes the number of slices and the 3rd and 4th dimesion denotes the width and height respectively. We crop each modality to (32,128,128) for computational purpose and stack each modality along the 0th axis. The segmentation masks contain 3 classes – ED,ET,NET/NCR. We resize and stack each class to form a tensor of shape (3,32,128,128).
Experimental Details
Loss functions
We use Dice loss as the objective function to train the model.
Training
We use Adam optimizer for optimizing the objective function. The learning rate is initially set to 0.001 and halved after every 100 epochs. We train the network until 300 epochs and the best weights are saved accordingly. We use NVIDIA Tesla P100 with 16 GB of VRAM to train the model.
Quantative Results
We evaluate the model on the basis of Dice Score Coefficient(DSC) and Intersection over Union(IoU) over three classes (WT+TC+ET).
Qualitative Results
Fig 1: Brain Complete Tumour Segmentation(blue indicates ground truth segmentation and red indicates predicted segmentation)
Statistical Inference
Fig 1: Validation Dice Score Coefficient(DSC)
Fig 2: Validation Dice Loss
Dependencies
- SimpleITK 2.0.2
- Pytorch 1.8.0
- CUDA 10.2
- TensorBoard 2.5.0
Installing
pip install SimpleITK
pip install tensorboard
Execution
python train.py
train.py contains code for training the model and saving the weights.
loader.py contains code for dataloading and train-test split.
utils.py contains utility functions.
evaluate.py contains code for evaluation.
Acknowledgments
[1] BraTS 3D UNet
[2] VNet
