This implementation based on BasicSR which is a open source toolbox for image/video restoration tasks and NAFNet
python 3.8.8
pytorch 1.9.0
cuda 11.3
git clone https://github.com/dhryougit/AFM.git
cd AFM
pip install -r requirements.txt
python setup.py develop --no_cuda_ext
We used NVIDIA RTX A6000 D6 48GB for trianing our models.
For training
python3 -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 basicsr/train.py -opt options/train/DnCNN.yml --name=DnCNN-afm-b --afm_type=AFM_B --seed=10 --afm_rate=0.8 --afm_easy_rate=0.3 --launcher pytorch
Training dataset : SIDD
Evaludation datasets : Poly, CC, HighISO, iPhone, Huawei.
Additioanl real-world noise datasets can be downloaded from "https://github.com/ZhaomingKong/Denoising-Comparison"
Dataset | Poly | CC | HighISO | iPhone | Huawei | OOD Avg. |
---|---|---|---|---|---|---|
PSRN | 37.75 | 36.84 | 39.17 | 40.65 | 38.39 | 38.56 |
SSIM | 0.9804 | 0.9830 | 0.9801 | 0.9777 | 0.9683 | 0.9779 |
Pre-trained model of our Dncnn trained on AFM-B can be downloaded from (https://drive.google.com/file/d/1uPJP2zNc4ViFc1QU7TXGAwrFIEScBJvL/view?usp=sharing)