Robust Equivariant Imaging (REI) in PyTorch

Robust Equivariant Imaging: a fully unsupervised framework for learning to image from noisy and partial measurements

Dongdong Chen, Julián Tachella, Mike E. Davies.

The University of Edinburgh, UK

In CVPR 2022 (oral)

Background

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.

Background: Equivariant Imaging (EI)

Figure 1: Equivariant imaging systems. If the set of signals is invariant to a certain set of transformations, the composition of imaging operator ( $A$ ) with the reconstruction function ( $f_\theta$ ) should be equivariant to these transformations.

EI is the first self-supervised learning framework that exploits the group invariance resent 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.
Given an inverse problem, EI learns the reconstruction function with NO need for either multiple forward operators or extra masking measurement data into multiple complementary/overlapped parts. EI is agnostic to neural network architecture. Please find our blog post and presentation video for a quick introduction of EI.

Robust Equivariant Imaging (REI)

Figure 2: REI training strategy. $x^{(1)}$ represents the estimated image, $T_g$ is the transformation, while $x^{(2)}$ and $x^{(3)}$ represent $T_gx^{(1)}$ and the estimate of $x^{(2)}$ from the (noisy) measurements $\tilde{y} = A (x^{(2)})$ respectively. The SURE loss aims to estimate the measurement consistency of clean measurement, REQ (robust equivariance) loss is the error (e.g. MSE) between $x^{(2)}$ and $x^{(3)}$ .

Motivation: while EI results are impressive and successfully solved the challenge of learning to image without groundtruth, its performance degrades with increasing measurement noise (Figure 2).
Main idea: we propose to employ 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).
Performance: REI can obatin considerable performance gains on linear (e.g. MRI, Inpainting) and nonlinear inverse problems (e.g. CT), thereby paving the way for robust unsupervised imaging with deep networks (Figure 4).
Remark 1: while we evaluated REI on the 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.
Remark 2: (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 $\sigma$ = 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 $A^{\dagger}y$ , EI, REI, Sup and the groundtruth on the non-linear CT (with MPG noise) and Inpainting (with Poisson noise) tasks, respectively.

Run the code

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):
- REI for the accelerated MRI task and the Inpainting task: rei_end2end.py
- REI for the low-dose and sparse-view CT task: rei_end2end_ct.py
- Find our implementation of SURE for Gaussian and Poisson noise models at: rei_end2end.py
- Find our implementation of SURE for Mixed Poisson-Gaussian noise model at: rei_end2end_ct.py
Download datasets from the below source and move them under the folders: ./dataset/mri, ./dataset/Urban100, and ./dataset/CT, repectively:
- fastMRI (only the subset 'Knee MRI'): https://fastmri.med.nyu.edu/
- Urban100: https://uofi.box.com/shared/static/65upg43jjd0a4cwsiqgl6o6ixube6klm.zip
- CT100: https://www.kaggle.com/kmader/siim-medical-images
Train: run the below scripts to train REI models:
- run ./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.
- or run 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
- step 1: implement their own forward model (physics of sensing model)
- step 2: determine the transformation group
- step 3: specify the path of new dataset

Citation

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}}

daydreamer2023 / REI