多阶段渐进式图像恢复（CVPR 2021）

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

---

新闻

我们很高兴地看到，我们的工作启发了 NTIRE 2021 挑战赛的获奖方案：

• 双像素散焦去模糊挑战赛 -- MRNet：用于双像素图像散焦去模糊的多级细化网络
• 图像去模糊挑战赛 -- HINet：用于图像恢复的半实例归一化网络

---

摘要： 图像恢复任务在重建图像时，需要在空间细节与高层次上下文信息之间取得复杂的平衡。本文提出了一种新颖的协同设计，能够最优地平衡这些相互竞争的目标。我们的主要方案是一种多阶段架构，逐步学习对退化输入的恢复函数，从而将整个恢复过程分解为更易于管理的步骤。具体而言，我们的模型首先使用编码器-解码器架构学习上下文特征，随后将其与保留局部信息的高分辨率分支相结合。在每个阶段，我们引入了一种新颖的逐像素自适应设计，利用原位监督注意力机制重新加权局部特征。这种多阶段架构的关键在于不同阶段之间的信息交换。为此，我们提出了一种双管齐下的方法：信息不仅按顺序从早期阶段传递到晚期阶段，还在特征处理模块之间建立了横向连接，以避免任何信息丢失。由此产生的紧密互联的多阶段架构被称为 MPRNet，在包括图像去雨、去模糊和去噪在内的多种任务的十个数据集上均表现出显著的性能提升。例如，在 Rain100L、GoPro 和 DND 数据集上，与当前最先进方法相比，PSNR 分别提高了 4 dB、0.81 dB 和 0.21 dB。

网络架构

MPRNet 总体框架	监督注意力模块（SAM）

安装

该模型基于 PyTorch 1.1.0 构建，并在 Ubuntu 16.04 环境中进行了测试（Python 3.7、CUDA 9.0、cuDNN 7.5）。

安装步骤如下：

conda create -n pytorch1 python=3.7
conda activate pytorch1
conda install pytorch=1.1 torchvision=0.3 cudatoolkit=9.0 -c pytorch
pip install matplotlib scikit-image opencv-python yacs joblib natsort h5py tqdm

安装预热调度器：

cd pytorch-gradual-warmup-lr; python setup.py install; cd ..

快速运行

要测试针对您自己的图像的预训练模型（去模糊、去雨、去噪），请运行：

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

以下是一个去模糊示例：

python demo.py --task Deblurring --input_dir ./samples/input/ --result_dir ./samples/output/

训练与评估

去模糊、去雨和去噪的训练与测试代码分别位于各自的目录中。

结果

实验针对不同的图像处理任务进行，包括图像去模糊、去雨和去噪。MPRNet 生成的图像可从 Google Drive 链接下载：去模糊、去雨，以及 去噪。

图像去模糊（点击展开）

合成数据集上的去模糊。	真实数据集上的去模糊。

图像去雨（点击展开）

图像去噪（点击展开）

引用

如果您使用 MPRNet，请考虑引用以下内容：

@inproceedings{Zamir2021MPRNet,
    title={Multi-Stage Progressive Image Restoration},
    author={Syed Waqas Zamir and Aditya Arora and Salman Khan and Munawar Hayat
            and Fahad Shahbaz Khan and Ming-Hsuan Yang and Ling Shao},
    booktitle={CVPR},
    year={2021}
}

联系方式

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

我们的相关工作

• 用于快速图像恢复与增强的特征增强学习，TPAMI 2022。论文 | 代码
• Restormer：面向高分辨率图像恢复的高效Transformer，CVPR 2022。论文 | 代码
• 用于真实图像恢复与增强的特征增强学习，ECCV 2020。论文 | 代码
• CycleISP：通过改进的数据合成实现真实图像恢复，CVPR 2020。论文 | 代码

MPRNet 快速上手指南

MPRNet (Multi-Stage Progressive Image Restoration) 是一个用于图像恢复（去模糊、去雨、去噪）的多阶段渐进式深度学习模型，发表于 CVPR 2021。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 16.04+)
• Python: 3.7
• PyTorch: 1.1.0
• CUDA: 9.0
• cuDNN: 7.5

注意：该工具基于较旧的 PyTorch 版本构建。如果您使用更新的 GPU 或驱动，可能需要调整 CUDA 版本或尝试在新版 PyTorch 上复现（可能存在兼容性风险）。

安装步骤

1. 创建 Conda 环境并安装基础依赖

conda create -n pytorch1 python=3.7
conda activate pytorch1
conda install pytorch=1.1 torchvision=0.3 cudatoolkit=9.0 -c pytorch

(国内用户推荐使用清华源加速安装)

conda install pytorch=1.1 torchvision=0.3 cudatoolkit=9.0 -c pytorch -c https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch/

2. 安装 Python 第三方库

pip install matplotlib scikit-image opencv-python yacs joblib natsort h5py tqdm

(国内用户推荐使用阿里源或清华源)

pip install matplotlib scikit-image opencv-python yacs joblib natsort h5py tqdm -i https://pypi.tuna.tsinghua.edu.cn/simple

3. 安装预热学习率调度器 (Warmup Scheduler)

克隆或进入 pytorch-gradual-warmup-lr 目录后执行：

cd pytorch-gradual-warmup-lr; python setup.py install; cd ..

基本使用

下载预训练模型后，您可以直接使用 demo.py 对本地图片进行推理。支持的任务包括：Deblurring (去模糊), Deraining (去雨), Denoising (去噪)。

命令格式

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

使用示例：图像去模糊

假设您已将去模糊的预训练模型放置好，并对 ./samples/input/ 目录下的图片进行处理，结果保存至 ./samples/output/：

python demo.py --task Deblurring --input_dir ./samples/input/ --result_dir ./samples/output/

只需将 --task 参数替换为 Deraining 或 Denoising，即可执行相应的图像恢复任务。