#### NOTE
New to DeepInverse? Get started with the basics with the
[5 minute quickstart tutorial](https://deepinv.org/auto_examples/basics/demo_quickstart.html.md#sphx-glr-auto-examples-basics-demo-quickstart-py)..
# Image reconstruction with a diffusion model
This code shows you how to use the DDRM diffusion algorithm Kawar *et al.*[1](#footcite-kawar2022denoising) to reconstruct images and also compute the
uncertainty of a reconstruction from incomplete and noisy measurements.
The DDRM method requires that:
* The operator has a singular value decomposition (i.e., the operator is a [`deepinv.physics.DecomposablePhysics`](https://deepinv.org/api/stubs/deepinv.physics.DecomposablePhysics.html.md#deepinv.physics.DecomposablePhysics)).
* The noise is Gaussian with known standard deviation (i.e., the noise model is [`deepinv.physics.GaussianNoise`](https://deepinv.org/api/stubs/deepinv.physics.GaussianNoise.html.md#deepinv.physics.GaussianNoise)).
```Python
import deepinv as dinv
from deepinv.utils.plotting import plot
import torch
import numpy as np
from deepinv.utils import load_example
```
## Load example image from the internet
This example uses an image of Messi.
```Python
device = dinv.utils.get_device()
x = load_example("messi.jpg", img_size=32).to(device)
```
```none
Selected GPU 0 with 24087.25 MiB free memory
```
## Define forward operator and noise model
We use image inpainting as the forward operator and Gaussian noise as the noise model.
```Python
sigma = 0.1 # noise level
physics = dinv.physics.Inpainting(
mask=0.5,
img_size=x.shape[1:],
device=device,
noise_model=dinv.physics.GaussianNoise(sigma=sigma),
)
```
## Define the MMSE denoiser
The diffusion method requires an MMSE denoiser that can be evaluated a various noise levels.
Here we use a pretrained DRUNET denoiser from the [denoisers](https://deepinv.org/user_guide/reconstruction/denoisers.html.md#denoisers) module.
```Python
denoiser = dinv.models.DRUNet(pretrained="download").to(device)
```
## Create the Monte Carlo sampler
We can now reconstruct a noisy measurement using the diffusion method.
We use the DDRM method from [`deepinv.sampling.DDRM`](https://deepinv.org/api/stubs/deepinv.sampling.DDRM.html.md#deepinv.sampling.DDRM), which works with inverse problems that
have a closed form singular value decomposition of the forward operator.
The diffusion method requires a schedule of noise levels `sigmas` that are used to evaluate the denoiser.
```Python
sigmas = np.linspace(1, 0, 100) if torch.cuda.is_available() else np.linspace(1, 0, 10)
diff = dinv.sampling.DDRM(denoiser=denoiser, etab=1.0, sigmas=sigmas, verbose=True)
```
## Generate the measurement
We apply the forward model to generate the noisy measurement.
```Python
y = physics(x)
```
## Run the diffusion algorithm and plot results
The diffusion algorithm returns a sample from the posterior distribution.
We compare the posterior mean with a simple linear reconstruction.
```Python
xhat = diff(y, physics)
# compute linear inverse
x_lin = physics.A_adjoint(y)
# compute PSNR
print(f"Linear reconstruction PSNR: {dinv.metric.PSNR()(x, x_lin).item():.2f} dB")
print(f"Diffusion PSNR: {dinv.metric.PSNR()(x, xhat).item():.2f} dB")
# plot results
error = (xhat - x).abs().sum(dim=1).unsqueeze(1) # per pixel average abs. error
imgs = [x_lin, x, xhat]
plot(imgs, titles=["measurement", "ground truth", "DDRM reconstruction"])
```
```none
0%| | 0/99 [00:00
## Create a Monte Carlo sampler
Running the diffusion gives a single sample of the posterior distribution.
In order to compute the posterior mean and variance, we can use multiple samples.
This can be done using the [`deepinv.sampling.DiffusionSampler`](https://deepinv.org/api/stubs/deepinv.sampling.DiffusionSampler.html.md#deepinv.sampling.DiffusionSampler) class, which converts
the diffusion algorithm into a fully fledged Monte Carlo sampler.
We set the maximum number of iterations to 10, which means that the sampler will run the diffusion 10 times.
```Python
f = dinv.sampling.DiffusionSampler(diff, max_iter=10)
```
## Run sampling algorithm and plot results
The sampling algorithm returns the posterior mean and variance.
We compare the posterior mean with a simple linear reconstruction.
```Python
mean, var = f(y, physics)
# compute PSNR
print(f"Linear reconstruction PSNR: {dinv.metric.PSNR()(x, x_lin).item():.2f} dB")
print(f"Posterior mean PSNR: {dinv.metric.PSNR()(x, mean).item():.2f} dB")
# plot results
error = (mean - x).abs().sum(dim=1).unsqueeze(1) # per pixel average abs. error
std = var.sum(dim=1).unsqueeze(1).sqrt() # per pixel average standard dev.
imgs = [
x_lin,
x,
mean,
std / std.flatten().max(),
error / error.flatten().max(),
]
plot(
imgs,
titles=[
"measurement",
"ground truth",
"post. mean",
"post. std",
"abs. error",
],
)
```
```none
0%| | 0/10 [00:00
* **References:**
* Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. *Advances in Neural Information Processing Systems*, 35:23593–23606, 2022.
**Total running time of the script:** (0 minutes 3.442 seconds)
🏷 Tags: Denoising, Inpainting
[`Download Jupyter notebook: demo_ddrm.ipynb`](https://deepinv.org/auto_examples/sampling/_downloads/9bf9c35d1850ffbe933dedaa0834a74c/demo_ddrm.ipynb)
[`Download Python source code: demo_ddrm.py`](https://deepinv.org/auto_examples/sampling/_downloads/25333f6f97b87a771a7d18541fd7a870/demo_ddrm.py)
[`Download zipped: demo_ddrm.zip`](https://deepinv.org/auto_examples/sampling/_downloads/9496ab938a025bc672041f7e6142ac5d/demo_ddrm.zip)
[Gallery generated by Sphinx-Gallery](https://sphinx-gallery.github.io)