【读点论文】All-In-One Image Restoration for Unknown Corruption用对比学习统一方法实现多种噪声图片的有效处理

All-In-One Image Restoration for Unknown Corruption

Abstract

在本文中，我们研究了图像恢复中的一个挑战性问题，即如何开发一种能够从各种未知损坏类型和程度中恢复图像的一体化方法 。为此，我们提出了一种一体化图像恢复网络（AirNet），由两个神经模块组成，称为基于对比的退化编码器（CBDE）和退化引导恢复网络（DGRN）。AirNet 的主要优势有两个方面。首先，它是一种一体化解决方案，可以在一个网络中恢复各种退化图像。其次，AirNet 不受损坏类型和级别的先验影响，仅使用观察到的损坏图像进行推理。这两个优势使 AirNet 在现实世界场景中享有更好的灵活性和更高的经济性，其中损坏的先验很难知道，并且退化会随着空间和时间而变化。大量的实验结果表明，所提出的方法在四个具有挑战性的数据集上优于 17 个图像恢复基线。代码可在 https://github.com/XLearning-SCU/2022-CVPR-AirNet 上找到。
论文地址：Li_All-in-One_Image_Restoration_for_Unknown_Corruption_CVPR_2022_paper.pdf (thecvf.com)。CVPR-2022

Introduction

单幅图像复原旨在从给定的劣化对应关系（例如嘈杂、下雨或朦胧的图像）中生成视觉上令人愉悦的高质量图像。在过去的几年中，图像复原已广泛应用于从自动驾驶到医学成像和监控等许多现实世界应用中。尽管在去噪、去模糊、去雨和去雾等特定领域已经取得了有希望的成果，但图像恢复在实践中遇到了以下障碍。一方面，有必要了解正确的损坏（即退化）以选择有竞争力的模型，因为几乎所有现有方法都只能处理特定的退化。一旦退化类型甚至损坏率发生变化，由于实际情况与模型构建或训练所采用的先验不一致，模型将获得不理想的性能。
另一方面，退化通常会在复杂的环境中发生变化，例如，自动驾驶汽车可能会连续甚至同时遭遇阴雨天气。总之，人们非常希望开发一种能够从各种未知(请注意，在本文中，"未知"指的是非特定的而不是看不见的损坏，"多重退化"是指给定的图像仅包含一个退化，但数据集将包含多个退化)损坏类型和程度中恢复图像的一体化方法，如下图所示。据我们所知，到目前为止，这种不具体的图像恢复问题几乎没有被触及。
- 我们的基本思想的说明。如图所示，大多数现有的多重降级方法通过将输入发送到专门设计的头部并使用相应尾部的输出来处理每个损坏。因此，它们需要提前获取损坏信息以指定校正的头部和尾部。不同的是，我们的一体化图像恢复网络（AirNet）不受损坏类型和级别的先验限制，因此在现实世界场景中具有更好的灵活性和更高的经济性。
为了解决上述问题，我们提出了由两个模块组成的一体化图像恢复网络（AirNet）。具体而言，基于对比的退化编码器（CBDE）旨在通过利用具有相同退化的图像的一致性和不同退化中存在的不一致性来学习退化表示 。在CBDE学习到的退化表示的指导下，退化引导恢复网络（DGRN）旨在恢复具有各种退化的图像。得益于CBDE和DGRN的结合，AirNet享有两个备受期待的优点，即i）它提供了一个一体化解决方案来恢复具有不同损坏类型和比率的图像；ii）它不受损坏类型和比率的先验知识的限制。值得注意的是，所提到的一体化解决方案在某些方面不同于现有的所谓的统一图像恢复方法。一方面，其他方法必须指定损坏类型和比率，而我们的方法则不需要。另一方面，它们通常将多个退化视为具有多个输入和输出头的多任务学习问题，其中每个输入和输出头对应于具有给定损坏率的预定损坏。相比之下，AirNet 是一个单通道网络，不区分不同的损坏类型和比率，因此具有更好的灵活性和更高的经济性。综上所述，本研究的贡献和创新之处如下：
- 据我们所知，AirNet 可能是首批以一体化方式从多个损坏中恢复图像的方法之一。由于我们的方法不需要任何降级信息来提前恢复，因此它可能更接近现实世界的情况。
- AirNet 采用双重方式工作，即从观察到的图像中对比学习退化表示，然后使用学习到的退化表示恢复干净的图像。需要指出的是，对比学习的成功在很大程度上依赖于正负对的构建。在本文中，我们展示了一种新方法，该方法可以有效捕捉多重退化的遗传特征。
- 在不失普遍性的前提下，我们进行了大量实验来验证 AirNet 在去噪、去雨和去雾方面的有效性，并与 17 条基线进行了比较。

在本节中，我们将简要回顾本文涉及的问题和方法的一些最新发展，即图像恢复和对比学习。

Image Restoration

根据本文的重点，现有的图像恢复方法可以分为两类，即单次退化图像恢复（IRSD）和多次退化图像恢复（IRMD）。
单一退化的图像恢复：IRSD 旨在从仅被具有固定损坏率的特定退化类型破坏的退化观测中恢复干净的图像。例如，作为开创性的深度去噪方法之一，DnCNN 无法处理多退化情况，即使在训练期间未见噪声比也会失败。其他图像恢复任务也面临着类似的挑战，例如去模糊、去雨和去雾。最近，一些研究显示出对不同退化的一定通用性。然而，他们需要针对不同的退化训练不同的模型，这并不是实践中预期的一体化解决方案。
多重退化图像恢复：最近，一些研究通过采用多输入和输出网络结构将注意力转向 IRMD。例如，[All in One Bad Weather Removal using Architectural Search] 提出了一种一体化模型来处理多种恶劣天气退化（例如雨、雾和雪），每种退化都由一个编码器专门处理。Chen 等人 [Pre-Trained Image Processing Transformer] 提出了一种基于 Transformer 的图像恢复方法，该方法使用多头和多尾架构来处理多重退化。与我们的方法最相似的方法可能是 [A General Decoupled Learning Framework for Parameterized Image Operators]。然而，该方法仍然需要知道输入的一些先验（例如噪声比和 JPEG 质量），以便以元学习的方式参数化网络。总而言之，尽管上述方法已经朝着 IRMD 迈进了一步，但它们仍然需要提前获得退化信息，以便将输入发送到校正头或生成元信息。

Contrastive Learning

对比学习是最先进的无监督表示学习方法，旨在最大化正对之间的相似性，同时最小化负对之间的相似性，其中正对和负对是通过数据增强获得的。最近，一些研究表明对比学习在图像恢复中的有效性。值得注意的是，尽管 DASR 和我们的 AirNet 都利用对比学习来捕获退化信息，但它们在某些方面存在显著差异。首先，正对和负对的定义不同。事实上，对比学习的成功很大程度上依赖于正对和负对的构建，因此这是许多研究的重点。其次，任务不同。简而言之，DASR 是专门为图像超分辨率设计的，而 AirNet 则被提出以一体化方式处理多重退化。第三，尽管任务不同，DASR 需要指定图像超分辨率尺度，而 AirNet 不需要任何退化参数。

The Proposed Method

在本节中，我们详细阐述了所提出的方法，该方法由基于对比的退化编码器（CBDE，fC (·)）和退化引导恢复网络（DGRN，fD(·)）组成，如图 2 所示。
- 所提出的 AirNet 的架构。（a）一体化图像恢复网络 (AirNet)；（b）基于对比的退化编码器 (CBDE)；（c）退化引导组 (DGG)；（d）退化引导模块 (DGM)。
对于给定的退化图像 x，AirNet 首先将其输入 f C ( ⋅ ) f_C (·) fC(⋅) 以学习潜在退化表示 z = f C ( x ) z = f_C (x) z=fC(x)。然后，x 和 z 进一步通过 $f_D(·) $ 以获得恢复的图像 y ′ = f D ( x , z ) y ' = f_D(x, z) y′=fD(x,z)。在不失普遍性的前提下，我们在本文中以三种流行的退化作为展示，即噪声、雾霾和雨水。接下来，我们将首先介绍整体损失函数，然后详细说明具有相应损失的两个子网络。

The Objective Function

为了消除观察到的图像中的损坏，我们提出以下目标函数：
- L = L R e c + L c l L = L_{Rec} + L_{cl} L=LRec+Lcl
- 其中 L R e c L_{Rec} LRec 是真实图像 y 与恢复的干净图像 y' 之间的重建损失。第二个损失 L c l L_{cl} Lcl 是 CBDE 的对比损失。
对于给定的退化图像 x， L R e c L_{Rec} LRec 旨在通过 AirNet 最小化 y 与恢复的干净图像 $y'= f(x) $之间的 L1 距离。从数学上讲，
- L R e c = 1 T ∑ i = 1 T ∣ f ( x i ) − y i ∣ L_{Rec} = \frac1 T\sum_{i=1}^ T |f(x_i) − y_i| LRec=T1i=1∑T∣f(xi)−yi∣
- 其中 T 是 x 像素的数量，i 是像素的索引。
与 L R e c L_{Rec} LRec 不同的是， L c l L_{cl} Lcl 是 CBDE 的特定损失，旨在学习不同退化的表示，同时保留它们可能存在的差异。更多细节将在第 3.2 节中介绍。

Contrastive-Based Degradation Encoder

基于对比的退化编码器旨在从输入 x 中提取潜在退化表示 z。为了使 AirNet 能够解决多个非特定退化问题，z 需要具备以下两个属性。首先，z 应该适应不同的退化。换句话说，对于具有不同退化的输入，即使图像内容相同，相应的 z 也应该不同 。为此，我们利用对比学习来学习 z，通过最大化具有相同退化的两个输入（即正样本）的一致性，同时最小化不同退化（即负样本）之间的一致性。具体来说，对于退化表示 q ， k + 和 k i − q，k^+ 和 k^−_i q，k+和ki− 分别是相应的正样本和负样本。然后，Lcl 可以重新表述为，
- L c l = − l o g e x p ( q ⋅ k + / τ ) ∑ i = 0 K e x p ( q ⋅ k i − / τ ) , ( 3 ) L_{cl} = − log\frac {exp(q · k ^+/τ ) }{\sum ^K _{i=0} exp(q · k^ − _i /τ )},(3) Lcl=−log∑i=0Kexp(q⋅ki−/τ)exp(q⋅k+/τ),(3)
- 其中 τ 是温度超参数，K 表示负样本的数量。
具体来说，对于给定的输入 x，我们从 x 中随机裁剪两个块，分别命名为 x q 和 x k + x_q 和 x_{k+} xq和xk+ 。由于同一图像中的退化应该是一致的，我们将 x q 和 x k + x_q 和 x_{k+} xq和xk+ 视为正对。相反，其他图像中的块被视为相对于 x q 的负 x k − x_q 的负 x_{k−} xq的负xk−。对于获得的对，我们将它们传递给 CBDE 并得到相应的中间表示 v q 、 v k + 和 v k − v_q、v_{k+} 和 v_{k−} vq、vk+和vk−，然后进一步输入到两层 MLP 中得到 q、k + 和 k −。为了学习一个保留不同退化区分度的退化空间，使用公式3 。
由于我们基于对比学习的解决方案，学习到的退化表示具有以下优势。确切地说，它不像现有方法那样依赖于明确定义损坏图像和干净图像之间关系的数学模型。因此，它避免了对这种先验知识的了解，其性能与确切定义无关。特别是，当由于混合多重退化或退化来自自然（例如雨和雾霾）而导致关系始终未知或不精确时，我们的方法更具竞争力。另一方面，我们的方法将不同的退化统一到同一个子空间中，同时保留它们的差异。相比之下，现有的单/多退化方法从不同的子空间学习不同退化的表示，从而失去了退化的可比性和关系。例如，与雾霾损坏相比，损坏率为 0.1 和 0.2 的高斯噪声在潜在空间中应该很接近。显然，我们的对比退化表示可以拥有这样的属性，这对于处理具有多重退化的数据至关重要。
其次，z 应尽可能保留空间结构以利于图像恢复。为此，我们采用 CBDE 第一层而不是最后一层的输出作为 z。换句话说，z 是张量而不是向量，因此可以保留上下文信息。此外，由于 z 的维度与输入和中间层的输出相同，因此可以灵活地与其他特征连接，并与现有的神经网络（如 DCN 和 SFT ）兼容。

Degradation-Guided Restoration Network

利用 CBDE 学习到的 z，使用 DGRN 从具有未知退化的输入中恢复干净的图像。如上图所示，DGRN 由五个退化引导组 (DGG) 构成，每个组又由五个退化引导块 (DGB) 组成。在每个 DGB 中，采用两个退化引导模块 (DGM) 在 z 的引导下恢复干净的图像。如上所述，DGM 是 DGRN 的基本模块，它由可变形卷积 (DCN) 层和空间特征变换 (SFT) 层组成。数学上地
- F D G M m ; b ; g = Φ D G M m ; b ; g ( F m − 1 ; b ; g , z ) = Φ D C N m ; b ; g ( F m − 1 ; b ; g ∣ z ) + Φ S F T m ; b ; g ( F m − 1 ; b ; g ∣ z ) ; F^{m;b;g}{DGM} = Φ^{m;b;g}{DGM}(F^{m−1;b;g}, z) = Φ^{m;b;g}{DCN} (F^{m−1;b;g}|z) + Φ^{m;b;g} {SF T} (F^{m−1;b;g}|_z); FDGMm;b;g=ΦDGMm;b;g(Fm−1;b;g,z)=ΦDCNm;b;g(Fm−1;b;g∣z)+ΦSFTm;b;g(Fm−1;b;g∣z);
- 其中 Φ D G M m ; b ; g Φ^{m;b;g}{DGM} ΦDGMm;b;g 是相对于第 g 个 DGG 的第 b 个 DGB 的第 m 个 DGM， F m − 1 ; b ; g F^{m−1;b;g} Fm−1;b;g 表示相对于第 g 个 DGG 的第 b 个 DGB 的第 (m−1) 个 DGM 的输出。 Φ D C N Φ{DCN} ΦDCN 和 Φ S F T Φ_{SF T} ΦSFT 分别是 DCN 层和 SFT 层。
DGM 的设计旨在实现以下两个目标。一方面，由于不同的退化应该有不同的感受野，因此人们强烈期望模型能够自适应不同的退化。为此，DGM 采用了可变形卷积（DCN），它可以根据调制偏移量和掩模动态调整感受野。具体而言，给定 K 个采样位置的可变形卷积核，让 w k w_k wk 和 p k ∈ { ( − 1 ; − 1 ) ; ( − 1 ; 0 ) ; ⋅ ⋅ ⋅ ; ( 1 ; 1 ) } p_k \in \{(-1; -1);(-1; 0); · · · ;(1; 1)\} pk∈{(−1;−1);(−1;0);⋅⋅⋅;(1;1)} 表示第 k 个位置的权重和预定义偏移量，则 DGM 中使用的 DCN 层定义为：
- Φ D C N m ; b ; g ( F m − 1 ; b ; g ∣ z ) = ∑ k = 1 K w k ⋅ F m − 1 ; b ; g ( p + p k + ∆ p k ) ⋅ ∆ m k ; Φ^{m;b;g}_{DCN} (F^{m−1;b;g}|z) =\sum^K _{k=1} w_k·F^{m−1;b;g}(p+p_k+∆p_k)·∆m_k; ΦDCNm;b;g(Fm−1;b;g∣z)=k=1∑Kwk⋅Fm−1;b;g(p+pk+∆pk)⋅∆mk;
- 其中 F m − 1 ; b ; g ( p ) F^{m−1;b;g}(p) Fm−1;b;g(p) 表示特征图 F m − 1 ; b ; g F^{m−1;b;g} Fm−1;b;g 中位置 p 处的特征。 ∆ p k ∆p_k ∆pk 和 ∆ m k ∆m_k ∆mk 分别是位置 k 的可学习偏移量和调制标量。在我们的实现中，AirNet 使用卷积层 conv(·) 学习 ∆ p k ∆p_k ∆pk 和 ∆ m k ∆m_k ∆mk，其输入是 F m − 1 ; b ; g F^{m−1;b;g} Fm−1;b;g 和 z 的串联。即，
- ( ∆ p k ; ∆ m k ) = c o n v ( c o n c a t ( F m − 1 ; b ; g , z ) ) ; (∆p_k; ∆m_k) = conv(concat(F^{m−1;b;g}, z)); (∆pk;∆mk)=conv(concat(Fm−1;b;g,z));
- 其中 concat(·) 是连接运算符。
另一方面，由于不同的退化图像具有不同的潜在分布，所提出的模型有望缩小分布差距，从而实现更强的多退化恢复能力。为此，DGM 采用 SFT 作为组件，根据 z 调整 F 的分布，即
- F S F T m ; b ; g = Φ S F T m ; b ; g ( F m − 1 ; b ; g ∣ z ) : F^{m;b;g}{SF T} = Φ^{m;b;g}{SF T}(F^{m−1;b;g}|z): FSFTm;b;g=ΦSFTm;b;g(Fm−1;b;g∣z):
- 具体来说，SFT 层旨在学习一个映射函数 M，该函数输出给定 z 的调制参数（γ 和 β）。然后，SFT 通过缩放和移动特征 F m − 1 ; b ; g F^{m−1;b;g} Fm−1;b;g 来进行仿射变换（γ 和 β）。从数学上讲，
- F S F T m ; b ; g = Φ S F T m ; b ; g ( F m − 1 ; b ; g ∣ γ , β ) = γ ⊙ F m − 1 ; b ; g + β ; F^{m;b;g}{SFT} = Φ^{m;b;g}{SFT} (F^{m−1;b;g}|γ, β) = γ\odot F^{m−1;b;g}+β; FSFTm;b;g=ΦSFTm;b;g(Fm−1;b;g∣γ,β)=γ⊙Fm−1;b;g+β;
- 其中 ⊙ \odot ⊙ 表示元素乘法，且 (γ; β) = M(z)。在我们的实验中，我们使用两个卷积层来实现 M。

Experiments

在本节中，我们将通过与 17 个基准进行比较，在四个广泛使用的数据集上评估所提出的方法。接下来，我们将首先介绍实验设置，然后展示基准上的定性和定量结果。最后，我们将进行一些消融研究来验证我们方法的有效性。

Experimental Settings

在本节中，我们介绍所使用的数据集、基线、评估指标和实施细节。
数据集：在实验中，我们使用以下六个数据集进行评估，即 BSD400、BSD68 、WED 和 Urban100 用于去噪；Rain100L 用于去雨；RESIDE 用于去雾。具体来说，BSD400 包含 400 张干净的自然图像，BSD68 包含 68 张自然图像。WED 包含从互联网收集的 4,744 张自然图像，Urban100 有 100 张干净图像。对于图像去噪，我们使用 BSD400 和 WED 的组合作为训练集，使用 BSD68 和 Urban100 的组合作为测试集 。通过向干净图像手动添加高斯白噪声来生成噪声图像，三个损坏级别为 σ = 15；25； 50. 对于图像去雨，我们在 Rain100L 上进行实验，该数据集包含 200 个雨天-干净训练对和 100 个测试图像对。对于图像去雾，我们在 RESIDE 数据集上进行实验，该数据集包含分别用于训练和测试的户外训练集 (OTS) 和合成目标测试集 (SOTS)。简而言之，OTS 包含 72,135 个户外雾天-干净图像对，SOTS 包含 500 个户外雾天-干净图像对。
基线：为了进行全面比较，我们将我们的方法与五种去噪方法、五种去雨方法、五种去雾方法、一种图像恢复方法和一种 IRMD 方法进行了比较。具体来说，去噪基线包含 CBM3D 、DnCNN 、IRCNN 、FFDNet 和 BRDNet 。去雨基线是 DIDMDN 、UMRL 、SIRR 、MSPFN 和 LPNet 。去雾基线是 DehazeNet 、MSCNN 、AOD-Net 、EPDN 和 FDGAN 。图像恢复基线是 MPRNet 。IRMD 基线是解耦学习 (DL) 。
为了全面展示我们方法的有效性，我们研究了两种不同的设置，即逐个训练 AirNet 以指定退化程度（OBO）和以一体化方式训练 AirNet 以所有退化程度（AIO）进行训练。换句话说，AIO 下的 AirNet 是在包含三种损坏（即噪声、雨水和雾霾）且退化程度不同（即 σ = 15；25；50）的所有数据集上训练的模型。
评估指标：使用两个流行的指标进行定量比较，即峰值信噪比（PSNR）和结构相似度（SSIM）。这些指标的值越高，表示方法的性能越好。
训练细节：我们在 NVIDIA GeForce RTX 2080Ti GPU 上使用 PyTorch 进行实验。为了优化 AirNet，我们使用默认 { β i } i = 1 2 \{β_i\}^ 2_ {i=1} {βi}i=12 的 ADAM 优化器，并将最大迭代次数设置为 1,500。为了热身，我们首先通过优化 L c l L_{cl} Lcl 训练 CBDE 100 次迭代。然后，我们使用 L 对整个网络进行 1,400 次迭代训练。学习率初始化为 0.001，然后在 60 个 epoch 后降至 0.0001。之后，每 125 个 epoch 后学习率降低一半。在实验中，我们使用 400×N 的批处理大小和 128 的补丁大小训练我们的模型，其中 N 是退化类型的数量。

Comparisons on Single Degradation

在本节中，我们展示了三个独立的图像恢复任务（即去噪、去雨和去雾）的定量和定性结果。
去噪：上表1 报告了在逐一设置下，在 BSD68 和 Urban100 上与五种去噪方法进行比较的结果。从结果中可以看出，AirNet 在几乎所有测试中都取得了最佳成绩。除了在定量评估中占据主导地位外，AirNet 在定性比较中也表现出优势，如下图所示。由于篇幅限制，我们在补充材料中留下了更多结果。
- 与 BSD68 数据库上的 SOTA 去噪方法进行比较。一些区域以彩色矩形突出显示，建议放大以获得更好的可视化和比较。
去雨：从下表和下图可以看出，AirNet 的表现也明显优于所有去雨基线。例如，AirNet 在 OBO 设置下的 PSNR 和 SSIM 分别比最佳方法高出 1.4 和 0.0074。
- Rain100L 数据集上图像去雨的定量结果。最佳结果以粗体显示。
- 在 Rain100L 数据库上对 SOTA 去雨方法进行比较。一些区域以彩色矩形突出显示，建议放大以获得更好的可视化和比较效果。
去雾如下表和下图所示，AirNet 在 PSNR 上略优于最佳基线。更具体地说，AirNet 在 PSNR 上比 FDGAN 高 0.03。然而，视觉结果表明 AirNet 可以恢复更多对人类有利的细节
- SOTS 数据集上图像去雾的定量结果。最佳结果以粗体显示。
- 在 SOTS 数据库上对 SOTA 去雾方法的比较。一些区域以彩色矩形突出显示，建议放大以获得更好的可视化和比较效果。

Comparisons on Multiple Degradations

AirNet 最吸引人的一点是它能够在一个一体化框架中处理不同的未知退化。在本节中，我们将进行实验来验证 AirNet 在这种设置下的有效性。为此，我们选择五种 IRSD 方法（即 BRDNet 、LPNet 、FDGAN 和 MPRNet ）和一种 IRMD 方法（即 DL ）作为基线。为了进行公平和广泛的比较，我们使用上述两种设置（即 One-By-One 和 All-In-One）重新训练这些方法。如下表所示，可以观察到 AirNet 在大多数情况下都优于所有基线。应该指出的是，虽然 DL 也可以处理多种退化，但它需要知道损坏类型和级别，以便可以指定网络的正确头部和尾部 。
- 在三个具有挑战性的数据集上的性能比较。最佳结果以粗体显示。

Results on Combined Degradations

在本节中，我们使用多种降级的不同组合来训练 AirNet，以分析损坏的数据集对性能的影响。如下表所示，降级越多，去噪就越困难，而去雨和去雾任务则不能得出相同的结论。有趣的是，去雨将有助于去噪，而去雾得益于所有降级的组合。未来有望进行更多的实证研究和理论分析。
- 退化组合消融研究。表中，"√"表示有退化的 AirNet，"-"表示不可用结果，最佳结果以粗体显示。

Results on Spatially Variant Degradation

在本节中，我们进行实验来证明 AirNet 对空间变化退化的有效性，即同一图像的不同区域具有不同的损坏程度。为此，我们合成了一个具有空间变化噪声的 BSD68 退化版本。具体来说，我们将每张干净图像分成四个区域，其中分别添加了 σ ∈ { 0 ； 15 ； 25 ； 50 } σ \in \{0；15；25；50\} σ∈{0；15；25；50} 的高斯噪声。下表表明，在这样的评估协议下，AirNet 在恢复潜在的干净图像方面也是有效的。
- 空间变异退化的定量结果。最佳结果以粗体显示。

Ablation Study

为了证明我们网络结构的有效性，我们通过删除 DCN 层和 SFT 层之一对 BSD68 进行了消融研究。从下表可以看出，DCN 层和 SFT 层对于提高 AirNet 的性能都很重要。
- 在 BSD68 和 Urban100 上进行的消融研究。最佳结果以粗体显示。

Conclusion

本文提出了一种不受损坏类型和损坏程度先验影响的一体化图像恢复网络 (AirNet)。同时，该方法是一种从不同损坏中恢复图像的一体化解决方案，可适用于先验难以预知或退化可能随时间和空间变化的各种实际场景。大量实验结果表明，AirNet 在定性和定量比较中都具有优越性。

Shortcomings and Broader Impact

尽管 AirNet 在三个图像恢复任务及其组合中通过实验表现出优越性，但它对其他损坏（如模糊和下雪）的表现如何尚不清楚。此外，还值得进一步探讨为什么不同的组合退化会导致单个任务的不同结果，如第 4.5 节所示。从更广阔的角度来看，尽管 AirNet 可以适应不同的损坏并避免在不同退化上采用同一算法的多个模型，但它仍然需要大量资源来优化该方法，从而导致碳排放并间接导致气候变暖。

AIRNet提出了一种较为简易的pipeline，以单一网络结构应对多种任务需求（不同类型，不同程度）。但在效果上看，ALL-In-One是不如One-By-One的，且本文方法的亮点是batch内选择patch进行对比学习。在与sota对比上，仅是Denoise任务精度占优，在Derain与Dehaze任务上，效果不如One-By-One的MPRNet方法。输入图像x交由CBDE提取到嵌入空间z，z与x输入到DGRN模块的DGG block中逐步优化，最终输出预测结果。模型代码在net\model.py AIRNet模型使用与代码分析（All-In-One Image Restoration Network）-CSDN博客

python 复制代码

from torch import nn
from net.encoder import CBDE
from net.DGRN import DGRN
class AirNet(nn.Module):
    def __init__(self, opt):
        super(AirNet, self).__init__()
        # Encoder
        self.E = CBDE(opt)  #编码特征值
        # Restorer
        self.R = DGRN(opt) #特征解码
    def forward(self, x_query, x_key):
        if self.training:
            fea, logits, labels, inter = self.E(x_query, x_key)
            restored = self.R(x_query, inter)
            return restored, logits, labels
        else:
            fea, inter = self.E(x_query, x_query)
            restored = self.R(x_query, inter)
            return restored

CBDE模块的功能是在模块内进行对比学习，核心是MoCo

python 复制代码

class CBDE(nn.Module):
    def __init__(self, opt):
        super(CBDE, self).__init__()
        dim = 256
        # Encoder
        self.E = MoCo(base_encoder=ResEncoder, dim=dim, K=opt.batch_size * dim)
    def forward(self, x_query, x_key):
        if self.training:
            # degradation-aware represenetion learning
            fea, logits, labels, inter = self.E(x_query, x_key)
            return fea, logits, labels, inter
        else:
            # degradation-aware represenetion learning
            fea, inter = self.E(x_query, x_query)
            return fea, inter

在AIRNet中的CBDE模块里的MoCo模块的关键代码如下，其在内部自行完成了正负样本的分配，最终输出logits, labels用于计算对比损失的loss。但其所优化的模块实际上是ResEncoder。MoCo模块只是在训练阶段起作用，在推理阶段是不起作用的。

python 复制代码

class MoCo(nn.Module):
    def forward(self, im_q, im_k):
        """
        Input:
            im_q: a batch of query images
            im_k: a batch of key images
        Output:
            logits, targets
        """
        if self.training:
            # compute query features
            embedding, q, inter = self.encoder_q(im_q)  # queries: NxC
            q = nn.functional.normalize(q, dim=1)
            # compute key features
            with torch.no_grad():  # no gradient to keys
                self._momentum_update_key_encoder()  # update the key encoder
                _, k, _ = self.encoder_k(im_k)  # keys: NxC
                k = nn.functional.normalize(k, dim=1)
            # compute logits
            # Einstein sum is more intuitive
            # positive logits: Nx1
            l_pos = torch.einsum('nc,nc->n', [q, k]).unsqueeze(-1)
            # negative logits: NxK
            l_neg = torch.einsum('nc,ck->nk', [q, self.queue.clone().detach()])
            # logits: Nx(1+K)
            logits = torch.cat([l_pos, l_neg], dim=1)
            # apply temperature
            logits /= self.T
            # labels: positive key indicators
            labels = torch.zeros(logits.shape[0], dtype=torch.long).cuda()
            # dequeue and enqueue
            self._dequeue_and_enqueue(k)
            return embedding, logits, labels, inter
        else:
            embedding, _, inter = self.encoder_q(im_q)
            return embedding, inter

DGRN模块的实现代码如下所示，可以看到核心是DGG模块，其不断迭代优化输入图像。

python 复制代码

class DGRN(nn.Module):
    def __init__(self, opt, conv=default_conv):
        super(DGRN, self).__init__()
        self.n_groups = 5
        n_blocks = 5
        n_feats = 64
        kernel_size = 3
        # head module
        modules_head = [conv(3, n_feats, kernel_size)]
        self.head = nn.Sequential(*modules_head)
        # body
        modules_body = [
            DGG(default_conv, n_feats, kernel_size, n_blocks) \
            for _ in range(self.n_groups)
        ]
        modules_body.append(conv(n_feats, n_feats, kernel_size))
        self.body = nn.Sequential(*modules_body)
        # tail
        modules_tail = [conv(n_feats, 3, kernel_size)]
        self.tail = nn.Sequential(*modules_tail)
    def forward(self, x, inter):
        # head
        x = self.head(x)
        # body
        res = x
        for i in range(self.n_groups):
            res = self.body[i](res, inter)
        res = self.body[-1](res)
        res = res + x
        # tail
        x = self.tail(res)
        return x

DGG模块内嵌DGB模块，DGB模块内嵌DGM模块，DGM模块内嵌SFT_layer模块与DCN_layer（可变性卷积）。

AIRNet中提到的loss包含两部分，其中Lrec是L1 loss，Lcl是Moco模块实现的对比损失。AIRNet的loss实现代码在train.py中，CE loss是针对CBDE（Moco模块）的输出进行计算，l1 loss是针对修复图像与清晰图片。AIRNet首先是训练CBDE模块，最后才训练CBDE模块+DGRN模块。

python 复制代码

    # Network Construction
    net = AirNet(opt).cuda()
    net.train()
    # Optimizer and Loss
    optimizer = optim.Adam(net.parameters(), lr=opt.lr)
    CE = nn.CrossEntropyLoss().cuda()
    l1 = nn.L1Loss().cuda()
    # Start training
    print('Start training...')
    for epoch in range(opt.epochs):
        for ([clean_name, de_id], degrad_patch_1, degrad_patch_2, clean_patch_1, clean_patch_2) in tqdm(trainloader):
            degrad_patch_1, degrad_patch_2 = degrad_patch_1.cuda(), degrad_patch_2.cuda()
            clean_patch_1, clean_patch_2 = clean_patch_1.cuda(), clean_patch_2.cuda()
            optimizer.zero_grad()
            if epoch < opt.epochs_encoder:
                _, output, target, _ = net.E(x_query=degrad_patch_1, x_key=degrad_patch_2)
                contrast_loss = CE(output, target)
                loss = contrast_loss
            else:
                restored, output, target = net(x_query=degrad_patch_1, x_key=degrad_patch_2)
                contrast_loss = CE(output, target)
                l1_loss = l1(restored, clean_patch_1)
                loss = l1_loss + 0.1 * contrast_loss
            # backward
            loss.backward()
            optimizer.step()

TrainDataset的实现代码在utils\dataset_utils.py中，首先找到__getitem__函数进行分析。以下代码为关键部分，删除了大部分在逻辑上重复的部分。TrainDataset一共支持5种数据类型，'denoise_15': 0, 'denoise_25': 1, 'denoise_50': 2,是不需要图像对的（在代码里面自动对图像添加噪声）；'derain': 3, 'dehaze': 4是需要图像对进行训练的。

python 复制代码

class TrainDataset(Dataset):
    def __init__(self, args):
        super(TrainDataset, self).__init__()
        self.args = args
        self.rs_ids = []
        self.hazy_ids = []
        self.D = Degradation(args)
        self.de_temp = 0
        self.de_type = self.args.de_type
        self.de_dict = {'denoise_15': 0, 'denoise_25': 1, 'denoise_50': 2, 'derain': 3, 'dehaze': 4}
        self._init_ids()
        self.crop_transform = Compose([
            ToPILImage(),
            RandomCrop(args.patch_size),
        ])
        self.toTensor = ToTensor()
    def __getitem__(self, _):
        de_id = self.de_dict[self.de_type[self.de_temp]]
        if de_id < 3:
            if de_id == 0:
                clean_id = self.s15_ids[self.s15_counter]
                self.s15_counter = (self.s15_counter + 1) % self.num_clean
                if self.s15_counter == 0:
                    random.shuffle(self.s15_ids)
            # clean_id = random.randint(0, len(self.clean_ids) - 1)
            clean_img = crop_img(np.array(Image.open(clean_id).convert('RGB')), base=16)
            clean_patch_1, clean_patch_2 = self.crop_transform(clean_img), self.crop_transform(clean_img)
            clean_patch_1, clean_patch_2 = np.array(clean_patch_1), np.array(clean_patch_2)
            # clean_name = self.clean_ids[clean_id].split("/")[-1].split('.')[0]
            clean_name = clean_id.split("/")[-1].split('.')[0]
            clean_patch_1, clean_patch_2 = random_augmentation(clean_patch_1, clean_patch_2)
            degrad_patch_1, degrad_patch_2 = self.D.degrade(clean_patch_1, clean_patch_2, de_id)
        clean_patch_1, clean_patch_2 = self.toTensor(clean_patch_1), self.toTensor(clean_patch_2)
        degrad_patch_1, degrad_patch_2 = self.toTensor(degrad_patch_1), self.toTensor(degrad_patch_2)
        self.de_temp = (self.de_temp + 1) % len(self.de_type)
        if self.de_temp == 0:
            random.shuffle(self.de_type)
        return [clean_name, de_id], degrad_patch_1, degrad_patch_2, clean_patch_1, clean_patch_2

TrainDataset返回的数据有：degrad_patch_1, degrad_patch_2, clean_patch_1, clean_patch_2。通过以下代码可以看出 clean_patch_1, clean_patch_2是来自于同一个图片，然后基于crop_transform变化，变成了2个对象

python 复制代码

            clean_img = crop_img(np.array(Image.open(clean_id).convert('RGB')), base=16)
            clean_patch_1, clean_patch_2 = self.crop_transform(clean_img), self.crop_transform(clean_img)
            # clean_name = self.clean_ids[clean_id].split("/")[-1].split('.')[0]
            clean_name = clean_id.split("/")[-1].split('.')[0]
            clean_patch_1, clean_patch_2 = random_augmentation(clean_patch_1, clean_patch_2)
 crop_transform = Compose([
            ToPILImage(),
            RandomCrop(args.patch_size),
        ])

random_augmentation的实现代码如下，可以看到只是随机对图像进行翻转或旋转，其目的是尽可能使随机crop得到clean_patch_1, clean_patch_2差异更大，避免裁剪出高度相似的patch。

python 复制代码

def random_augmentation(*args):
    out = []
    flag_aug = random.randint(1, 7)
    for data in args:
        out.append(data_augmentation(data, flag_aug).copy())
    return out
def data_augmentation(image, mode):
    if mode == 0:
        # original
        out = image.numpy()
    elif mode == 1:
        # flip up and down
        out = np.flipud(image)
    elif mode == 2:
        # rotate counterwise 90 degree
        out = np.rot90(image)
    elif mode == 3:
        # rotate 90 degree and flip up and down
        out = np.rot90(image)
        out = np.flipud(out)
    elif mode == 4:
        # rotate 180 degree
        out = np.rot90(image, k=2)
    elif mode == 5:
        # rotate 180 degree and flip
        out = np.rot90(image, k=2)
        out = np.flipud(out)
    elif mode == 6:
        # rotate 270 degree
        out = np.rot90(image, k=3)
    elif mode == 7:
        # rotate 270 degree and flip
        out = np.rot90(image, k=3)
        out = np.flipud(out)
    else:
        raise Exception('Invalid choice of image transformation')
    return out

degrad_patch来自于clean_patch，可以看到是通过D.degrade进行转换的。D.degrade相关的代码如下，可以看到只是对图像添加噪声。

python 复制代码

class Degradation(object):
    def __init__(self, args):
        super(Degradation, self).__init__()
        self.args = args
        self.toTensor = ToTensor()
        self.crop_transform = Compose([
            ToPILImage(),
            RandomCrop(args.patch_size),
        ])
    def _add_gaussian_noise(self, clean_patch, sigma):
        # noise = torch.randn(*(clean_patch.shape))
        # clean_patch = self.toTensor(clean_patch)
        noise = np.random.randn(*clean_patch.shape)
        noisy_patch = np.clip(clean_patch + noise * sigma, 0, 255).astype(np.uint8)
        # noisy_patch = torch.clamp(clean_patch + noise * sigma, 0, 255).type(torch.int32)
        return noisy_patch, clean_patch
    def _degrade_by_type(self, clean_patch, degrade_type):
        if degrade_type == 0:
            # denoise sigma=15
            degraded_patch, clean_patch = self._add_gaussian_noise(clean_patch, sigma=15)
        elif degrade_type == 1:
            # denoise sigma=25
            degraded_patch, clean_patch = self._add_gaussian_noise(clean_patch, sigma=25)
        elif degrade_type == 2:
            # denoise sigma=50
            degraded_patch, clean_patch = self._add_gaussian_noise(clean_patch, sigma=50)
        return degraded_patch, clean_patch
    def degrade(self, clean_patch_1, clean_patch_2, degrade_type=None):
        if degrade_type == None:
            degrade_type = random.randint(0, 3)
        else:
            degrade_type = degrade_type
        degrad_patch_1, _ = self._degrade_by_type(clean_patch_1, degrade_type)
        degrad_patch_2, _ = self._degrade_by_type(clean_patch_2, degrade_type)
        return degrad_patch_1, degrad_patch_2

项目中默认包含了All.pth，模型放到 ckpt/ 目录下。打开demo.py，将 subprocess.check_output(['mkdir', '-p', opt.output_path]) 替换为os.makedirs(opt.output_path,exist_ok=True)，避免在window上报错。

【读点论文】All-In-One Image Restoration for Unknown Corruption用对比学习统一方法实现多种噪声图片的有效处理

All-In-One Image Restoration for Unknown Corruption

Abstract

Introduction

Related Works

Image Restoration

Contrastive Learning

The Proposed Method

The Objective Function

Contrastive-Based Degradation Encoder

Degradation-Guided Restoration Network

Experiments

Experimental Settings

Comparisons on Single Degradation

Comparisons on Multiple Degradations

Results on Combined Degradations

Results on Spatially Variant Degradation

Ablation Study

Conclusion

Shortcomings and Broader Impact