Restormer：用于高分辨率图像恢复的高效Transformer（CVPR 2022 — 口头报告）

Syed Waqas Zamir、Aditya Arora、Salman Khan、Munawar Hayat、Fahad Shahbaz Khan以及Ming-Hsuan Yang

最新消息

• 2022年4月4日： 已集成至Huggingface Spaces 🤗，采用Gradio实现。体验在线演示：
• 2022年3月30日： 新增Colab演示。
• 2022年3月29日： Restormer被选为CVPR 2022的口头报告 :dizzy:
• 2022年3月10日： 训练代码已发布 :fire:
• 2022年3月3日： 论文已被CVPR 2022接收 :tada:
• 2021年11月21日： 测试代码和预训练模型已发布！

---

摘要： 由于卷积神经网络（CNN）能够从大规模数据中学习通用的图像先验，因此在图像恢复及相关任务中得到了广泛应用。近期，另一类神经网络架构——Transformer，在自然语言处理和高层视觉任务中表现出显著的性能提升。尽管Transformer模型克服了CNN的局限性（如感受野有限、对输入内容适应性不足），但其计算复杂度随空间分辨率呈二次增长，这使得它难以应用于大多数涉及高分辨率图像的恢复任务。在本工作中，我们通过在基础模块（多头注意力机制和前馈网络）中进行若干关键设计，提出了一种高效的Transformer模型，使其能够在捕捉长距离像素交互的同时，仍适用于大尺寸图像。我们的模型名为“恢复Transformer”（Restormer），在多项图像恢复任务上取得了当前最优的性能，包括去雨、单张图像运动模糊去除、散焦模糊去除（单张图像及双像素数据）以及图像去噪（高斯灰度/彩色噪声去除，以及真实场景图像去噪）。

---

网络架构

安装

请参阅INSTALL.md，了解运行Restormer所需的依赖项安装说明。

演示

要使用您自己的图片测试去雨、运动模糊去除、散焦模糊去除和去噪等任务的预训练Restormer模型，您可以选择Google Colab ，或按以下命令行方式操作：

python demo.py --task 任务名称 --input_dir 图片路径 --result_dir 保存结果路径

例如，对一个文件夹内的图片进行散焦模糊去除：

python demo.py --task 单张图像散焦模糊去除 --input_dir './demo/degraded/' --result_dir './demo/restored/'

又如，直接对一张图片进行散焦模糊去除：

python demo.py --task 单张图像散焦模糊去除 --input_dir './demo/degraded/portrait.jpg' --result_dir './demo/restored/'

训练与评估

去雨、运动模糊去除、散焦模糊去除和去噪任务的训练与测试说明分别位于各自目录中。下表汇总了相关链接，方便您快速导航：

任务	训练说明	测试说明	Restormer的可视化结果
去雨	链接	链接	下载
运动模糊去除	链接	链接	下载
散焦模糊去除	链接	链接	下载
高斯去噪	链接	链接	下载
真实场景去噪	链接	链接	下载

结果

我们针对不同的图像处理任务进行了实验，包括图像去雨、单张图像运动模糊去除、散焦模糊去除（分别针对单张图像和双像素数据）以及图像去噪（分别针对高斯噪声和真实场景数据）。

图像去雨（点击展开）

单张图像运动模糊去除（点击展开）

散焦模糊去除（点击展开）

S：单张图像散焦模糊去除。 D：双像素散焦模糊去除。

高斯图像去噪（点击展开）

上排：学习一个模型来处理多种噪声水平。 下排：为每种噪声水平单独训练一个模型。

灰度	彩色

真实图像去噪（点击展开）

引用

如果您使用 Restormer，请考虑引用以下文献：

@inproceedings{Zamir2021Restormer,
    title={Restormer: 高分辨率图像复原的高效Transformer}, 
    author={Syed Waqas Zamir 和 Aditya Arora 和 Salman Khan 和 Munawar Hayat 
            和 Fahad Shahbaz Khan 和 Ming-Hsuan Yang},
    booktitle={CVPR},
    year={2022}
}

联系方式

如有任何问题，请联系 waqas.zamir@inceptioniai.org

致谢： 本代码基于 BasicSR 工具箱和 HINet。

我们的相关工作

• 用于快速图像复原与增强的丰富特征学习，TPAMI 2022。论文 | 代码
• 多阶段渐进式图像复原，CVPR 2021。论文 | 代码
• 用于真实图像复原与增强的丰富特征学习，ECCV 2020。论文 | 代码
• CycleISP：通过改进的数据合成实现真实图像复原，CVPR 2020。论文 | 代码

Restormer 快速上手指南

Restormer 是一种高效的 Transformer 模型，专为高分辨率图像复原任务设计（CVPR 2022 Oral）。它在去雨、运动去模糊、散焦去模糊及图像去噪等任务上达到了业界领先（SOTA）的效果。

1. 环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 18.04+) 或 macOS
• Python: 3.8 或更高版本
• GPU: 支持 CUDA 的 NVIDIA 显卡（推荐显存 8GB 以上以处理高分辨率图像）
• 核心依赖:
  • PyTorch >= 1.7.1
  • torchvision
  • numpy
  • opencv-python
  • tqdm
  • einops
  • h5py

提示：国内开发者建议使用清华源或阿里源加速 Python 包的安装。

2. 安装步骤

第一步：克隆项目代码

git clone https://github.com/swz30/Restormer.git
cd Restormer

第二步：安装依赖库

您可以使用 pip 直接安装所需依赖。为了获得更快的下载速度，推荐使用国内镜像源：

# 使用清华源安装依赖
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

如果 requirements.txt 未包含所有必要组件，请手动安装核心库：

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

第三步：下载预训练模型

根据您需要执行的任务，从 Google Drive 下载对应的预训练模型权重，并放入项目根目录或指定文件夹。

• 去雨 (Deraining): 下载链接
• 运动去模糊 (Motion Deblurring): 下载链接
• 散焦去模糊 (Defocus Deblurring): 下载链接
• 去噪 (Denoising): 下载链接

(注：如果无法访问 Google Drive，可尝试项目提供的 Colab 演示或寻找国内网盘搬运资源)

3. 基本使用

安装完成后，您可以直接使用 demo.py 脚本对单张图片或文件夹进行图像复原。

命令格式

python demo.py --task Task_Name --input_dir path_to_images --result_dir save_images_here

使用示例

场景 1：对文件夹内的图片进行散焦去模糊 假设您有一组模糊图片位于 ./demo/degraded/ 目录，希望将修复后的图片保存至 ./demo/restored/：

python demo.py --task Single_Image_Defocus_Deblurring --input_dir './demo/degraded/' --result_dir './demo/restored/'

场景 2：直接处理单张图片 如果您只想处理单张名为 portrait.jpg 的图片：

python demo.py --task Single_Image_Defocus_Deblurring --input_dir './demo/degraded/portrait.jpg' --result_dir './demo/restored/'

支持的任务名称 (--task 参数)

在使用时，请将 Task_Name 替换为以下之一：

• Deraining (去雨)
• Motion_Deblurring (运动去模糊)
• Single_Image_Defocus_Deblurring (单图散焦去模糊)
• Dual_Pixel_Defocus_Deblurring (双像素散焦去模糊)
• Gaussian_Grayscale_Denoising (高斯灰度去噪)
• Gaussian_Color_Denoising (高斯彩色去噪)
• Real_Denoising (真实场景去噪)

---

更多高级用法（如重新训练模型或详细评估指标），请参考项目各子目录下的 README.md 文件。