-
JIA Shuman authoredJIA Shuman authored
Implementation of U-Net with Keras
This implementation builds a U-Net based on the model proposed by Isensee et al. in 2017: https://arxiv.org/abs/1802.10508
Original implementation: https://github.com/ellisdg/3DUnetCNN
Then, we implement distance related loss functions in the model.
Getting Started
Save datasets in folder: /train_set
Create for each subject a subfolder with image and ground trush mask.
e.g. ./train_set/subject1/img.nii.gz, ./train_set/subject1/mask.nii.gz
Prerequisites
Tensorflow, Keras, Numpy, SimpleItk, pydot, keras_contrib (tensorflow 1.6 with cuda 9.1)
conda activate **your_env**
conda install numpy
conda install -c conda-forge pydicom
conda install -c simpleitk simpleitk
conda install -c conda-forge tensorflow
conda install -c anaconda pydot
conda install keras
cd ***your_path_pkg_keras-contrib***
git clone https://www.github.com/keras-team/keras-contrib.git
cd keras-contrib
python setup.py install
Model
Check Configurations: inputs, image shape, validation split, model options, fit generator options, filenames.
The visualisation of net will be ploted to model.png.
Running
Run the training with
python Run.py
Datasets will be saved into .npy after validation split.
And then
the model will be saved as Net.hdf5 after each epoch,
training log of loss and val_loss is avaliable in training.log.
Calculate Dice Score for Validation Datasets
python Evaluate.py
This step saves prediction results for validation datasets in ./prediction
Also, average Dice Score (2*intersection/(gt+pred)) will appear in terminal.
2-Stage U-Net
python Crop.py
This step save cropped images around predicted segmentations in ./cropped
mkdir test_lowRedo
mv model.png training.log Net.hdf5 training_imgs.npy training_segs.npy training_subjects.txt validation_imgs.npy validation_segs.npy validation_subjects.txt test_lowRedo
save models and data
change options for input in Model.py
sed -i 's/train_set/cropped/g' Model.py
if use Contour loss:
python create_distance_map.py
sed -i 's/config\[\"option_loss\"\] = 1/config\[\"option_loss\"\] = 0/g' Model.py
Then, repeat the training and predicting process
python Run.py
python Evaluate.py
Authors
Original implementation: https://github.com/ellisdg/3DUnetCNN
Zihao Wang, Antoine Despinasse participated in this project.
Citation
If the code is used, please cite the paper https://hal.inria.fr/hal-01860285v4/document:
Shuman Jia, Antoine Despinasse, Zihao Wang, Hervé Delingette, Xavier Pennec, et al.. Automatically Segmenting the Left Atrium from Cardiac Images Using Successive 3D U-Nets and a Contour Loss. Statistical Atlases and Computational Modeling of the Heart (STACOM) workshop, Sep 2018, Granada, Spain. ⟨hal-01860285v4⟩