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

  0%|          | 0/200 [00:00<?, ?it/s]
  2%|▏         | 3/200 [00:00<00:07, 26.81it/s]
  4%|▎         | 7/200 [00:00<00:06, 31.49it/s]
  6%|▌         | 11/200 [00:00<00:05, 33.95it/s]
  8%|▊         | 15/200 [00:00<00:05, 35.17it/s]
 10%|▉         | 19/200 [00:00<00:05, 35.85it/s]
 12%|█▏        | 23/200 [00:00<00:04, 36.27it/s]
 14%|█▎        | 27/200 [00:00<00:04, 36.54it/s]
 16%|█▌        | 31/200 [00:00<00:04, 36.72it/s]
 18%|█▊        | 35/200 [00:00<00:04, 36.82it/s]
 20%|█▉        | 39/200 [00:01<00:04, 36.90it/s]
 22%|██▏       | 43/200 [00:01<00:04, 36.95it/s]
 24%|██▎       | 47/200 [00:01<00:04, 36.87it/s]
 26%|██▌       | 51/200 [00:01<00:04, 36.92it/s]
 28%|██▊       | 55/200 [00:01<00:03, 36.94it/s]
 30%|██▉       | 59/200 [00:01<00:03, 36.96it/s]
 32%|███▏      | 63/200 [00:01<00:03, 36.99it/s]
 34%|███▎      | 67/200 [00:01<00:03, 37.01it/s]
 36%|███▌      | 71/200 [00:01<00:03, 37.03it/s]
 38%|███▊      | 75/200 [00:02<00:03, 37.04it/s]
 40%|███▉      | 79/200 [00:02<00:03, 37.06it/s]
 42%|████▏     | 83/200 [00:02<00:03, 37.06it/s]
 44%|████▎     | 87/200 [00:02<00:03, 37.06it/s]
 46%|████▌     | 91/200 [00:02<00:02, 37.09it/s]
 48%|████▊     | 95/200 [00:02<00:02, 37.09it/s]
 50%|████▉     | 99/200 [00:02<00:02, 37.09it/s]
 52%|█████▏    | 103/200 [00:02<00:02, 37.08it/s]
 54%|█████▎    | 107/200 [00:02<00:02, 37.07it/s]
 56%|█████▌    | 111/200 [00:03<00:02, 37.09it/s]
 57%|█████▊    | 115/200 [00:03<00:02, 37.07it/s]
 60%|█████▉    | 119/200 [00:03<00:02, 37.08it/s]
 62%|██████▏   | 123/200 [00:03<00:02, 37.08it/s]
 64%|██████▎   | 127/200 [00:03<00:01, 37.10it/s]
 66%|██████▌   | 131/200 [00:03<00:01, 37.09it/s]
 68%|██████▊   | 135/200 [00:03<00:01, 37.09it/s]
 70%|██████▉   | 139/200 [00:03<00:01, 36.55it/s]
 72%|███████▏  | 143/200 [00:03<00:01, 36.45it/s]
 74%|███████▎  | 147/200 [00:04<00:01, 36.61it/s]
 76%|███████▌  | 151/200 [00:04<00:01, 36.69it/s]
 78%|███████▊  | 155/200 [00:04<00:01, 36.34it/s]
 80%|███████▉  | 159/200 [00:04<00:01, 36.53it/s]
 82%|████████▏ | 163/200 [00:04<00:01, 36.69it/s]
 84%|████████▎ | 167/200 [00:04<00:00, 36.80it/s]
 86%|████████▌ | 171/200 [00:04<00:00, 36.88it/s]
 88%|████████▊ | 175/200 [00:04<00:00, 36.95it/s]
 90%|████████▉ | 179/200 [00:04<00:00, 36.99it/s]
 92%|█████████▏| 183/200 [00:04<00:00, 37.00it/s]
 94%|█████████▎| 187/200 [00:05<00:00, 36.97it/s]
 96%|█████████▌| 191/200 [00:05<00:00, 37.00it/s]
 98%|█████████▊| 195/200 [00:05<00:00, 37.01it/s]
100%|█████████▉| 199/200 [00:05<00:00, 37.04it/s]
100%|██████████| 200/200 [00:05<00:00, 36.71it/s]
Measurement PSNR: 27.5 dB
Poisson2Sparse PSNR: 30.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.