即插即用系列 | PromptIR：MBZUAI提出一种基于Prompt的全能图像恢复网络

Title: PromptIR: Prompting for All-in-One Blind Image Restoration

PDF: https://arxiv.org/pdf/2306.13090v1.pdf

Code: https://github.com/va1shn9v/promptir

导读

图像恢复是从其受损版本中恢复高质量清晰图像的过程。基于深度学习的方法显著提升了图像恢复性能，然而，它们在不同类型和级别的退化上的泛化能力有限。这限制了它们在实际应用中的使用，因为需要针对每种具体的退化进行单独训练模型，并了解输入图像的退化类型才能应用相应的模型。本文介绍了一种基于提示的学习方法，称为PromptIR，用于全能图像恢复，可以有效地从各种类型和级别的退化中恢复图像。具体而言，本文方法使用提示来编码退化特定信息，并动态引导恢复网络。 这使得我们的方法能够推广到不同类型和级别的退化，并在图像去噪、去雨和去雾方面取得了最先进的结果。总的来说，PromptIR提供了一个通用且高效的插件模块，只需少量轻量级提示即可用于恢复各种类型和级别的受损图像，无需事先了解图像中存在的损坏信息。

引言

在图像获取过程中，常常会出现各种退化现象，如噪声、模糊、雾霾、雨水等，这些问题通常是由于相机的物理限制或环境条件不合适导致的。基于深度神经网络的方法在解决图像恢复问题时有不同的方法。一些方法在网络中引入了显式的任务特定知识来处理相应的恢复任务，如降噪、去模糊和去雾。然而，这些方法在特定的退化类型和程度之外缺乏泛化性。因此，迫切需要开发一种能够有效恢复各种类型和程度的退化图像的一体化方法。

最近的一种方法AirNet，通过采用对比学习范式解决了一体化恢复任务。这涉及训练一个额外的编码器来区分各种类型的图像退化。尽管AirNet取得了最先进的结果，但它在建模不同污染类型的完全解耦表示方面存在困难。此外，使用附加的编码器进行对比学习会导致更高的训练负担，因为需要两阶段的训练方法。

::: block-1

PromptIR提出了一个即插即用的提示模块，它隐式地预测与退化条件相关的提示，以引导具有未知退化的输入图像的恢复过程。来自提示的指导被注入到网络的多个解码阶段，其中包含少量可学习参数。这样可以学习一个全能的统一模型，可以在多个图像恢复任务（如去除雨水、去雾和降噪）中表现出色。
:::

为了克服这些挑战，本文提出了一种基于提示学习的方法来执行一体化图像恢复。该方法利用提示（一组可调参数），用于编码关于各种图像退化类型的重要区分信息（如下图2所示）。通过将提示与主恢复网络的特征表示相互作用，我们动态地增强表示，以获得具有退化特定知识的适应性，这种适应性使网络能够通过动态调整其行为有效地恢复图像。

::: block-1

该图显示了在PromptIR和最先进的AirNet中使用的退化嵌入的tSNE图。不同的颜色表示不同的退化类型。每个任务的嵌入更好地聚集在一起，显示了提示标记学习具有区分性的退化上下文的有效性，从而有助于恢复过程。
:::

本文主要亮点包括：

• 本文提出了一种基于提示的一体化恢复框架PromptIR，它仅依赖于输入图像来恢复清晰图像，而不需要任何关于图像中存在的退化的先验知识。
• 本文提示块是一个可轻松集成到任何现有恢复网络中的插件模块。它由提示生成模块（PGM）和提示交互模块（PIM）组成。提示块的目标是生成与输入条件相关的提示（通过PGM），这些提示具有有用的上下文信息，以指导恢复网络（通过PIM）有效地消除输入图像中的破坏。
• 本文实验证明了PromptIR的动态适应行为，在包括图像降噪、去雨和去雾在内的各种图像恢复任务上实现了最先进的性能。

方法

::: block-1

PromptIR方法在编码和解码阶段使用了UNet网络架构，其中包含了Transformer块。该框架的主要组件是提示块，由两个模块组成：提示生成模块（PGM）和提示交互模块（PIM）。提示生成模块使用输入特征Fl和提示组件生成与输入条件相关的提示P。然后，提示交互模块通过Transformer块使用生成的提示动态调整输入特征。提示与解码器特征在多个级别交互，以丰富特定于退化的上下文信息。
:::

PromptIR使用提示块来生成可学习的提示参数，并在恢复过程中利用这些提示来指导模型。框架通过逐级编码器-解码器将特征逐步转换为深层特征，并在解码器中引入提示块来辅助恢复过程。提示块在解码器的每个级别中连接，隐式地为输入特征提供关于退化类型的信息，以实现引导恢复。总体来说，PromptIR框架通过逐级编码和解码以及引入提示块的方式实现图像恢复任务。

Prompt Block

本文提出的PromptIR方法借鉴了在自然语言处理和计算机视觉任务中使用的基于提示的技术。在这些任务中，基于提示的技术已经被用于对在源任务上训练的大型固定模型进行参数高效微调，以适应目标任务。基于提示的技术之所以有效，是因为它们能够有效地将任务特定的上下文信息编码到提示组件中。在PromptIR中，提示组件是可学习的参数，与输入特征进行交互，以丰富它们的退化类型信息。提示块由两个关键组件组成：提示生成模块（PGM）和提示交互模块（PIM）

Prompt Generation Module (PGM)

提示组件 $P_c$ 是一组可学习的参数，与输入特征交互，嵌入了退化信息。一种直接的特征-提示交互方法是直接使用学习到的提示来校准特征。然而，这种静态方法可能会产生次优结果，因为它对输入内容是无知的。因此，本文提出了提示生成模块（PGM），它从输入特征中动态预测基于注意力的权重，并将这些权重应用于提示组件，生成与输入条件相关的提示 $P$。此外，PGM创建了一个共享空间，促进了提示组件之间的相关知识共享。

为了从输入特征 $F_l$ 生成提示权重，PGM首先对空间维度进行全局平均池化（GAP），生成特征向量 $v\in {\mathbb{R}}^{\hat{C}}$。接下来将 $v$ 通过通道缩减的卷积层，得到一个紧凑的特征向量，然后进行softmax操作，从而得到提示权重 $w\in {\mathbb{R}}^N$。最后使用这些权重对提示组件进行调整，接着应用一个 $3\times 3$ 的卷积层。总体而言，PGM的过程可以概括为：

$$P=\text{Conv3x3}\left(\sum _{c=1}^N{w}_i{P}_c\right),w_i=\text{Softmax}\left(\text{Conv1x1}\left(\text{GAP}(F_l)\right)\right)(2)$$

由于在推理阶段，恢复网络需要能够处理不同分辨率的图像，不能使用具有固定尺寸的提示组件 $P_c$。因此，作者对提示组件进行双线性插值操作，将其放大到与输入特征相同的尺寸。

Prompt Interaction Module (PIM)

PIM的主要目标是实现输入特征$F_l$和提示$P$之间的交互，以实现有指导的恢复过程。

在PIM中，沿着通道维度将生成的提示与输入特征进行拼接。接下来将拼接后的表示通过一个Transformer块进行处理，该块利用提示中编码的退化信息来转换输入特征。

本文的主要贡献是提示块，它是一个插件模块，与具体的架构无关。因此，在提出的PromptIR框架中，作者使用了现有的Transformer块，而不是开发一个新的块。Transformer块由两个顺序连接的子模块组成：多转置卷积头转置注意力（MDTA）和门控转置卷积前馈网络（GDFN）。MDTA在通道而不是空间维度上应用自注意操作，并具有线性复杂度。GDFN的目标是以可控的方式转换特征，即抑制信息较少的特征，只允许有用的特征在网络中传播。PIM的整体过程为：

$${\hat{F}}_l=\text{Conv3x3}\left(\text{GDFN}\left(\text{MDTA}[F_l;P]\right)\right)(3)$$

其中 $[;]$ 表示拼接操作。MDTA的公式为 $Y=W_{p}V\cdot \text{Softmax}(K\cdot Q/\alpha )+X$，其中 $X$ 和 $Y$ 分别表示输入和输出特征。$Q$、$K$ 和 $V$ 分别表示通过应用 1×1 点卷积后跟随 3×3 深度卷积在层归一化的输入特征图上获得的查询、键和值的投影。$W_p$ 是点卷积，$\alpha$ 是可学习的缩放参数，$(\cdot )$ 表示点积交互。GDFN的过程定义为 $Z=W_p^0\left(\varphi \left(W_d^1{W}_p^1(\text{LN}(Y))\right)\odot W_d^2{W}_p^2(\text{LN}(Y))\right)+Y$。其中，$W_d^{(\cdot )}$ 是 3×3 的深度卷积，$\odot$ 表示逐元素乘法，$\varphi$ 是GELU非线性激活函数，$\text{LN}$ 是层归一化。

实验

::: block-1

表1：全能恢复设置下的比较：使用单个模型在来自不同退化类型的图像组合数据集上训练。在不同任务上取平均值时，PromptIR相比之前的全能方法AirNet有0.86 dB的显著增益。
:::

::: block-1

图4：对SOTS数据集上的全能方法进行去雾比较。PromptIR生成的结果在视觉上比之前的最先进方法AirNet更好。
:::

::: block-1

图5：基于Rain100L数据集对全能方法进行图像去雨比较。PromptIR方法有效地去除雨滴，生成无雨痕的图像。
:::

结论

现有基于深度神经网络的图像恢复模型通常只适用于特定的退化类型，并且在其它退化类型上泛化能力有限。然而，在实际应用中，需要使用单一统一模型处理多种退化类型，而不是依赖于特定退化类型的模型，因为后者缺乏泛化能力并需要对输入中的具体退化类型有先验知识。为此，本文提出了一个即插即用的提示块，可以与输入特征进行交互，动态调整表示，使得恢复过程能够适应相关的各类退化任务。通过将提示块集成到最先进的恢复模型中，本文展示了提示块在全能图像恢复中的实用性，在图像去噪、去雨和去雾任务上取得了显著的改进。