.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "auto_examples/sampling/demo_ddrm.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        New to DeepInverse? Get started with the basics with the
        :ref:`5 minute quickstart tutorial <sphx_glr_auto_examples_basics_demo_quickstart.py>`.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_auto_examples_sampling_demo_ddrm.py:


Image reconstruction with a diffusion model
====================================================================================================

This code shows you how to use the DDRM diffusion algorithm :footcite:t:`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 :class:`deepinv.physics.DecomposablePhysics`).
* The noise is Gaussian with known standard deviation (i.e., the noise model is :class:`deepinv.physics.GaussianNoise`).

.. GENERATED FROM PYTHON SOURCE LINES 15-21

.. code-block:: Python

    import deepinv as dinv
    from deepinv.utils.plotting import plot
    import torch
    import numpy as np
    from deepinv.utils.demo import load_example








.. GENERATED FROM PYTHON SOURCE LINES 22-26

Load example image from the internet
--------------------------------------------------------------

This example uses an image of Messi.

.. GENERATED FROM PYTHON SOURCE LINES 26-32

.. code-block:: Python


    device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"

    x = load_example("messi.jpg", img_size=32).to(device)









.. GENERATED FROM PYTHON SOURCE LINES 33-37

Define forward operator and noise model
--------------------------------------------------------------

We use image inpainting as the forward operator and Gaussian noise as the noise model.

.. GENERATED FROM PYTHON SOURCE LINES 37-47

.. code-block:: 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),
    )









.. GENERATED FROM PYTHON SOURCE LINES 48-53

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 :ref:`denoisers <denoisers>` module.

.. GENERATED FROM PYTHON SOURCE LINES 53-56

.. code-block:: Python


    denoiser = dinv.models.DRUNet(pretrained="download").to(device)








.. GENERATED FROM PYTHON SOURCE LINES 57-64

Create the Monte Carlo sampler
--------------------------------------------------------------

We can now reconstruct a noisy measurement using the diffusion method.
We use the DDRM method from :class:`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.

.. GENERATED FROM PYTHON SOURCE LINES 64-69

.. code-block:: 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)








.. GENERATED FROM PYTHON SOURCE LINES 70-73

Generate the measurement
---------------------------------------------------------------------------------
We apply the forward model to generate the noisy measurement.

.. GENERATED FROM PYTHON SOURCE LINES 73-76

.. code-block:: Python


    y = physics(x)








.. GENERATED FROM PYTHON SOURCE LINES 77-81

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.

.. GENERATED FROM PYTHON SOURCE LINES 81-96

.. code-block:: 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"])




.. image-sg:: /auto_examples/sampling/images/sphx_glr_demo_ddrm_001.png
   :alt: measurement, ground truth, DDRM reconstruction
   :srcset: /auto_examples/sampling/images/sphx_glr_demo_ddrm_001.png
   :class: sphx-glr-single-img


.. rst-class:: sphx-glr-script-out

 .. code-block:: none


      0%|          | 0/9 [00:00<?, ?it/s]
     33%|███▎      | 3/9 [00:00<00:00, 26.78it/s]
     67%|██████▋   | 6/9 [00:00<00:00, 26.80it/s]
    100%|██████████| 9/9 [00:00<00:00, 26.91it/s]
    100%|██████████| 9/9 [00:00<00:00, 26.86it/s]
    Linear reconstruction PSNR: 8.55 dB
    Diffusion PSNR: 24.38 dB




.. GENERATED FROM PYTHON SOURCE LINES 97-104

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 :class:`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.

.. GENERATED FROM PYTHON SOURCE LINES 104-108

.. code-block:: Python


    f = dinv.sampling.DiffusionSampler(diff, max_iter=10)









.. GENERATED FROM PYTHON SOURCE LINES 109-113

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.

.. GENERATED FROM PYTHON SOURCE LINES 113-141

.. code-block:: 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",
        ],
    )




.. image-sg:: /auto_examples/sampling/images/sphx_glr_demo_ddrm_002.png
   :alt: measurement, ground truth, post. mean, post. std, abs. error
   :srcset: /auto_examples/sampling/images/sphx_glr_demo_ddrm_002.png
   :class: sphx-glr-single-img


.. rst-class:: sphx-glr-script-out

 .. code-block:: none


      0%|          | 0/10 [00:00<?, ?it/s]
     10%|█         | 1/10 [00:00<00:03,  2.69it/s]
     20%|██        | 2/10 [00:00<00:02,  2.69it/s]
     30%|███       | 3/10 [00:01<00:02,  2.70it/s]
     40%|████      | 4/10 [00:01<00:02,  2.70it/s]
     50%|█████     | 5/10 [00:01<00:01,  2.70it/s]
     60%|██████    | 6/10 [00:02<00:01,  2.70it/s]
     70%|███████   | 7/10 [00:02<00:01,  2.70it/s]
     80%|████████  | 8/10 [00:02<00:00,  2.70it/s]
     90%|█████████ | 9/10 [00:03<00:00,  2.70it/s]
    100%|██████████| 10/10 [00:03<00:00,  2.70it/s]
    100%|██████████| 10/10 [00:03<00:00,  2.70it/s]
    Iteration 9, current converge crit. = 8.09E-03, objective = 1.00E-01 
    Iteration 9, current converge crit. = 9.86E-02, objective = 1.00E-01 
    Linear reconstruction PSNR: 8.55 dB
    Posterior mean PSNR: 22.80 dB




.. GENERATED FROM PYTHON SOURCE LINES 142-145

:References:

.. footbibliography::


.. rst-class:: sphx-glr-timing

   **Total running time of the script:** (0 minutes 4.789 seconds)


.. _sphx_glr_download_auto_examples_sampling_demo_ddrm.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: demo_ddrm.ipynb <demo_ddrm.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: demo_ddrm.py <demo_ddrm.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: demo_ddrm.zip <demo_ddrm.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_