Equivariant Imaging: Learning Beyond the Range Space

Equivariant Imaging: Learning Beyond the Range Space

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

The University of Edinburgh

In ICCV 2021 (oral)

Figure: Learning to image from only measurements. Training an imaging network through just measurement consistency (MC) does not significantly improve the reconstruction over the simple pseudo-inverse (). However, by enforcing invariance in the reconstructed image set, equivariant imaging (EI) performs almost as well as a fully supervised network. Top: sparse view CT reconstruction, Bottom: pixel inpainting. PSNR is shown in top right corner of the images.

EI is a new self-supervised, 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.

Get quickly started

• Please find the presentation video and blog post for a quick introduction of EI.
• Please find the core implementation of EI at './ei/closure/ei.py' (ei.py).
• Please find the 30 lines code get_started.py and the toy cs example to get started with EI.
• Please find the follow-up works: robust EI (CVPR2022 oral) and theorems paper.

Overview

The problem: Imaging systems capture noisy measurements of a signal through a linear operator + . We aim to learn the reconstruction function where

• NO groundtruth data for training as most inverse problems don’t have ground-truth;
• only a single forward operator is available;
• has a non-trivial nullspace (e.g. ).

The challenge:

• We have NO information about the signal set outside the range space of or .
• It is IMPOSSIBLE to learn the signal set using alone.

The motivation:

We assume the signal set has a low-dimensional structure and is invariant to a groups of transformations (orthgonal matrix, e.g. shift, rotation, scaling, reflection, etc.) related to a group , such that and the sets and are the same. For example,

• natural images are shift invariant.
• in CT/MRI data, organs can be imaged at different angles making the problem invariant to rotation.

Key observations:

• Invariance provides access to implicit operators with potentially different range spaces: where and . Obviously, should also in the signal set.
• The composition is equivariant to the group of transformations : .

Figure: Learning with and without equivariance in a toy 1D signal inpainting task. The signal set consists of different scaling of a triangular signal. On the left, the dataset does not enjoy any invariance, and hence it is not possible to learn the data distribution in the nullspace of . In this case, the network can inpaint the signal in an arbitrary way (in green), while achieving zero data consistency loss. On the right, the dataset is shift invariant. The range space of is shifted via the transformations , and the network inpaints the signal correctly.

Equivariant Imaging: to learn by using only measurements , all you need is to:

• Define:

1. define a transformation group based on the certain invariances to the signal set.
2. define a neural reconstruction function , e.g. where is the (approximated) pseudo-inverse of and is a UNet-like neural net.

• Calculate:

1. calculate as the estimation of .
2. calculate by transforming .
3. calculate by reconstructing from its measurement .

• Train: finally learn the reconstruction function by solving: +

Requirements

• PyTorch (1.6)

All used packages are listed in the Anaconda environment.yml file. You can create an environment and run

conda env create -f environment.yml

Test

We 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. Then evaluate the trained models by running

python3 demo_test_inpainting.py

and

python3 demo_test_ct.py

Train

To train EI for a given inverse problem (inpainting or CT), run

python3 demo_train.py --task 'inpainting'

or run a bash script to train the models for both CT and inpainting tasks.

bash train_paper_bash.sh

Train your models

To train your EI models on your dataset for a specific inverse problem (e.g. inpainting), run

python3 demo_train.py --h

• Note: you may have to implement the forward model (physics) if you manage to solve a new inverse problem.
• Note: you only need to specify some basic settings (e.g. the path of your training set).

Citation

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