Jun-Jie Huang ([email protected]) and Pier Luigi Dragotti
Pytorch implementation for "WINNet: Wavelet-inspired Invertible Network for Image Denoising" (TIP'2022).
Overview of the proposed wavelet-inspired invertible network (WINNet). It consists of K levels of lifting inspired invertible neural networks (LINN) and denoising network. The forward transform of LINN non-linearly converts the input noisy image into coarse part (green) and detail parts (black). Denoising network will perform denoising operation on the detail part while the coarse version is decomposed again using a second level and the decomposition and denoising steps are repeated K times. The backward transform of the LINN will reconstruct the denoised image using the denoised detail parts and the original coarse part. The estimated noise level from the noise estimation network will be used to adjust the soft-thresholds of the soft-thresholding non-linearity to make the WINNet to adapt well to the current noise level.
- Python
- torchvision
- PyTorch>=1.0
- OpenCV for Python
- HDF5 for Python
- tensorboardX (Tensorboard for Python)
The 400 training images from the Berkeley segmentation dataset (BSD) of size 180 × 180 are used for training. The training patch size for non-blind and blind denoising setting is 40×40 and 50×50, respectively.
python WINNet_denoise_train.py --mode S --lr 0.001 --outf logs/WINNet_lvl_1_nlvl_25 --lvl 1 --noiseL 25
Note: For training with a single noise level, three noise levels are considered, i.e.,
python WINNet_denoise_train.py --mode B --lr 0.001 --outf logs/WINNet_lvl_1_nlvl_0_55 --lvl 1
Note: For blind image denoising scenario, the training noise level
python NENet_train.py --lr 0.001 --outf logs/NENet
Note: For training NENet, the training noise level
python WINNet_denoise_test.py --logdir pretrain/WINNet_lvl_1_nlvl_25 --lvl 1 --test_noiseL 25 --test_data Set12
python WINNet_wNE_blind_denoise_test.py --logdirdn pretrain/WINNet_lvl_1_nlvl_0_55 --logdirne pretrain/NENet --lvl 1 --test_data Set12 --test_noiseL 85
python WINNet_deblur_test.py --logdirdn pretrain/WINNet_lvl_1_nlvl_0_55 --logdirne pretrain/NENet --lvl 1 --test_data Set12 --test_noiseL 2.55
If you use any part of this code in your research, please cite our paper:
@article{huang2022WINNet,
author={Huang, Jun-Jie and Dragotti, Pier Luigi},
journal={IEEE Transactions on Image Processing},
title={{WINNet}: Wavelet-Inspired Invertible Network for Image Denoising},
year={2022},
volume={31},
number={},
pages={4377-4392},
doi={10.1109/TIP.2022.3184845}
}
@inproceedings{huang2021linn,
author={Huang, Jun-Jie and Dragotti, Pier Luigi},
title={{LINN}: Lifting Inspired Invertible Neural Networks for Image Denoising},
booktitle={2021 29th European Signal Processing Conference (EUSIPCO)},
year={2021},
pages={636-640},
organization={EURASIP},
doi={10.23919/EUSIPCO54536.2021.9615931}
}