【论文阅读】InstructIR: High-Quality Image Restoration Following Human Instructions

GitHub：https://github.com/mv-lab/InstructIR

摘要

图像恢复是一个基本问题，旨在从退化的图像中恢复出高质量的清晰图像。All-In-One 图像恢复模型能够利用退化特定的信息作为提示，引导恢复模型有效地恢复多种类型和不同程度的退化图像。本研究提出了首个使用人类编写指令来引导图像恢复模型的方法。通过自然语言提示，我们的模型能够在考虑多种退化类型的情况下，从退化的图像中恢复出高质量的图像。我们的方法，InstructIR，在多个恢复任务中（包括去噪、去雨、去模糊、去雾和低光图像增强）取得了最先进的结果，性能较以往的All-In-One恢复方法提升了+1dB。此外，我们的数据集和实验结果为基于文本引导的图像恢复和增强研究提供了一个全新的基准。

1. 研究背景

1. 图像恢复问题
图像恢复的目标是从退化图像中恢复出高质量的清晰图像。退化效果通常包括噪声、运动模糊、雾霾、低光范围等，这些退化现象普遍存在于计算机视觉中的低级任务中。这些效果可能由相机限制或恶劣的环境条件（如低光）引起。

2. 图像恢复的挑战
图像恢复是一个复杂的逆问题，因为对于给定的退化图像，恢复的过程可能有多个不同的解。很多传统方法专注于解决特定的退化问题，例如去噪、去模糊或去雾等。这些方法虽然在各自的任务上有效，但无法很好地处理其他类型的退化。为了更广泛地解决不同类型的退化问题，研究者开始提出多任务的图像恢复方法。

3. 多任务图像恢复模型的演变
一些方法使用通用的神经网络处理多种退化任务，但通常是针对每种退化类型独立训练一个模型。由于每种退化类型需要不同的模型，这种方法消耗了大量的计算资源。因此，最近的研究开始提出All-In-One（全能）图像恢复模型，这些模型能够考虑多种退化类型和不同的退化程度，通过使用单一的深度恢复模型来解决多个退化问题。例如，PromptIR和ProRes等方法，利用了学习到的引导向量（prompt embeddings）来帮助模型进行盲恢复，区别于传统的原始用户提示文本。

4. 文本引导图像恢复
与图像恢复相关的最新研究，如InstructPix2Pix，展示了使用文本提示引导图像生成和编辑模型的潜力。然而，这些方法并未涉及解决图像恢复中的逆问题。与之前的基于图像的退化分类方法不同，本文提出的InstructIR方法采用人类编写的自然语言指令来引导图像恢复模型。用户虽然可能缺乏特定领域的词汇，但通常能直观地了解需要修复的内容，这些信息可以用于引导模型的恢复过程。

5. 本文贡献
我们提出了一种全新的方法，InstructIR，该方法利用真实的人工编写指令来解决多任务图像恢复问题。通过综合实验，我们验证了文本指导在图像恢复和增强中的潜力，证明该方法在多个图像恢复任务（如去噪、去雨、去模糊、去雾和低光图像增强）中表现出色，超越了现有的全能恢复方法，提升了约**+1 dB**的性能。此外，我们的单一全能模型能够处理比许多现有方法更多的任务，并且在图像恢复任务上取得了最先进的结果。我们展示了使用该方法进行恢复的多个样本，证明了其广泛的适应性和强大的性能。

2. 跟随指令的图像恢复

本节将图像恢复任务视为一个监督学习问题，与以往的工作类似。具体步骤包括：

1. 生成训练数据：使用 GPT-4 基于自定义示例指令生成超过10,000个提示语，并解释了如何创建提示数据集（见第3.1节）。
2. 构建配对训练数据集：根据这些生成的提示语和相应的退化/清晰图像，构建了一个大规模的训练数据集。
3. 训练 InstructIR 模型：通过这些数据训练模型，并在多种类型的指令下对模型进行评估（包括真实的人工编写指令）。具体的文本编码器（第2.2节）和完整的模型（第2.3节）将会详细说明。

2.1 生成训练用的提示语

为什么使用指令？

受到 InstructPix2Pix 启发，本文采用人工编写的指令作为模型控制机制。与传统方法不同，用户无需提供清晰图像的示例或图像内容描述。指令提供了一种清晰且表达力强的交互方式，使用户能够准确地指出图像中的问题（即退化部分）。此外，指令的复杂度从模糊的指令（如“修复我的图像”）到精确的指令（如“去除噪点”）不等。

在训练过程中，系统将从不同退化类型的指令集中随机采样指令，进行自监督学习。当推理时，模型可以处理人类编写的指令，恢复或增强图像。训练时使用的提示语由GPT-4生成，并依据不同退化类型进行分类。

2.2 文本编码器

2.2.1 文本编码器的选择

文本编码器的作用是将用户输入的指令转换为固定大小的向量表示（文本嵌入）。与常用于图像生成的CLIP模型不同，退化类型的指令通常缺乏视觉内容，更多的是描述退化问题本身。为此，我们选择了一个纯文本的句子编码器（如 BGE-micro-v2），它在处理不同用户指令时具有较好的语义表达能力。

2.2.2 微调文本编码器

1. 文本编码器的作用

• 文本编码器 $E$ 将自然语言指令（例如“去除噪声”或“去模糊”）转换为固定大小的向量表示（embedding）。
• 这些向量表示包含指令的语义信息，作为图像修复模型的输入。

2. 为什么不训练整个编码器？

• 如果直接在小数据集上重新训练整个文本编码器 $E$，可能会导致过拟合，即在训练数据上表现良好，但在新数据上表现较差。
• 此外，文本编码器 $E$ 已经通过数百万条通用语言数据预训练过，具备良好的语义理解能力。完全重新训练可能破坏这种泛化能力。

3. 冻结编码器

• 为了避免过拟合，研究者冻结文本编码器，即不更新其权重（参数）。
• 重点放在适配文本编码器的输出，使其更适合图像修复任务。

4. 添加投影头

• 在冻结的文本编码器之上添加一个投影头，用来将原始的文本嵌入$E(t)$转换为图像修复模型所需的输入维度。

  投影的公式是：
  $\mathbf{e}=\mathrm{norm}(\mathbf{W}\cdot \mathrm{E}(t))$

  • $t$：输入的文本指令（如“去模糊我的图像”）。
  • $E(t)$：文本编码器生成的原始嵌入向量。
  • $W\in {\mathbb{R}}^{d_t\times d_v}$：一个可训练的矩阵，将文本嵌入维度$d_t$映射到修复模型输入维度$d_v$。
    -$\mathrm{norm}$：使用l2正则化对向量进行归一化，使其长度为1，保证嵌入向量的幅度一致性和稳定性。

5. 为什么要使用投影头？

• 轻量化：只训练矩阵$W$，参数较少，对数据需求也较低。
• 保留了预训练文本编码器的通用能力，同时对其输出进行了特定任务的优化。

6. 结果

• 这种方法确保生成的嵌入向量既能有效指导图像修复模型，又能保持对新指令的泛化能力，同时避免过拟合。
• 它在灵活性和效率之间取得了良好平衡。

2.2.3 结果可视化

1. 未训练的文本编码器的表现

• Figure 3a 显示，直接使用预训练文本编码器生成的嵌入，可以在某种程度上将不同任务的指令（如去雨、去噪）进行聚类。
• 但这些聚类并不理想，不同任务（degradation type）的边界可能重叠、不够清晰。
  

2. 经过投影头训练后的改进

• Figure 3b 显示，通过在文本编码器后添加一个投影头$W$，并对其进行训练，生成的嵌入聚类效果大大改善。
• 改善后的聚类特点：
  • 任务（去雨、去噪等）边界更加清晰：任务间的指令语义差异明显，聚类效果更好。
  • “超分辨率（super-res）”和“增强（enhancement）”的指令分布更广泛：这符合语言的逻辑，因为诸如“增加细节”这样的指令可能同时适用于增强、去模糊或去噪。

3. 投影头的作用

• 投影头$W$ 是一个简单的线性层，将文本嵌入的维度从$d_t=384$映射到修复模型的输入维度$d_v=256$。
• 它的作用：
  • 提取语义更符合图像修复需求的特定嵌入。
  • 共享的嵌入表示 $e$ 用于不同的任务模块。

4. 引入分类头的意图

• 为进一步提升训练效果和模型解释性，研究者引入了分类头 $C$。
• 分类头$C$的功能：
  • 将嵌入 $e$映射为不同任务的分类结果 $c$，其中 $c\in {\mathbb{R}}^D$，$D$是任务类别数（如去雨、去模糊、去噪等）。
  • 分类头 $C$ 是一个两层的多层感知机（MLP）。

5. 分类损失（Intent Classification Loss）

• 分类头的训练目标是通过添加分类损失，指导文本嵌入$e$ 更好地学习和区分不同任务类型。
• 具体实现：
  • 使用任务类别（如去模糊、去噪等）作为分类目标。
  • 优化投影头 $W$和分类头$C$，而不修改文本编码器 $E$本身。

6. 结果和性能

• 模型通过分类损失学到的嵌入不仅对图像修复模型有指导作用，还能更好地解释用户指令的意图。
• 经过几轮训练后，分类头对用户指令中涉及的退化类型（degradation type）达到了95%以上的分类准确率。

2.3 InstructIR

1. InstructIR 的结构
   InstructIR 包括两个主要组件：

• 文本编码器：负责将用户指令转换为嵌入表示，用于指导图像修复任务（在之前的 2.2 节详细讨论）。
• 图像模型：采用 NAFNet，一个基于 U-Net 结构的高效图像修复模型。NAFNet 用于处理图像特征，通过多个任务路由（task routing）技术实现单个模型处理多个任务。

2. 任务路由与多任务学习

• 任务路由（Task Routing） 是一种技术，允许模型根据输入任务选择和优先特定特征，这类似于专家混合模型（Mixture of Experts, MoE）。
• 为了支持多个图像修复任务（如去模糊、去噪、去雨等），InstructIR 提出了一个新模块，称为 Instruction Condition Block (ICB)，用以实现任务特定的特征转换。
  

3. Instruction Condition Block (ICB) 的机制

ICB 的工作流程如下：

1. 文本指导生成特征掩码：

   • 文本嵌入 $e$通过一个线性层 $W_c$，经过 Sigmoid 激活函数，生成任务相关的特征掩码$ m_c = \sigma(W_c \cdot e)$。
   • 这个掩码 $m_c$ 是一个维度为 $c$ 的“软二值掩码”，用于通道级的特征选择。

2. 特征通道的筛选与增强：

   • 掩码$m_c$ 与图像特征$F$ 进行逐通道乘法（channel-wise multiplication），实现特征筛选。
   • 筛选后的特征 $F_{l+1}$进一步传入 NAFBlock（卷积块）进行增强和处理。

3. 模块插入位置：

   • ICB 模块被嵌入到编码器和解码器的不同层级（encoder 和 decoder blocks），与 NAFBlocks 交替使用。

   • 在每一层 $l$，特征更新的公式为：
     $F_{l+1}=\text{ICB}(\text{Block}(F_l))$

   • 作用：

   • 掩码 $m_c$ 帮助模型根据任务选择最相关的通道特征，具有一定的解释性（权重高的特征对修复贡献更大）。

   • 通过这种方式，模型间接地学到多样化的滤波器，减少特征稀疏性。

4. 盲图像修复

• 盲修复的定义：

  • InstructIR 不依赖图像退化的显式信息（如噪声分布、模糊核、点扩散函数等）。
  • 模型通过图像和文本指令推断任务，因此被认为是一个盲图像修复模型。

• 与其他方法的对比：

  • 传统方法可能依赖于辅助的退化分类器来预测退化类型（如噪声种类）。
  • InstructIR 的独特性在于它直接结合指令和图像信息，无需显式的退化参数输入。

5. 关键优势

• 多任务适应性：ICB 模块允许模型动态调整特征，根据任务需求选择特定通道。
• 解释性：生成的掩码 $m_c$提供了解模型如何选择特征通道，权重越高的通道贡献越大。
• 高效性：只需训练线性层 $W_c$和卷积模块，计算成本相对较低。

3. 实验结果

模型评估和任务背景

We evaluate our model on 9 well-known benchmarks for different image restoration tasks: image denoising, deblurring, deraining, dehazing, real low-light enhancement, and photo-realistic image enhancement.

• 任务种类：模型适用于多种图像退化场景，例如降噪、去模糊、去雨条、去雾、低光增强和真实感图像增强。
• 基准测试：模型在 9 个著名的图像恢复任务基准数据集上进行了评估，覆盖了不同类型的图像退化。

We present extensive quantitative results in Table 2 and Table 3. We provide extensive comparisons with other all-in-one methods as well as task-specific methods.

• 结果展示：表 2 和表 3 展示了模型的定量评估结果，与其他“全能型”方法和任务专用方法做了全面比较。

Our single model successfully restores images considering different degradation types and levels.

• 统一模型：单个模型无需为特定任务单独优化，可以应对不同类型和程度的退化，证明其通用性。

---

实现细节

总体训练设计：模型支持端到端训练，图像模型无需预训练，文本模型使用预训练语言模型（句子编码器）。

文本编码器

We only need to train the text embedding projection and classification head (≈ 100K parameters). We initialize the text encoder with BGE-micro-v2, a distilled version of BGE-small-en.

• 关键模块：
  • 只需训练文本嵌入投影层和分类头，参数量约 10 万。
  • 文本编码器采用 BGE-micro-v2，这是 BGE-small-en 的简化版本。

---

The BGE encoders are BERT-like encoders pre-trained on large amounts of supervised and unsupervised data for general-purpose sentence encoding.

• BGE 模型：BERT 风格的编码器，适用于通用句子编码，基于大量监督和非监督数据进行预训练。

---

The BGE-micro model is a 3-layer encoder with 17.3 million parameters, which we freeze during training.

• 参数规模：BGE-micro 是一个 3 层编码器，参数量 1730 万，在训练时被冻结。

---

We also explore all-MiniLM-L6-v2 and CLIP encoders, however, we concluded that small models prevent overfitting and provide the best performance while being fast.

• 其他模型对比：尝试过其他小型编码器（如 MiniLM 和 CLIP），发现小模型能避免过拟合，性能最佳且速度快。

---

模型设计

• 主干网络：图像模型使用 NAFNet，包含 4 层编码器-解码器结构。

  • 编码器：每层块数分别为 [2, 2, 4, 8]。
  • 解码器：每层块数为 [2, 2, 2, 2]。
  • 中间层：编码器与解码器之间有 4 个中间块以增强特征。

• 跳跃连接：解码器跳跃连接采用“加法”而非传统的“拼接”，简化计算。

• 任务路由模块：在编码器和解码器中加入指令条件块（ICB），用于任务分配。

---

优化与训练策略

• 损失函数：

  • 图像恢复：使用 L1 损失（预测图像与真实图像的绝对误差）。
  • 文本分类：使用交叉熵损失（文本意图分类头）。

• 训练配置：

  • 批量大小：32。
  • 优化器：AdamW，学习率 5e-4。
  • 训练周期：500 轮，用 NVIDIA A100 约需 1 天。

• 数据增强：训练时对输入图像进行裁剪（256×256）和随机水平、垂直翻转。

• 参数规模：整个模型的参数量非常小：

  • 图像模型：16 M。
  • 文本投影层：100 k。
  • 语言模型：17 M。

推理性能：模型既适合标准 GPU 训练，也支持在计算资源受限环境下推理（如 Google Colab 上的 T4 GPU）。

多重退化任务实验结果

多任务恢复的初始设置

• 多任务类别：
  • 3D 设置：三种退化类型，包括降噪（denoising）、去雾（dehazing）、去雨条（deraining）。
    示例方法：AirNet。
  • 5D 设置：五种退化类型，新增去模糊（deblurring）和低光增强（low-light enhancement）。

---

5D 实验结果

3D 实验结果

3D 模型性能：表 3 显示了当任务数量减少到 3（即 3D 设置）时，InstructIR 的重建性能进一步提高，表现出更强的任务专用性。

---

指令的敏感性消融实验

实验设计

• 输入需求：InstructIR 需要输入退化图像和人类书写的指令（提示），因此构建了测试集，包含指令-图像配对数据。

• 性能依赖性：模型的性能与指令的精确性和模糊性密切相关。
  

• 表 4 数据对比了三种不同语言级别下的性能：

  1. Basic & Precise（基础且精确）：较详细的基础指令，性能最高。
  2. Basic & Ambiguous（基础但模糊）：模糊指令，性能下降。
  3. Real Users †（真实用户）：由业余摄影师撰写的精确指令，性能与第一种接近。

• 性能对比：较模糊的指令会导致图像增强（LOL）任务中性能显著下降。

---

5D 多任务模型（基础五任务）

InstructIR 的基础版本（5D）能够完成以下五项经典图像修复任务：

1. 去模糊（Deblurring）
2. 去噪声（Denoising）：例如处理高斯噪声（标准差为 15 的噪声）。
3. 去雨（Deraining）：移除图像中的雨滴效果。
4. 去雾（Dehazing）：处理图像中因雾霾导致的模糊效果。
5. 低光照增强（Low-light Enhancement）：增强图像的亮度和对比度。

实验表明，5D 模型在多个任务上的表现具有很高的性能指标，具体如表 6 所示：

• 在去模糊、去噪、低光照增强任务上，5D 模型均取得了优异的结果，尤其是 SSIM（结构相似性）和 PSNR（峰值信噪比）方面。
  

---

6D 多任务模型（增加单图像超分辨率 SISR）

新增任务：单图像超分辨率（SISR）

1. 任务描述：通过超分辨率技术将低分辨率图像恢复为高分辨率图像。
2. 训练数据：使用 DIV2K 数据集，这是超分辨率领域的标准数据集。
3. 生成低分辨率数据的方法：采用双三次（Bicubic）降采样，生成低分辨率（LR）图像；同时将高分辨率（HR）图像作为增强目标。
4. 模型特点：InstructIR 不直接进行图像上采样，而是通过细节增强优化高分辨率图像的质量。

性能分析

• 去模糊性能提升：增加 SISR 任务后，模型在去模糊任务上的性能有显著提高。这表明 SISR 和去模糊在图像细节增强上具有相关性，从而产生了协同效应。
• 对其他任务影响：新增任务对其他任务的性能没有明显的负面影响，具体表现如表 6 中的 6D 结果：
  • 去噪（Noise）：PSNR 为 31.39，与 5D 模型几乎相同。
  • 去模糊（Blur）：PSNR 提升至 29.73，SSIM 也有所提升。

---

7D 多任务模型（增加真实图像增强）

新增任务：真实图像增强（Real Image Enhancement）

1. 任务描述：通过模型增强图像的色彩、对比度等，使其更适合实际场景的使用，如社交媒体发布等。
2. 数据集：使用 MIT5K 数据集，这是一个大规模的真实图像增强数据集。
3. 任务特点：与前面任务（如去模糊、去噪）不同，真实图像增强不属于“逆问题”范畴，其目标是提升图像观感。

性能分析

• 新增任务性能：

  • 在 MIT5K 数据集上，InstructIR-7D 的表现优异，具体表现如表 5：
    • PSNR：24.65
    • SSIM：0.900
    • 色差（∆Eab，越低越好）：8.20
      

• 对其他任务的影响：

  • 添加真实图像增强任务后，对修复任务（如去模糊、去噪）的性能有轻微下降：
    • 去模糊任务：PSNR 从 29.73（6D） 降到 29.70（7D），SSIM 保持不变。
    • 去噪任务：PSNR 轻微下降至 31.37。
    • 低光照增强任务：PSNR 从 22.83 降至 22.81。
  • 尽管性能略有下降，但仍保持在较高水平。

---

指令控制的优势

1. 指令支持：

   • 模型可以通过自然语言指令实现图像修复和增强，例如：
     • “清除图像中的雨滴” → 进行去雨处理。
     • “提高分辨率” → 增强图像细节。
     • “调整图像亮度，使其适合发布在 Instagram 上” → 实现真实图像增强。

2. 连续指令执行：

   • 支持多次指令连续操作。例如：
     • 输入图像 → 去模糊 → 去雨 → 调整亮度。
   • 用户可以逐步优化图像，而无需重新生成。
     

3. 对抗性和分布外（OOD）指令处理：

   • 面对不合理或未见过的指令，模型会倾向于执行“恒等操作”（即不修改图像），这避免了对图像产生错误处理。
     

---

局限性

1. 复杂退化场景：模型在处理同时存在多种退化（如模糊和噪声）或未知分布的复杂图像时表现有限。
2. 未来改进方向：可通过：
   • 提供更真实的训练数据。
   • 扩展模型的复杂度（如更深的网络或更大的数据量）。

原文地址：https://blog.csdn.net/weixin_44329069/article/details/143802060

免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！