#### 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)..
# Plug-and-Play algorithm with Mirror Descent for Poisson noise inverse problems.
This is a simple example to show how to use a mirror descent algorithm for solving an inverse problem with Poisson noise.
See Hurault *et al.*[1](#footcite-hurault2023convergent) for more details.
The Mirror descent with RED denoiser writes
$$
x_{k+1} = \nabla \phi ( \nabla \phi^*(x_k) - \tau \nabla \distance{A(x_k)}{y} - \tau ( x_k - D_\sigma(x)))
$$
where $\phi$ is a convex Bergman potential, $\distance{A(x)}{y}$ is the data fidelity term and $D_\sigma(x)$ is a denoiser.
In this example, we use the DnCNN denoiser Zhang *et al.*[2](#footcite-zhang2017beyond). As the observation has been corrupted with Poisson noise, we use the [`deepinv.optim.PoissonLikelihood`](https://deepinv.org/api/stubs/deepinv.optim.PoissonLikelihood.html.md#deepinv.optim.PoissonLikelihood) data-fidelity term.
In Bolte *et al.*[3](#footcite-bolte2016descent), it is shown that, with this data-fidelity term, the right Bregman potential to use is Burg’s entropy [`deepinv.optim.bregman.BurgEntropy`](https://deepinv.org/api/stubs/deepinv.optim.BurgEntropy.html.md#deepinv.optim.BurgEntropy).
```Python
import deepinv as dinv
from pathlib import Path
import torch
from deepinv.optim.data_fidelity import PoissonLikelihood
from deepinv.optim.prior import RED
from deepinv.optim import MD
from deepinv.optim.bregman import BurgEntropy
from deepinv.utils import load_example
from deepinv.utils.plotting import plot, plot_curves
```
## Setup paths for data loading and results.
```Python
BASE_DIR = Path(".")
ORIGINAL_DATA_DIR = BASE_DIR / "datasets"
DATA_DIR = BASE_DIR / "measurements"
RESULTS_DIR = BASE_DIR / "results"
CKPT_DIR = BASE_DIR / "ckpts"
# Set the global random seed from pytorch to ensure reproducibility of the example.
torch.manual_seed(0)
img_size = 64
device = dinv.utils.get_device()
x_true = load_example("butterfly.png", img_size=img_size).to(device)
x = x_true.clone()
n_channels = 3 # 3 for color images, 1 for gray-scale images
operation = "deblurring"
# Degradation parameters
noise_level_img = 1 / 40 # Poisson Noise gain
# Generate the gaussian blur operator with Poisson noise.
physics = dinv.physics.BlurFFT(
img_size=(n_channels, img_size, img_size),
filter=dinv.physics.functional.gaussian_blur(),
device=device,
noise_model=dinv.physics.PoissonNoise(gain=noise_level_img),
)
```
```none
Selected GPU 0 with 24371.25 MiB free memory
```
## Define the PnP algorithm.
The chosen algorithm is here MD (Mirror Descent).
```Python
# Select the data fidelity term, here Poisson likelihood due to the use of Poisson noise in the forward operator.
data_fidelity = PoissonLikelihood(gain=noise_level_img)
# Set up the denoising prior. Note that we use a Gaussian noise denoiser, even if the observation noise is Poisson.
prior = RED(denoiser=dinv.models.DnCNN(depth=20, pretrained="download").to(device))
# Set up the optimization parameters
max_iter = 200 # number of iterations
stepsize = 1.0 # stepsize of the algorithm
sigma_denoiser = 0.05 # noise level parameter of the Gaussian denoiser
verbose = True # Logging parameters
# Define the unfolded trainable model.
model = MD(
prior=prior,
data_fidelity=data_fidelity,
stepsize=stepsize,
sigma_denoiser=sigma_denoiser,
early_stop=True,
max_iter=max_iter,
verbose=verbose,
bregman_potential=BurgEntropy(),
)
```
## Evaluate the model on the problem and plot the results.
The model returns the output and the metrics computed along the iterations.
For computing PSNR, the ground truth image `x_gt` must be provided.
```Python
y = physics(x)
x_lin = physics.A_adjoint(y)
# run the model on the problem.
with torch.no_grad():
x_model, metrics = model(
y, physics, x_gt=x, compute_metrics=True
) # reconstruction with PnP algorithm
# compute PSNR
print(f"Linear reconstruction PSNR: {dinv.metric.PSNR()(x, x_lin).item():.2f} dB")
print(f"PnP reconstruction PSNR: {dinv.metric.PSNR()(x, x_model).item():.2f} dB")
# plot images. Images are saved in RESULTS_DIR.
imgs = [y, x, x_lin, x_model]
plot(
imgs,
titles=["Input", "GT", "Linear", "Recons."],
save_dir=RESULTS_DIR / "images",
show=True,
)
# plot convergence curves. Metrics are saved in RESULTS_DIR.
plot_curves(metrics, save_dir=RESULTS_DIR / "curves", show=True)
```
```none
Linear reconstruction PSNR: 20.96 dB
PnP reconstruction PSNR: 23.56 dB
```
* **References:**
* Samuel Hurault, Ulugbek Kamilov, Arthur Leclaire, and Nicolas Papadakis. Convergent bregman plug-and-play image restoration for poisson inverse problems. *Advances in Neural Information Processing Systems*, 36:27251–27280, 2023.
* Kai Zhang, Wangmeng Zuo, Yunjin Chen, Deyu Meng, and Lei Zhang. Beyond a gaussian denoiser: residual learning of deep cnn for image denoising. *IEEE transactions on image processing*, 26(7):3142–3155, 2017.
* Jérôme Bolte, Heinz Bauschke, and Marc Teboulle. A descent lemma beyond lipschitz gradient continuity: first-order methods revisited and applications. *Mathematics of Operations Research*, 42:, 07 2016. [doi:10.1287/moor.2016.0817](https://doi.org/10.1287/moor.2016.0817).
**Total running time of the script:** (0 minutes 0.738 seconds)
🏷 Tags: Denoising
[`Download Jupyter notebook: demo_PnP_mirror_descent.ipynb`](https://deepinv.org/auto_examples/plug-and-play/_downloads/5194ed48e2d444717ea9f41443aea164/demo_PnP_mirror_descent.ipynb)
[`Download Python source code: demo_PnP_mirror_descent.py`](https://deepinv.org/auto_examples/plug-and-play/_downloads/e323befa10e13808c0c4ad9c8bde473c/demo_PnP_mirror_descent.py)
[`Download zipped: demo_PnP_mirror_descent.zip`](https://deepinv.org/auto_examples/plug-and-play/_downloads/624233af9b824e7eef55b5f3cb9ec585/demo_PnP_mirror_descent.zip)
[Gallery generated by Sphinx-Gallery](https://sphinx-gallery.github.io)