基于Swin-Conv-UNet和数据合成的实用盲图像去噪

[ArXiv论文] [在线演示] [发表论文]

以下结果是通过我们的SCUNet模型使用纯合成训练数据获得的！ 我们在训练过程中并未使用DND和SIDD的配对噪声/干净数据!

Swin-Conv-UNet (SCUNet) 去噪网络

提出的Swin-Conv-UNet（SCUNet）去噪网络的架构。SCUNet将Swin-Conv（SC）块作为UNet主干网络的主要构建块。在每个SC块中，输入首先通过1×1卷积，随后被均分为两个特征图组，每个组分别输入Swin Transformer（SwinT）块和残差3×3卷积（RConv）块；之后，SwinT块和RConv块的输出被拼接并经过1×1卷积以生成输入的残差。“SConv”和“TConv”分别表示步长为2的2×2卷积和步长为2的2×2转置卷积。

用于真实图像去噪的新数据合成流程

提出的配对训练块合成流程的示意图。对于高质量图像，执行随机洗牌的退化序列以生成噪声图像。同时，进行缩放和逆正向色调映射以生成对应的干净图像。然后对配对的噪声/干净训练块进行裁剪以训练深度盲去噪模型。请注意，由于泊松噪声是信号依赖的，虚线箭头“Poisson”表示干净图像用于生成泊松噪声。为解决颜色偏移问题，虚线箭头“Camera Sensor”表示在干净图像上执行逆正向色调映射。

通过我们提出的训练数据合成流程生成的噪声/干净块对。高质量图像块的大小为544×544。噪声/干净块的大小为128×128。

Web演示

在此处尝试SCUNet模型的Replicate网页演示

代码

1. 下载SCUNet模型

python main_download_pretrained_models.py --models "SCUNet" --model_dir "model_zoo"

2. 高斯去噪

   1. 灰度图像

   python main_test_scunet_gray_gaussian.py --model_name scunet_gray_25 --noise_level_img 25 --testset_name set12
   

   2. 色彩图像

   python main_test_scunet_color_gaussian.py --model_name scunet_color_25 --noise_level_img 25 --testset_name bsd68
   

3. 盲真实图像去噪

   python main_test_scunet_real_application.py --model_name scunet_color_real_psnr --testset_name real3
   

   python main_test_scunet_real_application.py --model_name scunet_color_real_gan --testset_name real3
   

高斯去噪结果

真实图像去噪结果

@article{zhang2023practical,
   author = {Zhang, Kai and Li, Yawei and Liang, Jingyun and Cao, Jiezhang and Zhang, Yulun and Tang, Hao and Fan, Deng-Ping and Timofte, Radu and Gool, Luc Van},
   title = {Practical Blind Image Denoising via Swin-Conv-UNet and Data Synthesis},
   journal = {Machine Intelligence Research},
   DOI = {10.1007/s11633-023-1466-0},
   url = {https://doi.org/10.1007/s11633-023-1466-0},
   volume={20},
   number={6},
   pages={822--836},
   year={2023},
   publisher={Springer}
}

SCUNet 快速上手指南

环境准备

• 系统要求：Python 3.8+，CUDA 11.3+（可选）
• 前置依赖：

  pip install torch torchvision opencv-python numpy
  # 推荐使用国内镜像源
  pip install -i https://pypi.tuna.tsinghua.edu.cn/simple torch torchvision opencv-python numpy
  

安装步骤

1. 下载预训练模型

   python main_download_pretrained_models.py --models "SCUNet" --model_dir "model_zoo"
   

2. 安装依赖库（如未安装）

   pip install -i https://pypi.tuna.tsinghua.edu.cn/simple torch torchvision opencv-python numpy
   

基本使用

基础去噪示例

# 灰度图像去噪
python main_test_scunet_gray_gaussian.py --model_name scunet_gray_25 --noise_level_img 25 --testset_name set12

# 彩色图像去噪
python main_test_scunet_color_gaussian.py --model_name scunet_color_25 --noise_level_img 25 --testset_name bsd68

实际图像去噪

python main_test_scunet_real_application.py --model_name scunet_color_real_psnr --testset_name real3
python main_test_scunet_real_application.py --model_name scunet_color_real_gan --testset_name real3

说明：所有命令均需在项目根目录执行，模型文件存储于model_zoo目录。