PromptIR：用于一体化盲图像恢复的提示方法（NeurIPS'23）

Vaishnav Potlapalli、Syed Waqas Zamir、Salman Khan 和 Fahad Shahbaz Khan

---

摘要：图像恢复是指从退化版本中恢复高质量的干净图像。基于深度学习的方法显著提升了图像恢复性能，然而，它们对不同退化类型和程度的泛化能力有限。这限制了其在现实世界中的应用，因为需要为每种特定退化训练单独的模型，并且必须知道输入图像的退化类型才能应用相应的模型。我们提出了一种基于提示的学习方法——PromptIR，用于一体化图像恢复，能够有效恢复各种类型和程度退化的图像。具体而言，我们的方法使用提示来编码退化特异性信息，然后利用这些信息动态引导恢复网络。这使得我们的方法能够泛化到不同的退化类型和程度，同时在图像去噪、去雨和去雾任务上仍能达到当前最先进的水平。总体而言，PromptIR提供了一个通用且高效的插件模块，只需少量轻量级提示即可在无需事先了解退化类型的情况下，恢复各类退化程度的图像。

---

网络架构

安装与数据准备

请参阅 INSTALL.md，了解运行本代码库所需的依赖安装及数据集准备步骤。

训练

在 data/ 目录下准备好训练数据后，使用以下命令启动模型训练：

python train.py

通过 de_type 参数可以选择要训练的退化类型组合。默认设置为所有三种退化类型（噪声、雨和雾）。

示例用法：若仅想训练去雨和去雾任务：

python train.py --de_type derain dehaze

测试

在 test/ 目录下准备好测试数据后，将模型检查点文件放置于 ckpt 目录中。预训练模型可从 这里 下载，也可在 releases 标签页中获取。执行评估时，使用以下命令：

python test.py --mode {n}

其中 n 是一个数字，用于设置要评估的任务：0 表示去噪，1 表示去雨，2 表示去雾，3 表示一体化设置。

示例用法：若要一次性测试所有退化类型，运行：

python test.py --mode 3

演示

要获得模型的可视化结果，可以使用 demo.py 脚本。将保存的模型文件放入 ckpt 目录后，运行：

python demo.py --test_path {path_to_degraded_images} --output_path {save_images_here}

示例用法：对某一目录下的图像进行推理：

python demo.py --test_path './test/demo/' --output_path './output/demo/'

直接对单张图像进行推理：

python demo.py --test_path './test/demo/image.png' --output_path './output/demo/'

若在运行 demo.py 时希望使用分块处理选项，需将 --tile 参数设置为 True。可通过 --tile_size 和 --tile_overlap 参数分别调整分块大小和重叠区域。

结果

PromptIR 框架在一体化设置下训练后的性能结果

表格

可视化结果

PromptIR 模型在一体化设置下评估的可视化结果可从 这里 下载。

引用

如果您使用了我们的工作，请考虑引用以下内容：

@inproceedings{potlapalli2023promptir,
  title={PromptIR: Prompting for All-in-One Image Restoration},
  author={Potlapalli, Vaishnav and Zamir, Syed Waqas and Khan, Salman and Khan, Fahad},
  booktitle={Thirty-seventh Conference on Neural Information Processing Systems},
  year={2023}
}

联系方式

如有任何问题，请联系 pvaishnav2718@gmail.com。

致谢：本代码基于 AirNet 和 Restormer 两个仓库开发而成。

PromptIR 快速上手指南

PromptIR 是一个基于提示学习（Prompt Learning）的通用图像复原框架，能够无需预知退化类型，即可对去噪、去雨、去雾等多种退化图像进行“多合一”的高质量复原。

环境准备

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

• 操作系统: Linux (推荐 Ubuntu 18.04+) 或 Windows (需配置 WSL2 或兼容环境)
• Python: 3.8 或更高版本
• GPU: 支持 CUDA 的 NVIDIA 显卡 (推荐显存 >= 8GB)
• 核心依赖:
  • PyTorch (建议 1.9.0+)
  • torchvision
  • numpy, pillow, tqdm, yaml 等常用科学计算库

注意：本项目代码基于 AirNet 和 Restormer 构建，具体依赖版本细节请参考项目根目录下的 INSTALL.md 文件。国内用户可使用清华源或阿里源加速 PyTorch 及相关包的安装。

安装步骤

1. 克隆仓库

   git clone https://github.com/vaishnavrao/PromptIR.git
   cd PromptIR
   

2. 安装依赖 根据 INSTALL.md 指示安装所需库。通常可以使用 pip 直接安装基础依赖（建议使用国内镜像源加速）：

   pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
   pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   

   (注：若项目中未提供 requirements.txt，请手动安装上述核心依赖)

3. 下载预训练模型 将预训练权重文件下载并放置于项目根目录下的 ckpt 文件夹中。

   • 下载地址：Google Drive 或查看 Releases 页面。
   • 确保文件名与代码中引用的名称一致（通常为 .pth 格式）。

4. 准备数据

   • 训练：将训练数据集整理好放入 data/ 目录。
   • 测试/演示：将待处理的退化图片放入指定目录（如 test/demo/）。

基本使用

1. 快速推理（Demo）

这是最简单的使用方式，可直接对单张图片或文件夹内的图片进行复原处理。

处理文件夹中的图片：

python demo.py --test_path './test/demo/' --output_path './output/demo/'

处理单张图片：

python demo.py --test_path './test/demo/image.png' --output_path './output/demo/'

高级选项（大图片分块处理）： 如果图片分辨率较大导致显存不足，可开启分块（Tiling）模式：

python demo.py --test_path './test/demo/' --output_path './output/demo/' --tile True --tile_size 512 --tile_overlap 32

2. 模型评估（Testing）

如果您需要量化评估模型在特定任务上的性能，请使用 test.py。

评估所有任务（去噪、去雨、去雾）：

python test.py --mode 3

仅评估单一任务：

• 0: 去噪 (Denoising)
• 1: 去雨 (Deraining)
• 2: 去雾 (Dehazing)

示例（仅评估去雨）：

python test.py --mode 1

3. 模型训练（Training）

如需从头训练或微调模型，请确保数据已准备好，然后运行：

默认训练（包含去噪、去雨、去雾三种退化）：

python train.py

指定退化类型训练： 例如，仅训练去雨和去雾模型：

python train.py --de_type derain dehaze