Poisson denoising using Poisson2Sparse#

This code shows how to restore a single image corrupted by Poisson noise using Poisson2Sparse, without requiring external training or knowledge of the noise level.

This method is based on the paper β€œPoisson2Sparse” Ta et al.[1] and restores an image by learning a sparse non-linear dictionary parametrized by a neural network using a combination of Neighbor2Neighbor Huang et al.[2], of the negative log Poisson likelihood, of the \(\ell^1\) pixel distance and of a sparsity-inducing \(\ell^1\) regularization function on the weights.

import deepinv as dinv
import torch

Load a Poisson corrupted image#

This example uses an image from the microscopy dataset FMD Zhang et al.[3].

# Seed the RNGs for reproducibility
torch.manual_seed(0)
torch.cuda.manual_seed(0)

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

physics = dinv.physics.Denoising(dinv.physics.PoissonNoise(gain=0.01, normalize=True))

x = dinv.utils.demo.load_example(
    "FMD_TwoPhoton_MICE_R_gt_12_avg50.png", img_size=(256, 256)
).to(device)
x = x[:, 0:1, :64, :64]
x = x.clamp(0, 1)
y = physics(x)
Selected GPU 0 with 3882.25 MiB free memory

Define the Poisson2Sparse model

backbone = dinv.models.ConvLista(
    in_channels=1,
    out_channels=1,
    kernel_size=3,
    num_filters=512,
    num_iter=10,
    stride=2,
    threshold=0.01,
)

model = dinv.models.Poisson2Sparse(
    backbone=backbone,
    lr=1e-4,
    num_iter=200,
    weight_n2n=2.0,
    weight_l1_regularization=1e-5,
    verbose=True,
).to(device)

Run the model#

Note that we do not pass in the physics model as Poisson2Sparse assumes a Poisson noise model internally and does not depend on the noise level.

x_hat = model(y)

# Compute and display PSNR values
learning_free_psnr = dinv.metric.PSNR()(y, x).item()
model_psnr = dinv.metric.PSNR()(x_hat, x).item()
print(f"Measurement PSNR: {learning_free_psnr:.1f} dB")
print(f"Poisson2Sparse PSNR: {model_psnr:.1f} dB")

# Plot results
dinv.utils.plot(
    [y, x_hat, x],
    titles=["Measurement", "Poisson2Sparse", "Ground truth"],
    subtitles=[f"{learning_free_psnr:.1f} dB", f"{model_psnr:.1f} dB", ""],
)
Measurement, Poisson2Sparse, Ground truth
  0%|          | 0/200 [00:00<?, ?it/s]
  0%|          | 1/200 [00:00<00:22,  8.96it/s]
  2%|▏         | 4/200 [00:00<00:10, 19.00it/s]
  4%|β–Ž         | 7/200 [00:00<00:08, 22.62it/s]
  5%|β–Œ         | 10/200 [00:00<00:07, 24.35it/s]
  6%|β–‹         | 13/200 [00:00<00:07, 25.28it/s]
  8%|β–Š         | 16/200 [00:00<00:07, 25.89it/s]
 10%|β–‰         | 19/200 [00:00<00:06, 26.29it/s]
 11%|β–ˆ         | 22/200 [00:00<00:06, 26.54it/s]
 12%|β–ˆβ–Ž        | 25/200 [00:01<00:06, 26.69it/s]
 14%|β–ˆβ–        | 28/200 [00:01<00:06, 26.74it/s]
 16%|β–ˆβ–Œ        | 31/200 [00:01<00:06, 26.79it/s]
 17%|β–ˆβ–‹        | 34/200 [00:01<00:06, 26.80it/s]
 18%|β–ˆβ–Š        | 37/200 [00:01<00:06, 26.87it/s]
 20%|β–ˆβ–ˆ        | 40/200 [00:01<00:05, 26.90it/s]
 22%|β–ˆβ–ˆβ–       | 43/200 [00:01<00:05, 26.90it/s]
 23%|β–ˆβ–ˆβ–Ž       | 46/200 [00:01<00:05, 26.93it/s]
 24%|β–ˆβ–ˆβ–       | 49/200 [00:01<00:05, 26.96it/s]
 26%|β–ˆβ–ˆβ–Œ       | 52/200 [00:02<00:05, 26.99it/s]
 28%|β–ˆβ–ˆβ–Š       | 55/200 [00:02<00:05, 27.02it/s]
 29%|β–ˆβ–ˆβ–‰       | 58/200 [00:02<00:05, 26.69it/s]
 30%|β–ˆβ–ˆβ–ˆ       | 61/200 [00:02<00:05, 26.79it/s]
 32%|β–ˆβ–ˆβ–ˆβ–      | 64/200 [00:02<00:05, 26.88it/s]
 34%|β–ˆβ–ˆβ–ˆβ–Ž      | 67/200 [00:02<00:04, 26.94it/s]
 35%|β–ˆβ–ˆβ–ˆβ–Œ      | 70/200 [00:02<00:04, 26.96it/s]
 36%|β–ˆβ–ˆβ–ˆβ–‹      | 73/200 [00:02<00:04, 27.02it/s]
 38%|β–ˆβ–ˆβ–ˆβ–Š      | 76/200 [00:02<00:04, 26.99it/s]
 40%|β–ˆβ–ˆβ–ˆβ–‰      | 79/200 [00:03<00:04, 27.01it/s]
 41%|β–ˆβ–ˆβ–ˆβ–ˆ      | 82/200 [00:03<00:04, 27.04it/s]
 42%|β–ˆβ–ˆβ–ˆβ–ˆβ–Ž     | 85/200 [00:03<00:04, 27.09it/s]
 44%|β–ˆβ–ˆβ–ˆβ–ˆβ–     | 88/200 [00:03<00:04, 27.11it/s]
 46%|β–ˆβ–ˆβ–ˆβ–ˆβ–Œ     | 91/200 [00:03<00:04, 27.14it/s]
 47%|β–ˆβ–ˆβ–ˆβ–ˆβ–‹     | 94/200 [00:03<00:03, 27.15it/s]
 48%|β–ˆβ–ˆβ–ˆβ–ˆβ–Š     | 97/200 [00:03<00:03, 27.15it/s]
 50%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆ     | 100/200 [00:03<00:03, 27.16it/s]
 52%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–    | 103/200 [00:03<00:03, 27.14it/s]
 53%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž    | 106/200 [00:04<00:03, 27.13it/s]
 55%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–    | 109/200 [00:04<00:03, 27.12it/s]
 56%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ    | 112/200 [00:04<00:03, 27.11it/s]
 57%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š    | 115/200 [00:04<00:03, 27.08it/s]
 59%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰    | 118/200 [00:04<00:03, 27.06it/s]
 60%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ    | 121/200 [00:04<00:02, 27.02it/s]
 62%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–   | 124/200 [00:04<00:02, 26.98it/s]
 64%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž   | 127/200 [00:04<00:02, 26.94it/s]
 65%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ   | 130/200 [00:04<00:02, 26.94it/s]
 66%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹   | 133/200 [00:05<00:02, 26.93it/s]
 68%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š   | 136/200 [00:05<00:02, 26.91it/s]
 70%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰   | 139/200 [00:05<00:02, 26.91it/s]
 71%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ   | 142/200 [00:05<00:02, 26.91it/s]
 72%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž  | 145/200 [00:05<00:02, 26.91it/s]
 74%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–  | 148/200 [00:05<00:01, 26.90it/s]
 76%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ  | 151/200 [00:05<00:01, 26.90it/s]
 77%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹  | 154/200 [00:05<00:01, 26.90it/s]
 78%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š  | 157/200 [00:05<00:01, 26.90it/s]
 80%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ  | 160/200 [00:06<00:01, 26.91it/s]
 82%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ– | 163/200 [00:06<00:01, 26.96it/s]
 83%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž | 166/200 [00:06<00:01, 26.99it/s]
 84%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ– | 169/200 [00:06<00:01, 27.01it/s]
 86%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ | 172/200 [00:06<00:01, 27.02it/s]
 88%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š | 175/200 [00:06<00:00, 27.01it/s]
 89%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰ | 178/200 [00:06<00:00, 26.93it/s]
 90%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ | 181/200 [00:06<00:00, 26.92it/s]
 92%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–| 184/200 [00:06<00:00, 26.85it/s]
 94%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž| 187/200 [00:07<00:00, 26.89it/s]
 95%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ| 190/200 [00:07<00:00, 26.96it/s]
 96%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹| 193/200 [00:07<00:00, 26.96it/s]
 98%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š| 196/200 [00:07<00:00, 27.02it/s]
100%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰| 199/200 [00:07<00:00, 27.04it/s]
100%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 200/200 [00:07<00:00, 26.70it/s]
Measurement PSNR: 27.5 dB
Poisson2Sparse PSNR: 30.1 dB
References:

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

Gallery generated by Sphinx-Gallery