/HeartSeg

code of DenseVoxNet and 3D-DSN

Primary LanguageMATLAB

DenseVoxNet and 3D-DSN

  • Automatic 3D Cardiovascular MR Segmentation with Densely-Connected Volumetric ConvNets, MICCAI 2017
  • 3D Deeply Supervised Network for Automated Segmentation of Volumetric Medical Images, MedIA 2017

Introduction

This repository includes the code (training and testing) for our papers about DenseVoxNet and 3D-DSN. The code is based on 3D-CNN for volumetric segmentation, and the 3D operation is implemented with 3D-Caffe library. You can simply extend this code to other volumetric segmentation tasks.

Installation

  1. Follow the instruction of 3D-Caffe to install Caffe library. Rememebr to use git checkout 3D-Caffe to switch the correct branch and use make matcaffe to build Matlab interface.

  2. As an alternative, you can

git clone https://github.com/yulequan/HeartSeg
cd 3D-Caffe
cp Makefile.config.example Makefile.config
vim Makefile.config
##uncomment USE_CUDNN := 1 if you want to use CuDNN
make -j8
make matcaffe
  1. Install Matlab toolbox NIfTI_tools in order to read nii data format in Matlab.

Usage

  1. Download HVSMR dataset (phase 2) and put them in folder data.

  2. Prepare the hdf5 data to train the model.

cd code
#modify parameters in prepare_h5_data.m file
matlab prepare_h5_data.m
  1. Train the model
cd DenseVoxNet or cd 3D-DSN # Different folder for different network architecture
sh start_train.sh
  1. Test the model
cd code
#modify parameters and model path in test_model.m
matlab test_model.m

Note

  • We use HDF5DataLayer to read data. You need to generate the hdf5 data from original data type. See prepare_h5_data.m for more details.

  • Due to the limited GPU memory, we use patch-based training/testing strategy. The patch size is 64 by 64 by 64. You can modify it in network architecture and corresponding .m file.

  • In order to get the whole prediction, we provided two schemes (average fusion and major voting). Please see papers for more detail.

  • Generally, the batch normalization will accerate the training. We had better set the batch size (>1) to properly use batchnorm.

Citation

If you find this code is useful for your research, please consider citing:

@article{yu2017automatic, author = {Yu, Lequan and Cheng,Jie-Zhi and Dou, Qi and Yang, Xin and Chen, Hao and Qin, Jing and Heng, Pheng-Ann}, title = {Automatic 3D Cardiovascular MR Segmentation with Densely-Connected Volumetric ConvNets}, Journal = {MICCAI}, year = {2017} }

and/or

@article{dou20173d, author = {Dou, Qi and Yu, Lequan and Chen, Hao and Jin, Yueming and Yang, Xin and Qin, Jing and Heng, Pheng-Ann}, title = {3D deeply supervised network for automated segmentation of volumetric medical images}, journal = {Medical Image Analysis}, publisher = {Elsevier}, year = {2017} }

Questions

Please contact 'ylqzd2011@gmail.com' or 'lqyu@cse.cuhk.edu.hk'