jjhuangcs / winnet

Official Code of "Wavelet-inspired Invertible Network for Image Denoising"

WINNet: Wavelet-inspired Invertible Network for Image Denoising

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.

1. Dependencies

• Python
• torchvision
• PyTorch>=1.0
• OpenCV for Python
• HDF5 for Python
• tensorboardX (Tensorboard for Python)

2. Usage

2.1. Training

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.

Training WINNet with known noise level

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., $\sigma_N$ = 15, 25 and 50.

Training WINNet with blind noise level [0, 55]

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 $\sigma_N$ is uniformly drawn from [0, 55].

Training Noise Estimation Network (NENet)

python NENet_train.py --lr 0.001 --outf logs/NENet

Note: For training NENet, the training noise level $\sigma_N$ is uniformly drawn from [0, 55].

2.2. Testing

Testing for denoising with known noise level

python WINNet_denoise_test.py --logdir pretrain/WINNet_lvl_1_nlvl_25 --lvl 1 --test_noiseL 25 --test_data Set12

Testing for blind denoising

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

Testing for deblurring

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

Citation

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}
}

winnet's Issues

model

Hello, I have some questions that have been showing no module named iUNet.I want to know what's going on, thank you!

train

Hi, I was wondering about the data set setup for training NEnet, does it use the same data set for training WINNET?
I also want to ask if this training does not use Gpus? What is the training time?
The other thing that I have a problem with is that when I was training, the data enhancement was stuck in the generated H5 file size of about 18.7 gigabytes so it wouldn't increase and then it wouldn't go down, and I wanted to know what was going on.
I am really looking forward to your reply.Thank you!

Question about the multi-resolution coefficients of the learned wavelet

Thanks for the inspiring work. I would like to know how to set up an alternative wavelet version with multi-resolution coefficients. I naively set stride = 2 for each splitting module (e.g., DCTSplit with stride = 2 and appropriate paddings), which leads to the right shape but the wrong value of coefficients, e.g., the first-level low-frequency approximation drops the image content.