This repository implements three different models and analyses the results on the BRATS 2020 dataset downloaded from the link. This work is done as part of the course project MBP 1413H: Biomedical Applications of AI.
This project aims at two goals
-
Brain tumour segmentation
- ResUNet
- Attention U-Net
- UNet Transformer(UNETR)
-
Uncertainty prediction
- Monte Carlo Dropout
BRATS dataset contains multimodal MRI scans with extensive annotation, comprising GD-enhancing Tumor, peritumoral edema, and necrotic and non-enhancing tumor core. For our experiments, we group labels to make new set of labels as follows:-
- Tumor core label was created by joining non enhancing tumor core and GD enhancing tumor
- Whole tumor label was created by joining non enhancing tumor core, peritumoral edema, and GD enhancing tumor
- Enhancing tumor label was created by just considering class GD enhancing tumor
- Background label was created by considering all the pixels not included in the Whole tumor label
Using these labels, we performed multi label segmentation using three models. We report the dice score for our three experiments.
Dice | TC | ET | WT | AVG |
---|---|---|---|---|
Resunet (WithNoAug) | 0.4 | 0.79 | 0.68 | 0.62 |
Resunet (WithAug) | 0.70 | 0.74 | 0.68 | 0.69 |
Attention(WithAug) | 0.72 | 0.73 | 0.71 | 0.72 |
Attention (WithNoAug) | 0.65 | 0.69 | 0.66 | 0.67 |
UnetR(WithAug) | 0.76 | 0.84 | 0.72 | 0.77 |
UnetR(WithNoAug) | 0.50 | 0.65 | 0.69 | 0.61 |
Please refer to the report for more details.
These instructions will get you a copy of the project up and running on your local machine for development and testing purposes. Note that your local environment should be cuda enabled.
Libraries needed
monai
torchmetrics
pytorch
torchvision
segmentation_models_pytorch
pathlib
The dataset can be downloaded from the link. For running Attention U-Net experiments, one can directly use the downloaded dataset stored in h5
format. For running UResnet and UNETR, the dataset should be in nifty
format.
1 - remove_input.py
- This script runs ablation study by removing channels and replacing them with random noise or 02 - train_unetr.py
- This script trains transformer unet3 - transform.py
- This script convertshdf5
tonifti
format. Run this script before training the network4 - uncertainity.py
- This script gets the uncertaininty maps for a given model using Monte Carlo Dropout5 - utils\
- This directory contains helper functions for dataloading, metric calculations, visualizations, uncertaininty calculation, etc.
1 - Hyperparameters\
- This directory is used bytrain_attunet_multilabel.py
. It containsjson
files with hyperparameter settings for experiment2 - evaluate.py
- This script gets the metric scores and uncertainty map for a given model and dataset3 - remove_input_attn.py
- This script runs ablation study by removing channels and replacing them with random noise or 04 - utils\
- This directory contains helper functions for dataloading, metric calculations, visualizations, uncertaininty calculation, etc.
1 - train_Resunet.py
- This script trains ResUnet2 - uncertainity.py
- This script gets the uncertaininty maps for a given model using Monte Carlo Dropout3 - utils\
- This directory contains helper functions for dataloading, metric calculations, visualizations, uncertaininty calculation, etc.
# To train UNETR
# With Slurm Scripts
>>> sbatch -p gpu /path/to/script
# Directly run python
>>> python ./UNet_transformer/train_unetr.py
# Directly run python
>>> python ./Attention_UNET/train_attunet_multilabel.py\
-p [str: HYPERPARAMETER LOCATION]\
-n [str: NAME OF EXPERIMENT]\
-l [str: DATASET LOCATION MODIFIER (compute canada)]\
-m [bool: MULTI GPU (True/False)]\
-r [str: RESUME FROM CHECKPOINT(default: None)]\
-s [bool: SAVE CHECKPOINTS (True/False)]\
-w [bool: USE WANDB (True/False)]\
# Directly run python
>>> python ./UResnet/train_Resunet.py
The following people contributed to this project.