Attention Mechanism Enhanced Kernel Prediction Networks (AME-KPNs)

Offical in Attention-Mechanism-Enhanced-KPN

The unofficial implementation of AME-KPNs in PyTorch, and paper is accepted by ICASSP 2020 (oral), it is available at http://arxiv.org/abs/1910.08313.

Data

Use SIDD dataset to train.

Have two folder : noisy image and ground true image

Input folder have struct :

/
    /noise
        /[scene_instance]
            /[image].PNG
    /gt
        /[scene_instance]
            /[image].PNG

Train

This repo. supports training on multiple GPUs.

Train

CUDA_VISIBLE_DEVICES=0,1 python train_eval_syn.py --noise_dir ../image/noise/ --gt_dir ../image/gt/ --image_size 512 --batch_size 1 --save_every 100 --loss_every 10 -nw 4 -c -m -ckpt att_kpn --model_type attKPN --restart

If no `--restart`, the train process would be resumed.
Train Deep Guide Filter

CUDA_VISIBLE_DEVICES=0,1 python train_eval_syn_DGF.py —noise_dir ../image/noise/ —gt_dir ../image/gt/ —image_size 512 —batch_size 1 —burst_length 16 —save_every 100 —loss_every 10 -nw 4 -c -m -ckpt att_kpn —model_type attKPN —restart```

## Eval
Eval

CUDA_VISIBLE_DEVICES=0,1 python test.py —noise_dir ../image/noise/ —gt_dir ../image/gt/ —image_size 512 -nw 4 -c -m -ckpt att_kpn —model_type attKPN```

Eval with custome data : data have two folder image : *noise* and *gt*. 
Image will save in folder *-s* after predicted.

CUDA_VISIBLE_DEVICES=1 python test_custom_DGF.py -n ../FullTest/noisy/ -g ../FullTest/clean/ -b 4 -c -ckpt att_kpn_dgf_4_new -m attKPN -s img/att_kpn_dgf_4_new

### News
- Support KPN (Kernel Prediction Networks), MKPN (Multi-Kernel Prediction Networks)
- The current version supports training on color images.
- Add Deep Guide Filter
- Add noise estimate model to end-to-end denoising model
- Add KPN_Wave : replace **polling** layer by **Wavelet pooling**, **Upsampling** by **inverse wavelet pooling**. (*Att_KPN_Wavelet_DGF*)
- Add NonKPN model
- Add synthetic data loader
## Name 
*_custom  : load image from unstruct folder, print or save image for report
*_split : load one image and split image into burst image. 
*_DGF : model with Deep Guide Filter
*_noise : model with noise estimate 
### Requirements

pip install -r requirments.txt

### Citation

https://github.com/z-bingo/Attention-Mechanism-Enhanced-KPN

@article{zhang2019attention,
title={Attention Mechanism Enhanced Kernel Prediction Networks for Denoising of Burst Images},
author={Bin Zhang and Shenyao Jin and Yili Xia and Yongming Huang and Zixiang Xiong},
year={2019},
journal={arXiv preprint arXiv:1910.08313}
}
```