Mentions légales du service

Skip to content
Snippets Groups Projects

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⟩