johncolby / unet_brats

3D U-Net deep learning segmentation of BraTS 2018 data with mxnet

README

This python package contains a 3D U-Net algorithm (a type of deep learning convolutional neural network based on Ronneberger et al., 2015) for the segmentation of voxelwise medical imaging data. It is geared toward the BraTS 2018 glioma brain tumor dataset, although the underlying network is broadly applicable to other voxelwise multi input channel and multi output class applications. It is implemented in the mxnet deep learning framework, using code adapted from the The Straight Dope book.

Install

1. Make sure a version of mxnet is installed with GPU support. For example:

   pip install mxnet-cu92mkl

2. Download this repository, either with git clone <URL> (where <URL> is the git repository URL) or by clicking the download link in the git web interface.

3. Install the companion python module. From the command line:

   cd /path/to/unet_brats/
   pip install .
   

   Alternatively, install directly from the git repository like:

   pip install git+<URL>

Usage

Example Jupyter notebooks for training and testing/inference are included in the notebooks directory.

Performance

This project aims to create a good implementation of a simple "vanilla" U-Net, which may be used as a basic foundation for more complex analyses. Nevertheless, despite the focus on simplicity, it achieves competitive performance out-of-the-box, for example achieving top 10 results in whole tumor segmentation at the time of this writing. From the BraTS 2018 Leaderboard:

Team	Dice_ET	Dice_WT	Dice_TC
Deep_Rad_UCSF	0.78829	0.90918	0.84132

Example output