Skip to content

The paddlepaddle implemention of NAFNet for denoising

License

Notifications You must be signed in to change notification settings

kongdebug/NAFNet-pd

Repository files navigation

NAFNet

The paddlepaddle implemention of NAFNet

Simple Baselines for Image Restoration 论文复现

官方源码

AI Studio 脚本项目地址

1. 简介

NAFNet

NAFNet的网络设计和特点如上图所示,采用带跳过连接的UNet作为整体架构,同时修改了Restormer块中的Transformer模块,并取消了激活函数,采取更简单有效的simplegate设计,运用更简单的通道注意力机制

2. 复现精度

验收标准:SIDD PSNR: 40.30 SSIM: 0.961

复现结果:SIDD PSNR: 40.20 SSIM: 0.959

3. 数据集、复现模型权重、文件结构

数据集

下载数据:

利用官方代码,得到SIDD数据集并解压到SIDD_Data下

  • mkdir ./SIDD_Data

  • 下载训练集 train, 将 (./SIDD_Medium_Srgb/Data) 移动到 ./SIDD_Data/

  • 下载测试集 val (ValidationNoisyBlocksSrgb.mat and ValidationGtBlocksSrgb.mat) ,移动到 ./SIDD_Data/

  • 结构如下:

    ./SIDD_Data/Data
    ./SIDD_Data/ValidationNoisyBlocksSrgb.mat
    ./SIDD_Data/ValidationGtBlocksSrgb.mat
  • python sidd_data_preprocessing.py

    • 将训练图片剪裁为 512x512 patches.
    • 处理后的数据集可于AI Studio找到

:为简化操作,可直接使用处理后上传的SIDD数据;NAFNetHINET均是旷视科技的模型,对SIDD数据均以步长384切割成512×512大小的图像块,故此处数据处理参考hinet_paddle进行处理

复现模型权重

AI Studio:下载链接,数据集中的model_best.pdparams即为复现得到的最优权重

文件结构

NAFNet-pd
    |-- dataloaders
    |-- SIDD_Data
         |-- train                 # SIDD-Medium 训练数据
         |-- val                   # SIDD 测试数据
    |-- SIDD_patches
         |-- train_mini            # 小训练数据,用于TIPC测试
         |-- val_mini              # 小测试数据,用于TIPC测试
    |-- logs                       # 训练日志
    |-- test_tipc                  # TIPC: Linux GPU/CPU 基础训练推理测试
    |-- networks
         |-- NAFNet_arch.py        # NAFNet模型代码
    |-- utils                      # 一些工具代码
    |-- config.py                  # 配置文件
    |-- export_model.py            # 预训练模型的导出代码
    |-- infer.py                   # 模型推理代码
    |-- LICENSE                    # LICENSE文件
    |-- losses.py                  # 损失函数
    |-- predict.py                 # 模型预测代码
    |-- README.md                  # README.md文件
    |-- sidd_data_preprocessing.py # SIDD数据预处理代码
    |-- test_denoising_sidd.py     # 测试SIDD数据上的指标
    |-- train.py                   # TIPC训练测试代码
    |-- train_denoising_1card.py   # 单机单卡训练代码
    |-- train_denoising_4cards.py  # 单机多卡训练代码
    |-- training_1card.yml         # 单机单卡训练配置文件
    |-- training_4cards.py         # 单机多卡训练配置文件

4. 环境依赖

PaddlePaddle == 2.2.2, 若用paddle2.3.2则paddle.cumsum()函数有问题,会带来错误的推理结果

5. 快速开始

GPU数量改变时,须保证

total_batchsize*iter == 8gpus*8bs*400000iters

与官方保持一致

单机单卡

python train_denoising_1card.py

配置文件为 training_1card.yml

单机四卡

python -m paddle.distributed.launch train_denoising_4cards.py

此处为用四张卡,配置文件为 training_4cards.yml

日志说明

由于训练模型采用的是脚本任务训练,本身脚本任务就有相应的日志记录,均保存在了./logs文件夹下

模型评估

在 SIDD 测试数据上作测试

python test_denoising_sidd.py --weight ../data/data168981/model_best.pdparams 

输出如下:

PSNR: 40.2024
SSIM: 0.9590

模型预测

在 SIDD 小验证集上作预测,结果存放在 results/ 文件夹下

python predict.py --model_ckpt ../data/data168981/model_best.pdparams --data_path ./SIDD_patches/val_mini/ --save_path results/ --save_images

输出结果如下:

PSNR on test data 40.7486, SSIM on test data 0.9809

预测的部分图片如图所示:

噪声图像 预测输出

推理过程:

需要安装 reprod_log:

pip install reprod_log

模型动转静导出:

python export_model.py --model-dir ../data/data168981/model_best.pdparams --save-inference-dir ./inference_output

最终在./inference_output/文件夹下会生成下面的3个文件:

inference_output
  |----model.pdiparams     : 模型参数文件
  |----model.pdmodel       : 模型结构文件
  |----model.pdiparams.info: 模型参数信息文件

模型推理:

python infer.py --model-dir inference_output --use-gpu True --benchmark False --clean-dir=./SIDD_patches/val_mini/gt_crops/ValidationBlocksSrgb_0.png --noisy-dir=./SIDD_patches/val_mini/input_crops/ValidationBlocksSrgb_0.png

输出结果如下:

image_name: ./SIDD_patches/val_mini/input_crops/ValidationBlocksSrgb_0.png, psnr: 41.94810134125529

6. TIPC

首先安装AutoLog(规范化日志输出工具)

pip install  https://paddleocr.bj.bcebos.com/libs/auto_log-1.2.0-py3-none-any.whl

在linux下,进入 NAFNet-pd 文件夹,运行命令:

  • 获取小批量数据
bash test_tipc/prepare.sh ./test_tipc/configs/NAFNet/train_infer_python.txt 'lite_train_lite_infer'
  • 训练、测试、导出推理一体化
bash test_tipc/test_train_inference_python.sh ./test_tipc/configs/NAFNet/train_infer_python.txt 'lite_train_lite_infer'

7. 致谢

感谢NAFNet-officialMIRNet_paddle以及hinet_paddle分享了他们的代码,在本次复现过程中提供了帮助,以及AI Studio提供的算力与答疑支持。

About

The paddlepaddle implemention of NAFNet for denoising

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published