This repository contains official source code for the method proposed in: E₁D₃ U-Net for Brain Tumor Segmentation: Submission to the RSNA-ASNR-MICCAI BraTS 2021 Challenge.

Download the BraTS dataset. The structure of the dataset should be as follows (using nomenclature from BraTS 2021 dataset for the sake of demonstration, however this works for other BraTS dataset as well):

dataset/
    - BraTS2021_AAAAA/
        BraTS2021_AAAAA_flair.nii.gz
        BraTS2021_AAAAA_t1.nii.gz
        BraTS2021_AAAAA_t1ce.nii.gz
        BraTS2021_AAAAA_t2.nii.gz
        BraTS2021_AAAAA_seg.nii.gz
    - BraTS2021_AAAAB/
        BraTS2021_AAAAB_flair.nii.gz
        BraTS2021_AAAAB_t1.nii.gz
        BraTS2021_AAAAB_t1ce.nii.gz
        BraTS2021_AAAAB_t2.nii.gz
        BraTS2021_AAAAB_seg.nii.gz
    ...

Preprocess the dataset as follows:

python data_preprocess.py --data_dir "path_to_dataset"

[For training datasets only]: Crop and save data as .npy files via:

python data_crop_npy.py --src_folder "path_to_src_folder" --dst_folder "path_to_dst_folder"

Perform train/val split on the training dataset as needed. Place the newly generated splits in different folders, as the network's training session differentiates between training and validation splits according to where it is placed (defined in config.yaml file).

Training/Testing sessions can be executed in two ways: anaconda and docker. Training/Testing sessions are guided by the configuration defined in the config.yaml file.

If you have Anaconda set-up on your system, install PyTorch (tested with 1.8.1, CUDA 11.1). Install the rest of the dependencies with:

pip install -r requirements.txt

Move to the e1d3 folder to execute the following commands.

Training:

python train.py --config config.yaml --gpu 0

Testing:

python test.py --config config.yaml --gpu 0

If you have Docker set-up with GPU support, build a docker image as follows:

docker build -t brats_e1d3 .

To execute the container for training/testing, provide paths in absolute format.

Training:

docker run -it --rm --gpus all -v "train_data_path":"train_data_path" -v "val_data_path":"val_data_path" -v "model_save_path":"model_save_path" -v "config_path":"config_path" brats_e1d3 --train --config "config_path/config.yaml" --gpu 0

Testing:

docker run -it --rm --gpus all -v "test_data_path":"test_data_path" -v "model_save_path":"model_save_path" -v "config_path":"config_path" brats_e1d3 --test --config "config_path/config.yaml" --gpu 0

[Note]: The paths set internally in the docker container should match those provided in config.yaml, as those will only be visible to the training/testing session.

If you found any part of this work useful, please cite as follows:

@misc{bukhari2021e1d3,
    title={E1D3 U-Net for Brain Tumor Segmentation: Submission to the RSNA-ASNR-MICCAI BraTS 2021 Challenge},
    author={Syed Talha Bukhari and Hassan Mohy-ud-Din},
    year={2021},
    eprint={2110.02519},
    archivePrefix={arXiv},
    primaryClass={eess.IV},
    url={https://arxiv.org/abs/2110.02519},
}