Skip to content

KurtLabUW/brats2023_updated

Repository files navigation

An Optimization Framework for BraTS 2023

This repository contains the implementation of our paper (https://arxiv.org/abs/2402.07008). The code is developed by Kurtlab for the 2023 International Brain Tumor Segmentations (BraTS) Cluster of Challenges (see here.)

Screen Shot 2024-03-07 at 11 15 38 AM

The aim of the BraTS Challenges is to push the state-of-the-art in brain tumor segmentation algorithms. That is, algorithms that can segment tumor subregions from multiparametric MRI scans of the brain.

We submitted to three challenges:

  1. Adult gliomas.
  2. BraTS-Africa - gliomas from patients in sub-Saharan Africa.
  3. Meningiomas.

Our team ranked in the top 10 of all those who submitted, including 7th place in the BraTS-Africa challenge.

Implementation

Our baseline model is based on the Optimized U-Net presented by Futrega et al. from NVIDIA.

We developed a framework of optimization techniques to boost model performance. These include:

  1. Preprocessing the MRI data with Z-score normalization and contrast rescaling.
  2. Postprocessing our predictions to remove small connected components to minimize the number of false positives predicted, to account for the new performance metrics introduced in BraTS 2023.
  3. A transfer learning technique in which specific layers of the model are frozen before training is continued on a new dataset.

We found that this transfer learning technique successfully allowed us to transfer knowledge from the large Challenge 1 dataset of adult gliomas to the smaller and lower resolution Challenge 2 dataset for BraTS-Africa.

We performed several ablation studies to fine-tune these techniques.

Code

The dependencies for our code are provided in the requirements.txt file.

We provide a brief explanation of each subpackage:

  • datasets - the dataset class for loading training and test data for BraTS;
  • losses - loss functions used in model training;
  • model_routines - commonly run routines for model training, validation and inference;
  • processing - preprocessing, postprocessing and plotting code;
  • utils - many utility functions used in model training and inference.

Usage

To ensure the code runs with the relative imports working correctly, navigate outside the root folder and use python -m to run the Python module as a script. For example, if you wanted to run the training with validation model routine, you would run the following line.

python -m brats2023_updated.model_routines.train_with_val

Contributors

This research was conducted by Tianyi Ren, Ethan Honey and Harshitha Rebala under the supervision of Dr Mehmet Kurt. The code was developed by Ethan Honey and Harshitha Rebala. The updated version of the code and this repository was put together by Ethan Honey. The code for the edge loss function was developed by Agamdeep Chopra.

Acknowledgements

We would particularly like to acknowledge the work of Futrega et al. and NVIDIA for their Optimized U-Net model that was a crucial part of our foundation, as well as their accompanying code.

We acknowledge MONAI for their implementations of the Dice Score and 95% Hausdorff Distance metrics. We used these in our validation code.

We acknowledge the cc3d Python library. We used their connected component tools in our postprocessing strategies.

We acknowledge Abderezaei et al. for their development of the edge loss function. We utilized this in training our model.

Finally, we would like to acknowledge the organizers of BraTS 2023, for presenting the challenges, sharing the datasets and supporting us in the development of our algorithms.

References

To-do

  • Evaluate updated model performance on BraTS-GLI and BraTS-Africa
  • Produce plots of predicted segmentations for the updated model