joeylitalien/noise2noise-pytorch

noise2noise‑pytorch – Denoising without clean data

What it is

• An unofficial PyTorch re‑implementation of the Noise2Noise paper (Lehtinen et al., 2018). The method trains an image‑to‑image network to map noisy inputs to noisy targets, learning to restore images even when no clean reference is available.

Key components

• A U‑Net‑style convolutional network (the code contains a small bug where encoder and decoder share weights, but it still works).
• Training script (train.py) supporting several synthetic noise models:
  • Gaussian (σ controlled via --noise-param)
  • Poisson (λ via --noise-param)
  • Text overlay (probability p via --noise-param)
  • Monte‑Carlo rendering noise (see MonteCarlo.md).
• Testing script (test.py) that loads a checkpoint and visualises noisy / denoised / clean images.
• Optional plotting of loss/PSNR curves during training (--plot-stats).

How to get started

1. Install the required Python packages (PyTorch 0.4.1, torchvision 0.2.0, NumPy, Matplotlib, Pillow, OpenEXR) via pip3 install -r requirements.txt.
2. Prepare data – any image collection works. The README shows how to split the COCO 2017 validation set into train, valid, and test folders, or you can point the scripts at your own directories.
3. Train a model, e.g. for Gaussian noise:

   python3 train.py \
     --train-dir data/train --train-size 1000 \
     --valid-dir data/valid --valid-size 200 \
     --ckpt-save-path ckpts \
     --nb-epochs 10 \
     --batch-size 4 \
     --loss l2 \
     --noise-type gaussian \
     --noise-param 50 \
     --crop-size 64 \
     --plot-stats \
     --cuda
   

   Replace --noise-type and --noise-param for other corruptions.
4. Test a saved checkpoint:

   python3 test.py \
     --data data/test \
     --load-ckpt ckpts/gaussian/n2n.pt \
     --noise-type gaussian \
     --noise-param 50 \
     --crop-size 256 \
     --show-output 3 \
     --cuda
   

   The script will display side‑by‑side montages of noisy input, denoised output, and (if available) clean targets.

Results shown in the repo

• Gaussian noise (σ = 25) denoised images reach ~32 dB PSNR, close to the clean‑target baseline.
• Text overlay (p = 0.25) denoised images reach ~28 dB PSNR.
• Training was performed on an old GTX 780 for a few hundred epochs.

Known limitations / open issues

• The U‑Net implementation mistakenly shares encoder weights with the decoder; fixing it would require a small code change.
• Activation functions are reversed (LeakyReLU used where the paper suggests ReLU and vice‑versa); the current setup works but may not be optimal.
• Poisson noise handling is not fully aligned with the original paper because Poisson noise is data‑dependent; users may need to adapt the TensorFlow reference implementation for a proper treatment.

Who might use this

• Researchers or hobbyists experimenting with self‑supervised image restoration.
• Developers who need a lightweight PyTorch baseline for denoising tasks where clean ground‑truth data are unavailable.

References

• Lehtinen, J. et al. “Noise2Noise: Learning Image Restoration without Clean Data”, ICML 2018.
• COCO dataset citation for the example data.

---

The repository is a genuine, research‑oriented code base for a machine‑learning image‑restoration technique.