A Tensorflow Implementation of Brain Tumor Segmentation using Topological Loss
This is a Tensorflow implementation of Topological and Smoothing losses. Source Code is taken from here
This repository provides source code and pre-trained models for brain tumor segmentation with BraTS dataset. The method is detailed in [1].
This implementation is based on NiftyNet and Tensorflow. While NiftyNet provides more automatic pipelines for dataloading, training, testing and evaluation, this naive implementation only makes use of NiftyNet for network definition, so that it is lightweight and extensible. A demo that makes more use of NiftyNet for brain tumor segmentation is proivde at https://cmiclab.cs.ucl.ac.uk/CMIC/NiftyNet/tree/dev/demos/BRATS17
If you use any resources in this repository, please cite the following papers:
An example of brain tumor segmentation result.
A CUDA compatable GPU with memoery not less than 6GB is recommended for training. For testing only, a CUDA compatable GPU may not be required.
Tensorflow. Install tensorflow following instructions from https://www.tensorflow.org/install/
NiftyNet. Install it by following instructions from http://niftynet.readthedocs.io/en/dev/installation.html or simply typing:
pip install niftynet
data_root/BRATS2015_Training
or data_root/Brats17TrainingData
and the validation set will be in data_root/BRATS2015_Validation
or data_root/Brats17ValidationData
.The trainig process needs 9 steps, with axial view, sagittal view, coronal view for whole tumor, tumor core, and enhancing core, respectively.
The following commands are examples for BraTS 2017. However, you can edit the corresponding *.txt
files for different configurations.
python modified_train.py config17/train_wt_ax.txt
python modified_train.py config17/train_wt_sg.txt
python modified_train.py config17/train_wt_cr.txt
python modified_train.py config17/train_tc_ax.txt
python modified_train.py config17/train_tc_sg.txt
python modified_train.py config17/train_tc_cr.txt
python modified_train.py config17/train_en_ax.txt
python modified_train.py config17/train_en_sg.txt
python modified_train.py config17/train_en_cr.txt
python util/rename_variables.py
You may need to edit this file to set different parameters. As an example for Brats 2015, after running this command, you will see a model named model15/msnet_tc32sg_init
that is copied from model15/msnet_tc32_20000.ckpt
. Then just set start_iteration=1 and model_pre_trained=model15/msnet_tc32sg_init in config15/train_tc_sg.txt
.
Write a configure file that is similar to config15/test_all_class.txt
or config17/test_all_class.txt
and
set the value of model_file to your own model files. Run:
python test.py your_own_config_for_test.txt
Calcuate dice scores between segmentation and the ground truth, run:
python util/evaluation.py
You may need to edit this file to specify folders for segmentation and ground truth.