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)

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", ""],
)

  0%|          | 0/200 [00:00<?, ?it/s]
  1%|          | 2/200 [00:00<00:12, 16.44it/s]
  2%|▏         | 4/200 [00:00<00:11, 17.21it/s]
  3%|▎         | 6/200 [00:00<00:11, 17.43it/s]
  4%|▍         | 8/200 [00:00<00:10, 17.54it/s]
  5%|▌         | 10/200 [00:00<00:10, 17.61it/s]
  6%|▌         | 12/200 [00:00<00:10, 17.64it/s]
  7%|▋         | 14/200 [00:00<00:10, 17.68it/s]
  8%|▊         | 16/200 [00:00<00:10, 17.70it/s]
  9%|▉         | 18/200 [00:01<00:10, 17.72it/s]
 10%|█         | 20/200 [00:01<00:10, 17.73it/s]
 11%|█         | 22/200 [00:01<00:10, 17.76it/s]
 12%|█▏        | 24/200 [00:01<00:09, 17.69it/s]
 13%|█▎        | 26/200 [00:01<00:09, 17.66it/s]
 14%|█▍        | 28/200 [00:01<00:09, 17.35it/s]
 15%|█▌        | 30/200 [00:01<00:09, 17.14it/s]
 16%|█▌        | 32/200 [00:01<00:09, 16.91it/s]
 17%|█▋        | 34/200 [00:01<00:09, 16.96it/s]
 18%|█▊        | 36/200 [00:02<00:09, 17.19it/s]
 19%|█▉        | 38/200 [00:02<00:09, 17.33it/s]
 20%|██        | 40/200 [00:02<00:09, 17.47it/s]
 21%|██        | 42/200 [00:02<00:08, 17.57it/s]
 22%|██▏       | 44/200 [00:02<00:08, 17.62it/s]
 23%|██▎       | 46/200 [00:02<00:08, 17.68it/s]
 24%|██▍       | 48/200 [00:02<00:08, 17.69it/s]
 25%|██▌       | 50/200 [00:02<00:08, 17.70it/s]
 26%|██▌       | 52/200 [00:02<00:08, 17.74it/s]
 27%|██▋       | 54/200 [00:03<00:08, 17.51it/s]
 28%|██▊       | 56/200 [00:03<00:08, 17.59it/s]
 29%|██▉       | 58/200 [00:03<00:08, 17.63it/s]
 30%|███       | 60/200 [00:03<00:07, 17.66it/s]
 31%|███       | 62/200 [00:03<00:07, 17.69it/s]
 32%|███▏      | 64/200 [00:03<00:07, 17.71it/s]
 33%|███▎      | 66/200 [00:03<00:07, 17.72it/s]
 34%|███▍      | 68/200 [00:03<00:07, 17.72it/s]
 35%|███▌      | 70/200 [00:03<00:07, 17.29it/s]
 36%|███▌      | 72/200 [00:04<00:07, 16.96it/s]
 37%|███▋      | 74/200 [00:04<00:07, 16.66it/s]
 38%|███▊      | 76/200 [00:04<00:07, 16.61it/s]
 39%|███▉      | 78/200 [00:04<00:07, 16.93it/s]
 40%|████      | 80/200 [00:04<00:06, 17.17it/s]
 41%|████      | 82/200 [00:04<00:06, 17.36it/s]
 42%|████▏     | 84/200 [00:04<00:06, 17.48it/s]
 43%|████▎     | 86/200 [00:04<00:06, 17.58it/s]
 44%|████▍     | 88/200 [00:05<00:06, 17.67it/s]
 45%|████▌     | 90/200 [00:05<00:06, 17.67it/s]
 46%|████▌     | 92/200 [00:05<00:06, 17.70it/s]
 47%|████▋     | 94/200 [00:05<00:05, 17.73it/s]
 48%|████▊     | 96/200 [00:05<00:05, 17.73it/s]
 49%|████▉     | 98/200 [00:05<00:05, 17.73it/s]
 50%|█████     | 100/200 [00:05<00:05, 17.74it/s]
 51%|█████     | 102/200 [00:05<00:05, 17.73it/s]
 52%|█████▏    | 104/200 [00:05<00:05, 17.73it/s]
 53%|█████▎    | 106/200 [00:06<00:05, 17.76it/s]
 54%|█████▍    | 108/200 [00:06<00:05, 17.75it/s]
 55%|█████▌    | 110/200 [00:06<00:05, 17.76it/s]
 56%|█████▌    | 112/200 [00:06<00:05, 17.41it/s]
 57%|█████▋    | 114/200 [00:06<00:05, 16.81it/s]
 58%|█████▊    | 116/200 [00:06<00:05, 16.53it/s]
 59%|█████▉    | 118/200 [00:06<00:05, 16.39it/s]
 60%|██████    | 120/200 [00:06<00:04, 16.78it/s]
 61%|██████    | 122/200 [00:07<00:04, 17.06it/s]
 62%|██████▏   | 124/200 [00:07<00:04, 17.27it/s]
 63%|██████▎   | 126/200 [00:07<00:04, 17.42it/s]
 64%|██████▍   | 128/200 [00:07<00:04, 17.52it/s]
 65%|██████▌   | 130/200 [00:07<00:03, 17.61it/s]
 66%|██████▌   | 132/200 [00:07<00:03, 17.64it/s]
 67%|██████▋   | 134/200 [00:07<00:03, 17.66it/s]
 68%|██████▊   | 136/200 [00:07<00:03, 17.70it/s]
 69%|██████▉   | 138/200 [00:07<00:03, 17.71it/s]
 70%|███████   | 140/200 [00:08<00:03, 17.71it/s]
 71%|███████   | 142/200 [00:08<00:03, 17.74it/s]
 72%|███████▏  | 144/200 [00:08<00:03, 17.75it/s]
 73%|███████▎  | 146/200 [00:08<00:03, 17.68it/s]
 74%|███████▍  | 148/200 [00:08<00:02, 17.72it/s]
 75%|███████▌  | 150/200 [00:08<00:02, 17.73it/s]
 76%|███████▌  | 152/200 [00:08<00:02, 17.74it/s]
 77%|███████▋  | 154/200 [00:08<00:02, 17.58it/s]
 78%|███████▊  | 156/200 [00:08<00:02, 17.23it/s]
 79%|███████▉  | 158/200 [00:09<00:02, 16.07it/s]
 80%|████████  | 160/200 [00:09<00:02, 15.37it/s]
 81%|████████  | 162/200 [00:09<00:02, 16.01it/s]
 82%|████████▏ | 164/200 [00:09<00:02, 16.50it/s]
 83%|████████▎ | 166/200 [00:09<00:02, 16.87it/s]
 84%|████████▍ | 168/200 [00:09<00:01, 16.97it/s]
 85%|████████▌ | 170/200 [00:09<00:01, 17.19it/s]
 86%|████████▌ | 172/200 [00:09<00:01, 17.37it/s]
 87%|████████▋ | 174/200 [00:10<00:01, 17.29it/s]
 88%|████████▊ | 176/200 [00:10<00:01, 17.27it/s]
 89%|████████▉ | 178/200 [00:10<00:01, 17.43it/s]
 90%|█████████ | 180/200 [00:10<00:01, 17.52it/s]
 91%|█████████ | 182/200 [00:10<00:01, 17.60it/s]
 92%|█████████▏| 184/200 [00:10<00:00, 17.65it/s]
 93%|█████████▎| 186/200 [00:10<00:00, 17.68it/s]
 94%|█████████▍| 188/200 [00:10<00:00, 17.69it/s]
 95%|█████████▌| 190/200 [00:10<00:00, 17.72it/s]
 96%|█████████▌| 192/200 [00:11<00:00, 17.73it/s]
 97%|█████████▋| 194/200 [00:11<00:00, 17.68it/s]
 98%|█████████▊| 196/200 [00:11<00:00, 17.50it/s]
 99%|█████████▉| 198/200 [00:11<00:00, 17.03it/s]
100%|██████████| 200/200 [00:11<00:00, 16.68it/s]
100%|██████████| 200/200 [00:11<00:00, 17.37it/s]
Measurement PSNR: 27.3 dB
Poisson2Sparse PSNR: 31.0 dB

References:

[1]

Calvin-Khang Ta, Abhishek Aich, Akash Gupta, and Amit K Roy-Chowdhury. Poisson2sparse: self-supervised poisson denoising from a single image. In International Conference on Medical Image Computing and Computer-Assisted Intervention, 557–567. Springer, 2022.

[2]

Tao Huang, Songjiang Li, Xu Jia, Huchuan Lu, and Jianzhuang Liu. Neighbor2neighbor: self-supervised denoising from single noisy images. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, 14781–14790. 2021.

[3]

Yide Zhang, Yinhao Zhu, Evan Nichols, Qingfei Wang, Siyuan Zhang, Cody Smith, and Scott Howard. A poisson-gaussian denoising dataset with real fluorescence microscopy images. In CVPR. 2019.