BrainTumorSegmentation

Introduction

Code for implementation of the model published here.

Abstract

Brain tumors are one of the deadliest forms of cancer with a mortality rate of over 80%. A quick and accurate diagnosis is crucial for increasing the chances of survival. However, in medical analysis, the manual annotation and segmentation of brain tumors are complicated. Multiple MRI modalities are typically analyzed as they provide unique information about the tumor regions. Although using several MRI modalities helps segment brain tumors, they tend to increase overfitting and computation in computer models. This paper proposed a region of interest detection algorithm that was implemented during data preprocessing to locate salient features and remove extraneous MRI data. This decreased the input size of the model, allowing for more aggressive data augmentations and deeper neural networks. Following the preprocessing of the MRI modalities, a fully convolutional autoencoder segmented the different brain MRIs using channel-wise attention and attention gates. Subsequently, test-time augmentations and an energy-based model were used for voxel-based uncertainty predictions. Experimentation was conducted on the BraTS 2019, 2020, and 2021 benchmarks, and the proposed models achieved a state-of-the-art segmentation performance with mean dice scores of 84.55, 88.52, and 90.82 on each respective dataset. Additionally, qualitative results were used to assess the segmentation models and uncertainty predictions. Extensive experimentation showed that the proposed framework with cropping and attention mechanisms achieved state-of-the-art performance when compared to its contemporaries.

Setup

Getting Started

1. Clone the repository

   git clone https://github.com/WeToTheMoon/BrainTumorSegmentation.git
   

2. Download the dataset(s) being tested with
   • BraTS2019
   • BraTS2020
   • BraTS2021

Dependencies

It is recommended to create a conda or venv when working with this project. A full list of requirements can be found in the requirements.txt file. Python 3.10.12 was used.

Creating Dataset

Datasets consist of two main types, binary and cropped datasets. The binary dataset is created using the binary model and is used in order to create the cropped dataset which is used by the multiclass model. All of these datasets are wrapped using the MRIDataset class in order to easily validate and access its data.

In order to create the dataset, a helper method has been provided. Simply call the method with the location of the previously downloaded dataset and the desired output path.

from utils.dataset_helpers import create_new_dataset

input_dataset_path = ...
output_dataset_path = ...

# Already handles training the binary model
create_new_dataset(input_dataset_path, output_dataset_path)

Results

After creating the dataset, the multiclass model can be run

python train_multiclass --dataset_dir <cropped_dataset_path> --weights <location_to_save_model_weights>

or

from train_multiclass import train as train_multiclass_model

cropped_dataset_path = ...
multiclass_model_weights = ...
train_multiclass_model(cropped_dataset_path, multiclass_model_weights)