Bayesian Image Reconstruction using Deep Generative Models
Razvan Marinescu, Daniel Moyer, Polina Golland
MIT CSAIL
- Pre-print: arxiv
- Code:
- Tensorflow
- PyTorch (+ variational inference)
- Demos through Colab notebooks:
Abstract
Machine learning models are commonly trained end-to-end and in a supervised setting, using
paired (input, output) data. Examples include recent super-resolution methods that train on
pairs of (low-resolution, high-resolution) images. However, these end-to-end approaches
require re-training every time there is a distribution shift in the inputs (e.g., night
images vs daylight) or relevant latent variables (e.g., camera blur or hand motion). In
this work, we leverage state-of-the-art (SOTA) generative models (here StyleGAN2) for
building powerful image priors, which enable application of Bayes' theorem for many
downstream reconstruction tasks. Our method, Bayesian Reconstruction through Generative
Models (BRGM), uses a single pre-trained generator model to solve different image restoration
tasks, i.e., super-resolution and in-painting, by combining it with different forward
corruption models. We keep the weights of the generator model fixed, and reconstruct the
image by estimating the Bayesian maximum a-posteriori (MAP) estimate over the input latent
vector that generated the reconstructed image. We further use variational inference to
approximate the posterior distribution over the latent vectors, from which we sample multiple
solutions. We demonstrate BRGM on three large and diverse datasets: (i) 60,000 images from
the Flick Faces High Quality dataset (ii) 240,000 chest X-rays from MIMIC III and (iii) a
combined collection of 5 brain MRI datasets with 7,329 scans. Across all three datasets and
without any dataset-specific hyperparameter tuning, our simple approach yields performance
competitive with current task-specific state-of-the-art methods on super-resolution and
in-painting, while being more generalisable and without requiring any training. Our source
code and pre-trained models are available online:
News
- Next:
- 3D reconstruction for medical scans
- More reconstruction tasks: e.g. MRI compressed sensing
- June 2021: Reimplemented the method in Pytorch, and switched to StyleGAN-ADA.
- May 2021: Added variational inference extension for sampling multiple solutions. Updated methods section in arXiv paper. Also included qualitative comparisons against Deep Image Prior in supplement.
- Feb 2021: Updated methods section in arXiv paper. We now start from the full Bayesian formulation, and derive the loss function from the MAP estimate (in appendix), and show the graphical model. Code didn’t change in this update.
- Dec 2020: Code and pre-trained models published on github
- Nov 2020: Uploaded article pre-print to arXiv.
If you use our model, please cite:
@article{marinescu2020bayesian,
title={Bayesian Image Reconstruction using Deep Generative Models},
author={Marinescu, Razvan V and Moyer, Daniel and Golland, Polina},
journal={arXiv preprint arXiv:2012.04567},
year={2020}
}
Main results, as twitter thread
BRGM - Curated tweets by RazMarinescu
Mistake fixed in the pre-print (full thread on Twitter)
In science, it's very important to admit mistakes and correct them. I recently made a claim in my Bayesian image reconstruction paper that turned out to be wrong /1#ScienceMistakes
— Razvan Marinescu (@RazMarinescu) February 23, 2021
Funding
This project received funding from the NIH grants NIBIB NAC P41EB015902 and NINDS R01NS086905.