On the Posterior Distribution in Denoising: Application to Uncertainty Quantification Official Implementation [ICLR 2024]
Hila Manor and Tomer Michaeli
Technion - Israel Institute of Technology
This repository contains the code release for On the Posterior Distribution in Denoising: Application to Uncertainty Quantification Official Implementation.
Animations_Intro.mp4
python -m pip install -r requirements.txt
Note that for DDPM (faces), their code uses Open-MPI which sometimes have problems installing on machines with conda installed.
We support a number of pre-trained models, and one can clone their repo and download their checkpoints as necessary.
As this is a simple network we built and trained, the checkpoint is already included in the repo.
This repo contains the implementations for multiple denoisers that we support: DnCNN
, IRCNN
, SwinIR
.
-
Clone their repo (provide this path to the
--model_zoo
parameter)git clone https://github.com/cszn/KAIR.git
-
Follow the instructions in
KAIR/model_zoo/README.md
to download the wanted checkpoints or here for SwinIR.- We used
colorDN_DFWB_s128w8_SwinIR-M
models, but the interface should be able to use most versions)
- We used
The original official implementation of Noise2Void is not in python. Nevertheless, the authors later published Probabilistic Noise2Void, and with its (python) GitHub implementation also included the python version of N2V.
This version of the code needs some fixes in their original code, and therefore is provided in the repo. Our trained checkpoint over the FMD data and the training notebooks are also included in the local pn2v repo.
- Run the
GetData
notebook to download the FMD dataset, extract the images and preprocess them for n2v.
We use Label-Efficient Semantic Segmentation with Diffusion Models's checkpoint for DDPM trained on the entire FFHQ dataset, and tested on celebA (as is usually done with the faces domain and diffusion models).
The relevant version of guided_diffusion used is already included in this repo, and therefore:
- Follow their
download_checkpoint.sh
to downloadffhq.pt
, and place it inDDPM_FFHQ
. - If needed, follow their
download_datasets.sh
to download celebA images.
python main.py -e <number of eigenvectors> -p <context size around the patch> -t <subspace iters> -c <small constant> -o <output folder> -d <denoiser model> -i <input image path>
Use --help
for more information on the parameters and other options, such as low_acc
for finding EVs only quickly (without calculating the moments for the marginal distribution), or use_poly
to try and fit a polynomial for the moments calculation.
Use -v
to calculate the higher-order moments and estimate the density along the PCs.
FacesAnimations.mp4
If you use this code for your research, please cite our paper:
@inproceedings{
manor2024posterior,
title={On the Posterior Distribution in Denoising: Application to Uncertainty Quantification},
author={Hila Manor and Tomer Michaeli},
booktitle={The Twelfth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=adSGeugiuj}
}