Noise2Noise: Learning Image Restoration without Clean Data
This is an unofficial PyTorch implementation of Noise2Noise (Lehtinen et al. 2018).
- PyTorch (0.4.1)
- Torchvision (0.2.0)
- NumPy (1.14.2)
- Matplotlib (2.2.3)
- Pillow (5.2.0)
- OpenEXR (1.3.0)
To install the latest version of all packages, run
pip3 install --user -r requirements.txt
This code was tested on Python 3.6.5 on macOS High Sierra (10.13.4) and Ubuntu 16.04. It will fail with Python 2.7.x due to usage of 3.6-specific functions. Note that training and testing will also fail on Windows out of the box due to differences in the path resolvers (
The authors use ImageNet, but any dataset will do. COCO 2017 has a small validation set (1 GB) which can be nicely split into train/valid for easier training. For instance, to obtain a 4200/800 train/valid split you can do:
mkdir data && cd data mkdir train valid test wget http://images.cocodataset.org/zips/val2017.zip unzip val2017.zip && cd val2017 mv `ls | head -4200` ../train mv `ls | head -800` ../valid
You can also download the full datasets (7 GB) that more or less match the paper, if you have the bandwidth:
mkdir data && cd data mkdir train valid test wget http://images.cocodataset.org/zips/test2017.zip wget http://images.cocodataset.org/zips/val2017.zip unzip -j test2017.zip -d train unzip -j val2017.zip -d valid
Add your favorite images to the
data/test folder. Only a handful will do to visually inspect the denoiser performance.
python3 train.py --h for list of optional arguments, or
examples/train.sh for an example.
By default, the model train with noisy targets. To train with clean targets, use
--clean-targets. To train and validate on smaller datasets, use the
--valid-size options. To plot stats as the model trains, use
--plot-stats; these are saved alongside checkpoints. By default CUDA is not enabled: use the
--cuda option if you have a GPU that supports it.
The noise parameter is the maximum standard deviation σ.
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
The noise parameter is the Poisson parameter λ.
python3 train.py --loss l2 \ --noise-type poisson \ --noise-param 50 \ --cuda
The noise parameter is the approximate probability p that a pixel is covered by text.
python3 train.py \ --loss l1 \ --noise-type text \ --noise-param 0.5 \ --cuda
Monte Carlo rendering noise
See other read-me file.
Model checkpoints are automatically saved after every epoch. To test the denoiser, provide
test.py with a PyTorch model (
.pt file) via the argument
--load-ckpt and a test image directory via
--show-output option specifies the number of noisy/denoised/clean montages to display on screen. To disable this, simply remove
python3 test.py \ --data ../data \ --load-ckpt ../ckpts/gaussian/n2n.pt \ --noise-type gaussian \ --noise-param 50 \ --crop-size 256 \ --show-output 3 \ --cuda
python3 test.py --h for list of optional arguments, or
examples/test.sh for an example.
Gaussian model was trained for 100 epochs with a train/valid split of 2000/400. Text model was trained with 1000/200 due to the overhead introduced when corrupting with text. Both models were trained on an old NVIDIA GTX 780.
|Gaussian noise (σ = 25)|
|Noisy input (20.34 dB)||Denoised (32.68 dB)||Clean targets (32.49 dB)||Ground truth|
|Text overlay (p = 0.25)|
|Noisy input (15.07 dB)||Denoised (28.10 dB)||Clean targets (27.79 dB)||Ground truth|
- Activation functions should be LeakyReLUs everywhere except last layer where it should be ReLU. The current implementation (and pretrained models) assumes the opposite due to me originally misreading the paper. It was reported that using the correct version yields unstable training without batch norm; I haven't tested it yet.
- It is unclear how to deal with Poisson noise since it is data-dependent and thus nonadditive. See official TensorFlow implementation to adapt properly.
Jaakko Lehtinen, Jacob Munkberg, Jon Hasselgren, Samuli Laine, Tero Karras, Miika Aittala,and Timo Aila. Noise2Noise: Learning Image Restoration without Clean Data. Proceedings of the 35th International Conference on Machine Learning, 2018.
Tsung-Yi Lin, Michael Maire, Serge Belongie, Lubomir Bourdev, Ross Girshick, James Hays, Pietro Perona, Deva Ramanan, C. Lawrence Zitnick, and Piotr Dollár. Microsoft COCO: Common Objects in Context. arXiv:1405.0312, 2014.
I would like to acknowledge Yusuke Uchida for his Keras implementation of Noise2Noise. Although Keras and PyTorch are very different frameworks, parts of his code did help me in completing this implementation.