ResPr-UNet-3D-Denoising-Efficient-Pipeline-TF-keras

This repository contains data, code and results for the paper **A residual U-Net network with image prior for 3D image denoising, Proc. Eur. Signal Process. Conf. EUSIPCO, pp. 1264-1268, 2020. ** hal-02500664: https://hal.archives-ouvertes.fr/hal-02500664/document

We propose a modified U-Net architecture (ResPrU-Net) that exploits a prior image for 3D image denoising. The prior image is concatenated with the input and is connected to the output with a residual connection, as shown in the image below

The idea behind ResPrU-Net is based a previously proposed variational method SPADE (JFPJ Abascal et al, A sparse and prior based method for 3D image denoising, Proc. Eur. Signal Process. Conf. EUSIPCO, 2019, hal-02056591). In this work, we extend the idea of exploiting a prior image to learning approaches. SPADE -- a modification of TV method by including a functional that promotes sparsity wrt to the prior image -- was shown to improve convergence and maintain image texture of the prior image, being less sensitive to noise and artefacts. ResPrU-Net provided better results than U-Net for high noise.

In this repository, we provide and compare the proposed method to 2D and 3D variational and deep learning approaches.

Learning approaches are trained using efficient tensorflow pipelines for image datasets, based on keras and tf.data (see below for more details).

Tensorboard callbacks have been used to visualize losses, metrics, denoised images and histograms of the learned weights during training. TensorFlow Profiler is also used to profile the execution of the code.

Data

Data provided in this repository consist of human thorax CT scans (kits19 dataset, https://kits19.grand-challenge.org/) corrupted with additive Gaussian noise.

A subsample of kits19 data can be directly dowloaded from https://www.dropbox.com/sh/3sozaz9m7pyxcsa/AAAx1VULqc3T4rYC55vTiuJta?dl=0 (in png format, data with loosy conversion from float64 to unit8; fetch original data from the kits19 link):

'\Data\Train': All slices for 6 subjects.
'\Data\Test': All slices for 2 subjects.
'Kits19_train_data_200subj_2kimgs_2pcnoise_crop_rnd.npz': 2k slices randomly selectec from 200 subjects.
'Kits19_test_data_2subj_2pcnoise_crop_rnd.npz': Data used for test corresponding to two subjects.

Realistic synthetic and experimental spectral CT knee data used in the publication will be provided later on.

Installation

There are two options to try out the code:

Use the google colab notebook for training or assessing results. Just open on your google drive and it is ready to use the free google cloud service with free GPU!
Use provided Python code under tensorflow with keras. The python code has been tested under Python 3.8 and tensorflow 2.4.1. Installation guidelines below.

Different TF versions required specific cudatoolkit version. With anaconda one can handle several cuda versions in one machine. To install the tested version on Windows (should work also in linux):

conda create -n tf2 anaconda python=3.8 
conda install cudatoolkit=11.0 cudnn=8.0 -c=conda-forge
pip install --upgrade tensorflow-gpu==2.4.1 
python
import tensorflow as tf 
tf.test.is_gpu_available()

For TF compatibilities: https://www.tensorflow.org/install/source_windows For setting up Tensorflow-GPU with Cuda and Anaconda on Windows: https://towardsdatascience.com/setting-up-tensorflow-gpu-with-cuda-and-anaconda-onwindows-2ee9c39b5c44

Summary of results

2D learning-based approaches

The following figures show results for U-Net (with 32 and 64 filters on the first layer; 332,641 parameters) and a simple 3-layer ConvNet (37,633 parameters), for 2 % additive Gaussian noise. Unless specified, training data is comprised of 2000 slices taking randomly from all 200 subjects. Similar results were obtained by training on 10000 (10k) slices. Worse results corresponded to considering only 6 subjects (6sub, ~2k slices). U-Net 32 trained on 1h30 min and the ConvNet in 28 min. Train data file: (Kits19_train_data_200subj_2kimgs_2pcnoise_crop_rnd.npz. Test data file: Kits19_test_data_2subj_2pcnoise_crop_rnd.npz.

The following figures show results for U-Net and a simple 3-layer ConvNet (trained on 2k slices), split Bregman TV and BM3D, for 5 % additive Gaussian noise.

3D learning-based approaches

The following figures show results for 2D CNN and U-Net and 3D U-Net for 10 % additive Gaussian noise. Figure from Demo_assess_compare_models.ipynb.

Repository files

The repository contains the following files:

Python functions

Train_model_denoising2D.ipynb: Google colab notebook to train UNet model: Fetch data, build pipeline and train the network using google colab free cloud service with GPU
Demo_assess_compare_models.ipynb: Google colab notebook to assess and compare 2D and 3D models
Train_model_denoising2D.py: Demo to train several 2D CNN models (simple CNN, simple ResNet, U-Net)
Train_model_denoising3D.py: (Demo to train 3D models (3D CNN and 3D U-Net)
Train_model_denoising_ResPrUNet.py: (Demo to train a 3D model for ResPr-UNet (under construction)
Demo_assess_2D_learning_methods.py: Demo to compared the trained 2D models
Demo_assess_3D_learning_methods.py: Demo to compared the trained 3D models
Train and assess ResPrUNet: Under construction
Demo_assess_sparsity_models.ipynb: Demo to assess sparsity models: TV and BM3D.

Trained models

Data and models for notebook comparison: Models trained on 10 % noise. Can be fetched from the notebook. Link: https://www.dropbox.com/s/occmmmvpsttby26/data_models_denoising_comparison.zip?dl=0
kits19_200subj_data2k_Convnet_64_noisepc2_model_best.h5, kits19_200subj_data2k_Convnet_64_noisepc5_model_best.h5: Simple ConvNet (3 layers, 64 filters) trained on 2% and 5% additive Gaussian noise for 2,000 slices randomly selected from 200 subjects.
kits19_200subj_data2k_UNet_32_noisepc2_model_best.h5, kits19_200subj_data2k_UNet_32_noisepc5_model_best.h5: U-Net (32 filters on the first layer) trained on 2% and 5% additive Gaussian noise for 2,000 slices randomly selected from 200 subjects.
kits19_200subj_32x256x256_3DUNet_32_noisepc5_model_best.h5: 3D U-Net (32 filters on the first layer) trained on 5% additive Gaussian noise, 200 subjects, randonmly cropped patches of 32x256x256, for 300 epochs.

Data API

Tensorflow data API automatizes the data pipeline, chaining transformations (preprocessing and data augmentation), shuffling data. While the model is being trained for a given batch on the GPU, data are being prepared to be ready for next batch on the GPU. On the tests carried out here, using .chache (keep images in memory) and .prefetch (fetch and prepare data for next iteration), iterations were 1.6 times faster. Using tf datasets, data (image slices) are automatically fetched and pre-processed from the given directories.

The current implementation is based on the following sources:
-Aurélien Géron, Hands-On Machine Learning with Scikit-Learn, Keras, and TensorFlow, 2nd Edition, O'Reilly Media, Inc. ISBN: 9781492032649

-https://www.tensorflow.org/api_docs/python/tf/data/Dataset

-https://medium.com/deep-learning-with-keras/tf-data-build-efficient-tensorflow-input-pipelines-for-image-datasets-47010d2e4330

-https://keras.io/examples/vision/super_resolution_sub_pixel/

-https://www.wouterbulten.nl/blog/tech/data-augmentation-using-tensorflow-data-dataset/

If you use this code, please reference any of the following two publications: A residual U-Net network with image prior for 3D image denoising, Proc. Eur. Signal Process. Conf. EUSIPCO, pp. 1264-1268, 2020. JFPJ Abascal et al, Material Decomposition in Spectral CT Using Deep Learning: A Sim2Real Transfer Approach, IEEE Access, vol. 9, pp. 25632-25647, 2021.

If you need to contact the author, please do so at juanabascal78@gmail.com

jabascal/ResPr-UNet-3D-Denoising-Efficient-Pipeline-TF-keras