Uformer：一种用于图像修复的通用U型Transformer（CVPR 2022）

Zhendong Wang, Xiaodong Cun, Jianmin Bao, Wengang Zhou, Jianzhuang Liu, Houqiang Li

论文链接：[Arxiv] [CVPR]

更新：

• 2024.10.24 修复失效链接。
• 2022.07.06 上传了Uformer的新代码和模型。
• 2022.04.09 上传了Uformer在去噪（SIDD、DND）、运动模糊去除（GoPro、HIDE、RealBlur-J/-R）以及散焦模糊去除（DPDD）方面的结果。
• 2022.03.02 Uformer已被CVPR 2022接收！ :fire:
• 2021.11.30 在Arxiv链接中更新了Uformer。新的代码、模型和结果将随后上传。
• 2021.10.28 发布了Uformer32在SIDD和DND上的结果。
• 2021.09.30 发布了用于SIDD去噪的预训练Uformer16模型。
• 2021.08.19 发布了一个预训练模型（Uformer32）！并添加了用于计算FLOP/GMAC的脚本。
• 2021.07.29 添加了一个用于测试任意分辨率图像上预训练模型的脚本。

---

在本文中，我们提出了一种高效且有效的基于Transformer的架构——Uformer。该架构采用Transformer模块构建层次化的编码器-解码器网络，用于图像修复任务。Uformer包含两个核心设计，使其非常适合这一任务。第一个关键要素是局部增强窗口Transformer模块，其中我们使用非重叠窗口自注意力机制来降低计算复杂度，并在前馈网络中引入深度可分离卷积，以进一步提升其捕捉局部上下文的能力。第二个关键要素是探索了三种跳跃连接方案，以有效地将信息从编码器传递到解码器。凭借这两大设计，Uformer具备强大的能力来捕捉图像修复所需的有用依赖关系。我们在多个图像修复任务上的大量实验表明，Uformer具有显著优势，包括图像去噪、去雨、去模糊和去摩尔纹等。我们期望我们的工作能够激励更多研究者探索基于Transformer的架构在低层视觉任务中的应用。

软件包依赖

该项目基于PyTorch 1.9.0、Python3.7和CUDA11.1构建。软件包依赖可通过以下命令安装：

pip install -r requirements.txt

预训练模型

• Uformer_B：SIDD | GoPro

预训练模型的结果

• Uformer_B：SIDD | DND | GoPro | HIDE | RealBlur-J | RealBlur-R | DPDD

数据准备

去噪

对于SIDD的训练数据，可以从官方网址下载SIDD-Medium数据集。 然后通过以下命令生成用于训练的补丁：

python3 generate_patches_SIDD.py --src_dir ../SIDD_Medium_Srgb/Data --tar_dir ../datasets/denoising/sidd/train

对于SIDD和DND的评估，可以从这里下载数据。

去模糊

若要在 GoPro 数据集上进行训练，并在 GoPro、HIDE、RealBlur-J 和 RealBlur-R 数据集上进行评估，您可以从这里下载数据。

然后将所有去噪数据放入 ../datasets/denoising 目录下，所有去模糊数据放入 ../datasets/deblurring 目录下。

训练

去噪

要在 SIDD 数据集上训练 Uformer，您可以通过以下命令开始训练：

sh script/train_denoise.sh

去模糊

要在 GoPro 数据集上训练 Uformer，您可以通过以下命令开始训练：

sh script/train_motiondeblur.sh

评估

要评估 Uformer，您可以运行：

sh script/test.sh

针对每个数据集的评估，您需要取消注释相应的行。

计算开销

我们提供了一个简单的脚本，用于自行计算 FLOPs。该脚本已添加到 model.py 文件中。您可以更改配置并运行：

python3 model.py

此仓库中手动计算的 GMACs 与主论文中的结果略有差异，但并不影响结论。我们将在后续对论文进行修正。

引用

如果您在研究中使用了本项目，请考虑引用以下文献：

@InProceedings{Wang_2022_CVPR,
    author    = {Wang, Zhendong and Cun, Xiaodong and Bao, Jianmin and Zhou, Wengang and Liu, Jianzhuang and Li, Houqiang},
    title     = {Uformer: A General U-Shaped Transformer for Image Restoration},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    month     = {June},
    year      = {2022},
    pages     = {17683-17693}
}

致谢

本代码大量借鉴了 MIRNet 和 SwinTransformer 的实现。

联系方式

如有任何问题或建议，请联系我们（王振东：ZhendongWang6@outlook.com，岑晓东：vinthony@gmail.com）。

Uformer 快速上手指南

Uformer 是一个基于 Transformer 的通用 U 型架构，专为图像复原任务（如去噪、去模糊、去雨等）设计。本指南将帮助你快速配置环境并运行预训练模型。

环境准备

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

• 操作系统: Linux (推荐) 或 macOS
• Python: 3.7+
• PyTorch: 1.9.0
• CUDA: 11.1 (如需 GPU 加速)
• 硬件: 支持 CUDA 的 NVIDIA 显卡（推荐显存 8GB 以上以运行大模型）

安装步骤

1. 克隆项目代码

   git clone https://github.com/ZhendongWang6/Uformer.git
   cd Uformer
   

2. 安装依赖包 建议使用国内镜像源（如清华源）以加快下载速度：

   pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   

   注意：如果自动安装失败，请手动安装对应版本的 torch 和 torchvision。

3. 下载预训练模型 根据任务需求下载对应的预训练权重，并放入项目根目录或指定文件夹：

   • 图像去噪 (SIDD): 下载链接
   • 运动去模糊 (GoPro): 下载链接

基本使用

1. 数据准备

在使用前，需将测试数据整理到相应目录。以去噪任务为例：

• 下载 SIDD 或 DND 测试数据集。
• 将数据放置于 ../datasets/denoising 目录下（具体结构参考原仓库说明）。

2. 运行测试/推理

修改 script/test.sh 脚本，取消注释你需要评估的数据集对应行，然后执行：

sh script/test.sh

该脚本会自动加载预训练模型并对指定数据集进行图像复原处理，结果通常保存在输出目录中。

3. (可选) 计算模型复杂度

如需计算模型的 FLOPs/GMACs，可直接运行：

python3 model.py

提示：对于自定义图像的任意分辨率测试，原仓库提供了相关脚本支持，具体可参考 scripts 目录下的测试逻辑进行调整。