Dongdong Chen, Julián Tachella, Mike E. Davies.
The University of Edinburgh, UK
In CVPR 2022 (oral)
Deep networks provide state-of-the-art performance in multiple imaging inverse problems ranging from medical imaging to computational photography. However, most existing networks are trained with clean signals which are often hard or impossible to obtain. This work aims to solve the challenge: learn the reconstruction function from noisy and partial measurements alone. Please find our presentation video for a quick introduction.
Figure 1: Equivariant imaging systems. If the set of signals is invariant to a certain set of transformations, the composition of imaging operator () with the reconstruction function () should be equivariant
to these transformations.
self-supervised
learning framework that exploits the group invariance
present in signal distributions to learn a reconstruction function from partial measurement data alone (Figure 1). EI is end-to-end
and physics-based
learning framework for inverse problems with theoretical guarantees which leverages simple but fundamental priors about natural signals: symmetry
and low-dimensionality
. EI is agnostic to neural network architecture
. Please find our blog post and presentation video for a quick introduction of EI.Figure 2: REI training strategy. represents the estimated image, is the transformation, while and represent and the estimate of from the (noisy) measurements respectively. The SURE
loss aims to estimate the measurement consistency of clean measurement, REQ
(robust equivariance) loss is the error (e.g. MSE) between and .
Stein's Unbiased Risk Estimator (SURE)
to obtain a fully unsupervised training loss that is robust to noise, i.e. have an unbiased SURE estimator to the measurement consistency loss of clean measurements. With the SURE loss and the EI objective, our proposed REI framework can learn to image from noisy partial measurements alone (Figure 3, Figure 4). Gaussian
, Poisson
and Mixed Poisson-Gaussian (MPG)
models, SURE can handle many other models including non-exponential ones, see Raphan et al. for a detailed list. By this repo, we believe one can implement other noise models accordingly without giant changes.(R)EI is agnostic to neural network architecture
-- one can employ (R)EI to train any existed imaging networks to achieve fully unsupervised learning to image without changing the architectures. In addition to our demonstrated applications of REI (EI) on image inpainting, CT and MRI image reconstruction tasks, REI (EI) can be used to achieve new and fully unsupervised learning solutions to other inverse problems in computer vision and scientific imaging tasks, especially the cases when no groundtruth data is available for training.Figure 3: Motivation. The performance of EI degrades with increasing noise. From top to bottom: reconstructions of EI, supervised (Sup) baseline, and the proposed REI on 4× accelerated MRI with Gaussian
noise level = 0.01, 0.1, 0.2. PSNR values are shown in the top right corner of the images
Figure 4: More results. From top to bottom: reconstruction of , EI, REI, Sup and the groundtruth on the non-linear CT (with MPG
noise) and Inpainting (with Poisson
noise) tasks, respectively.
We collected some Frequently Asked Questions, please find the above Q & A.
Requirements: configure the environment by following: environment.yml to run Inpainting and CT experiments. To run MRI experiments, please install the 'fastmri' package by pip install fastmri
.
Find the implementation of Robust Equivariant Imaging (REI):
accelerated MRI
task and the Inpainting
task: rei_end2end.pylow-dose and sparse-view CT
task: rei_end2end_ct.pySURE
for Gaussian
and Poisson
noise models at: rei_end2end.pySURE
for Mixed Poisson-Gaussian
noise model at: rei_end2end_ct.pyDownload datasets from the below source and move them under the folders: ./dataset/mri
, ./dataset/Urban100
, and ./dataset/CT
, repectively:
Train: run the below scripts to train REI models:
./demo_scripts/demo_mri.py
, ./demo_scripts/demo_inpainting.py
, ./demo_scripts/demo_ct.py
to train REI for MRI, Inpainting, and CT tasks, respectively.train_bash.sh
to train REI models on all tasks.bash train_bash.sh
Test: run 'demo_test.py' to test the performance (PSNR) of a trained model on a specific task.
python3 demo_test.py
We also provide the trained models used in the paper which can be downloaded at Google Drive. Please put the downloaded folder 'ckp' in the root path.
To solve a new inverse problem, one only needs to
If you use this code for your research, please cite our papers.
@inproceedings{chen2021equivariant,
title = {Equivariant Imaging: Learning Beyond the Range Space},
author = {Chen, Dongdong and Tachella, Juli{\'a}n and Davies, Mike E},
booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
month = {October},
year = {2021},
pages = {4379-4388}}
@inproceedings{chen2022robust,
title = {Robust Equivariant Imaging: a fully unsupervised framework for learning to image from noisy and partial measurements},
author = {Chen, Dongdong and Tachella, Juli{\'a}n and Davies, Mike E},
booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
year = {2022}}