[CR]厚云填补_MemoryNet-优快云博客

Memory augment is All You Need for image restoration

Abstract

图像恢复数据低级视觉任务，大部分CNN方法被设计为黑盒，缺乏透明度和内部美观性，虽然也提出了一些将传统优化算法与DNN相结合的方法，但是都存在一定的局限性。本文提出一种三粒度记忆层与对比学习的MemoryNet。具体来说，将样本分为正、负、实三个样本进行对比学习，记忆层能够保留图像的深层特征，对比学习使学习到的特征收敛到平衡。在Derain/Deshadow/Deblur任务上的实验表明，这些方法对提升恢复的性能有效。此外，本文提出的模型在三个不同的图像质量退化类型的数据集上均取得了显著的PSNR、SSIM增益，有力证明了恢复的图像在感知上是真实的。

1 Introduction

图像恢复是一项低级视觉任务，指的是针对退化图像的恢复。常见的退化类型包括额外的噪声、模糊等。近年来，计算机视觉的快速发展使得处理越来越多的退化任务成为了可能，包括超分辨率、单幅图像去雾、图像去阴影和图像去雨。

图像恢复存在无数个可行的解决方案。图像先验用于将解决方案空间限制为有效/自然图像。然而，设计这样的先验是一项具有挑战性的任务，而且往往无法推广。为了解决这个问题，最近最先进的方法采用卷积神经网络CNN通过捕获自然图像来隐式学习更一般的先验大规模数据统计。基于CNN的方法的模型设计是其性能优于其他方法的主要原因。随着生成对抗网络的成功，一些用于扭曲图像恢复的生成模型被引入，如图像修复、图像阴影去除、图像去雨和图像去云。以图像去阴影为例，这些方法将失真恢复问题公式化为寻找合适的扭曲和预测密集网络，从而获得了最先进的结果。然而，现有的努力主要集中在模型结构上，一个被广泛忽视的机会是记忆学习方法。

我们认为图像复原研究存在两个问题：

第一，当模型收敛时，一个限制问题暴露出来：性能无法得到明显提升，此时模型处于过拟合状态。出乎意料的是，这种限制并不是通过简单地增加更多的层数就能解决的。事实上，当给定训练数据集时，它的分布是客观固定的。因此，如何在保持原有底层网络不变的情况下，让网络学到更多的潜在信息，是一个非常难以解决的问题。
第二，阴影去除不同于将一整块区域涂成白色的图像修复，它必须对阴影的区域进行一定的操作，同时保留原始图像的特征，这很难实现，因为你需要考虑阴影区域下的真实图像信息。

因此，我们设计了MemoryNet，具体分为记忆增强和对比学习。一方面，为了解决上述第一个问题，本文引入了一个新颖的记忆增强模块MA（memory augment），它额外建模了一个可学习的潜在属性变量，用于记忆全局范围内代表性结构的原型模式，该范围通常涵盖了多样化的身份样本。通过增加这种记忆可能性的预测，模型对未见类的置信度也得到了提高。MA还传达了从先前查看的样本中学习到的额外域级低频信息，用于协作决策，避免了类似GAP的语义抽象。更详细地说，我们通过读取记忆将特征重新收集到预定义的由粗到细的原型索引中，以进行进一步的相似性测量。与在两幅图像中搜索相关的判别区域不同，这种原型对齐是轻量级的，类似于多级哈希。记忆增强网络很好地增强了图像去阴影的应用效果，证明了其在消融研究中的实用性。另一方面，为了解决提到的第二个问题，我们设计了一个弱监督学习的对比度学习网络，如图2所示。我们将图像去阴影任务定义为三分类问题，分别对应正样本（干净样本）、标准样本（去阴影样本）、负样本（阴影样本），以全局特征作为锚点，使模型能够区分全局特征和局部特征是否来自同一张图像。因此，对比学习通过对比学习上下边界，将锚定图像约束为封闭图像，有助于恢复网络接近正图像，同时避开负图像。

我们对图像去阴影/图像去雨/图像去模糊三个任务进行了大量实验，均取得了令人满意的结果。本文的结果总结如图1所示。这项工作的主要贡献是：

图1 MemoryNet的图像阴影去除和图像去雨效果。

本文设计了一种新型端到端网络MemoryNet，用于图像恢复，可生成上下文丰富且空间精确的输出。
本文设计了一种新型记忆增强层，模拟了一个可学习的潜在属性变量，以记住全局代表性的结构原型模式。
对三种典型的图像恢复任务（即合成图像去阴影、真实图像去雨和图像去模糊）进行了广泛的实验，结果表明，我们提出的MemoryNet在保持有吸引力的计算复杂度的同时取得了很好的效果。此外，我们还提供了详细的消融研究、定性结果和泛化测试。

2 相关工作

2.1 图像阴影去除

监督阴影去除

本文提出一个数据集SRD，最早的弱监督DeshadowNet最大的特点是全自动端到端实现阴影去除（一个新的阴影去除数据集）。为了共同利用彼此进步的优势，ST-CGAN提供了一种多任务视角，不同于所有其他现有方法，因为其以端到端的方式一起学习检测和消除。通过估计线性变换函数，SID和DSC创建深度网络来照亮阴影区域，用于阴影（半阴影）。RSI-GAN创建了一个广泛的框架，使用多个GAN挖掘照明和残差数据以去除阴影。DHAN使用双层聚合网络来消除边界伪影。Auto-exposure旨在挖掘阴影和非阴影区域的上下文信息。

弱监督阴影去除

根据 Mask-shadowGAN的说法，先前的深度学习方法用于解决阴影去除难题，这些方法是监督配对数据。然而，获取匹配的数据集可能是一个挑战。根据LG-shadow的说法，在实际操作中，CNN 训练更倾向于非配对数据，因为它更容易训练。SpA-Former 建议使用基于注意力机制的 Transformer 网络，与 Transformer 一起学习阴影空间注意力图。

无监督学习阴影去除

G2R 利用了阴影照片经常同时具有阴影和非阴影区域这一事实。通过使用此技术，可以裁剪一组阴影和非阴影块，以提供未配对的数据用于网络训练，从而提供三个子网络模块的可能性：阴影生成、阴影去除和阴影细化。阴影去除任务由 TC-GAN以无监督方式执行。将基于循环一致性的双向映射方法与基于 GAN 的无监督阴影去除方法进行了对比。

2.2 图像去雨

监督去雨

SPANet在生成器和鉴别器中引入了注意力机制，它也会在多个时间步骤上生成注意力图，并确定原始图中网络需要注意的区域：雨点及其周围环境。NEDNet的编码器-解码器架构引用了非局部增强，成功去除了不同密度的雨水，同时完美地保留了图像细节。Yang提出了一种新的流程：要完成去雨，首先检测雨的位置，然后估计雨线，最后去除背景层。Ren提供了一个起始点：该模型由六个阶段组成，每个阶段分为两个模型，并以初始雨水图和其前一步生成的去雨图的拼接作为输入。

半监督去雨

Wei和Huang提出了半监督模型，可以记录不同的降雨退化原型，并通过自监督学习对其进行更新。

无监督去雨

Guo提出了无监督注意力机制引导的雨水提取模型E。对有雨和无雨地图的空间域都使用了注意力机制，并使用带有两个约束分支的CycleGAN循环结构来去雨。

2.3 记忆模块的发展

Dong首次将Memory模块引入到异常检测任务中，认为Encoder可以看作是query generator；Decoder的输入是Memory模块生成的大小为（H，W，C）的新特征图，用来重建生成的图像，可以理解为重建Encoder的特征图，生成的新特征图包含更多的信息量的正常帧信息，以便在测试时更好地区分正常帧和异常帧，达到视频异常检测的目的。MMOS是一个去雨任务，认为最重要的环节是中间的Memory模块，用于对雨水退化过程建模/存储不同的雨水模式。具体来说，z(x)相当于一个查询，用于在记忆中找到最相关的项，并将其与软注意力相结合，作为去雨任务的指导。Chen是将记忆模块移植到行人重新识别任务。

3 网络结构

图2 MemoryNet的结构，分为记忆增强网络和对比学习。

MemoryNet 的灵感来自，如图 2 所示。该网络主要由两部分组成。第一部分是记忆增强，由两个编码器和解码器网络以及随后的残差网络组成。本文使用编码器-解码器来学习多尺度上下文信息，而最后阶段则在原始图像分辨率上进行操作以保留精细的空间细节。第二部分是对比学习网络。

我们没有简单地级联多个阶段，而是在每两个阶段之间添加一个监督注意模块。在真实图像的监督下，我们的模块会重新调整前一阶段的特征图，然后将其传递到下一阶段。此外，我们引入了一种跨阶段特征融合机制，其中较早子网络的中间多尺度上下文特征有助于巩固后一子网络的中间特征。

3.1 增强记忆

如何检测异常区域

将模糊称为一种异常模式，因此应该检测哪里异常，并将异常细化为正常状态。在这个考虑下，我们结合了异常检测和完成的代理。这个基线对于Transformer框架是可行的。在这里将阴影图像定义为异常状态，将干净图像定义为正常。遵循传统的异常检测编码器和解码器结构，首先将干净的图像发送到记忆增强的编码器/解码器以记录正常模式。不需要改变原来的框架基本编码器/解码器，只需要在其中插入记忆和监督。在训练这个异常检测阶段之后，编码器和解码器满足重建能力，而记忆记录正常模式。当我们发送一个模糊的图片时，这个模型无法将其恢复到干净的状态。

如何将异常转换为正常

基本思想类似于上下文编码器/解码器，因为我们将去阴影视为类似于完成的代理任务。在训练阶段，我们已经知道编码器在去阴影方面很弱，因为我们只用干净的图像进行训练。因此，我们引入了一个潜在的上下文回归器来恢复模糊区域。

增强记忆

图3 解码器和编码器的记忆增强结构

记忆增强的结构如图 3 所示。由于 CNN 模型的表征能力强，当输入为不规则帧时，重构误差很小，导致输出错误。通过在编码器和解码器中添加记忆模块解决了这个问题，这使得编码器和解码器能够记录正常的帧属性，并削弱 CNN 的表征能力，以区分正常和异常帧。在本研究中，我们重新制定了基于分类的排名检索的最后一层 CNN 的可解释概率处理。然后，为了减少领域偏差，我们构建了一个分层的记忆调整和对齐模块。

记忆模块包含N个用度量 $M\in \mathbb{R}^{N\times C}$ 记录的原型，具有固定的特征维度C，N是记忆项的数量，是一个根据需要调整的超参数，记忆寻址计算每个查询相对于所有记忆项的权重，然后使用一个基于注意机制的访问记忆的寻址算子，即记忆读取器，然后使用将每幅图像分配给备用原型：

其中 $f_{i}$ 和 $m_{j}$ 分别是来自输入 $f$ 的特征和原型切片原型度量 M。 $w_{ij}$ 是标准化权重，用于测量 $f_{i}$ 和 $m_{j}$ 之间的余弦相似度 $d(-,-)$ 。因此，可以从特征 $f$ 计算出分配的原型 $h$ ，如下所示：

构建了图2所示的记忆增强。MA(Memory augment)由原型M(即部分、实例、语义)组成的层次语义组成，避免过度抽象(avoid over-abstraction)。实例和语义原型总结了先前的低级别原型。因此，尽管在原型的记忆插件中跨越各种语义多样性，但共享M以表示所有通用概念样本。具体而言，我们将原型M定义为 $2\times (P\times I\times S\times N_{c})\times C$ 形状，其中P、I、S分别针对部分、实例、语义级别预先定义的原型编号，并且 $N_{c}$ 是类别编号。在总结语义原型之前，每个部分和实例原型均以两种模式的复制。因此，对于内部模式间隙，我们将各个模式保持在较低级别的代表模式部分和实例原型，然后在语义级别将其共同对齐。如图2所示，每个高级别原型项目都可以总结其较低范围。例如，实例原型子量表 $M_{ins}$ 的 $i_{th}$ 行 $m_{ins,i}$ ，可以看作是加权子段 $M_{part}$

在训练这个异常检测阶段后，编码器和解码器符合重建能力，而记忆记录正常模式。当我们发送模糊的图像时，该模型无法将其恢复清晰。因此，每个阶段的编码输入是：

SFe代表千层特征融合产物。表示在记忆增强的这个阶段，则重建异常特征和正常特征，并使用约束，我们称之为 $L_{recon}$ ，定义为：

3.2 损失函数设计

MemoryNet的损失函数如下所示：

Y表示地面图像，而 $L_{char}$ 是Charbonnier损失：

对于所有实验，恒定的 $^{\circ}A$ 经验设置为

此外， $L_{edge}$ 是边缘损失，定义为：

4 实验

实验细节

我们的CR-MemoryNet是一种端到端的可训练模型，不需要任何预训练。使用Pytorch 1.8.0和NVIDIA GTX 3090 GPU实施的。在本文中，我们选择了三个评估指标，PSNR和SSIM以及RMSE。

阴影去除

本文使用的数据集为ISTD。我们从经验上使用Adam优化器来优化我们的网络。在我们的实验中，我们将第一个动量值设置为0.9，第二动量值为0.999，衰减为 $5\times 10^{-4}$ 。ISTD包括540个测试三元组和1330个训练三联组，分为阴影、阴影掩码和无阴影图片。在训练和测试中，SRD分别包含2680和408对图像。

真实场景雨去除

本文使用DeRainDrop数据集进行训练和测试。提供了861个图像对用于训练，有两个测试数据集（teatA和testB）。TeatA是testB的子集，其中包含58对良好的对齐图像。TestB具有249个图像对，其中一小部分的图像对不完全对齐。

图像去模糊

对于图像去模糊，我们用GoPro的2103个图像对训练了模型，与预定义的模糊核不同，这两个数据集是在涉及real-world degradation的真实场景中产生的，例如camera response functions和human consciousness dynamic blurring。

4.1 与最先进的方法进行对比

阴影去除

将我们的方法与现有方法进行比较，包括Yang[37]，Guo [35]，Gong [36]，Deshad Ownet [23]，STC-Gan [22]，DSC [39]，Mask-Shadowgan [43]，Ris-Gan [40]，DHAN [41]，DHAN [41]，SID [38]，SID [38]，SID [38]，LG-SHADOW [LG-SHADOW [46]，GAST，G2]，G2 [48] 48 [48]，D.48 [48]，D.48 [48]，D.48 [48] Auto-Exp [42]，Spa-Former [50]，Canet [44]。我们在LAB color space中采用均方根误差RMSE、结构相似性指数SSIM、峰值信噪比PSNR作为评价指标。表1分别展示了ISTD数据集上不同阴影去除方法的RMSE、SSIM、PSNR的值。定量比较如图4所示。MemoryNet无论是部分阴影还是非部分阴影还是无阴影区域都在PNSR上取得了最佳性能，RMSE都可以在无阴影区域达到最佳性能，一般都超过了SOTA。SID和G2R可能会错误地处理相对较暗的非阴影区域，从而造成一些误解。事实证明，即使其他网络输入包含阴影掩码，模型也无法充分利用阴影掩码的信息。Auto-Exp和CANet使用Triplet数据集（输入、掩码、标签），在indicators方面表现良好。但是，是否有必要在实际应用中获取阴影的掩码。从应用的角度看，这没有实际价值。本文中的训练只需要成对数据，而测试只需要一个简单的阴影图，这具有很大的实际意义。

表1 ISTD上图像阴影去除的性能比较。（红色标记的代表排名第一，蓝色标记代表排名第二）。

图4 GOPRO数据集上图像去模糊的视觉性能比较。

图像去雨

如表2和图5所示，我们展示了DeRainDrop数据集上TestA和TestB上去雨的PSNR/SSIM的得分方式。将我们的方法与现有的方法进行比较，包括CMFNet、D-DAM、BPP、Maxim、IDT。我们的MemoryNet在testB上获得了最佳的SSIM分数（0.84）和第二最佳的PSNR的得分（25.38db），在testA上获得了最佳SSIM分数（0.904）和最佳PSNR（24.63db）。该图5展示了DeRainDrop中TestB图像的可视化结果，该结果很好地证明了我们的方法有效地去除了雨滴，而恢复的图像在视觉上比其他模型更接近真实图像。为了与MMOS相比，这是一个使用记忆模块的去雨网络，我们对其进行了去雨测试，但是结果表明，MMOS在实际去雨数据集上不可行，我们认为这可能是由于在实际数据计算中，噪声数据的使用与目标网络生成的伪标签进行配对并未成功。

表2 RainDrop（TestA和TestB）上去雨的性能比较。红色标记的代表排名第一，蓝色标记的代表排名第二。

图5 RainDrop数据集上图像的视觉性能比较。

图像去模糊

在表3和图6中，我们展示了去模糊任务的PSNR/SSIM分数方法。我们与包括GAO、DBGAN、MT-RNN、MPRNet和DGUNet在内的几种非常有竞争力的算法进行了比较，尽管我们的记忆网络无法达到最佳性能，但评估得分仍可以满足，表3中介绍了定量评估结果。这意味着所提出的模型可以处理发生的degradation，而DGUNet的结果和指标是最好的，这在很大程度上要归功于其梯度下降策略，他们将其整合到近端梯度下降（proximal gradient descent）算法的梯度下降步骤中，以驱动其处理复杂的真实图像degradation。

表3 GOPRO数据集上图像去模糊的性能比较。红色标记的表示排名第一，蓝色标记代表排名第二。

图6 GOPRO数据集上图像去模糊的视觉性能比较。

4.2 MemoryNet消融实验

记忆增强的定量对比

为了验证本文提出的记忆增强和对比学习的有效性，我们对ISTD数据集进行了消融实验，如图7和表4所示。由于每个内存项都可以在所有查询中计算余弦相似度，并且相似性均以完成，因此我们替换了我们的记忆层，一个用于三分支，一个用于两个分支，一个用于单分支。进行消融实验，以检查其是否成功以增加记忆，并且表4中显示，在本研究中，记忆增强更适合三阶段恢复网络。为了更清晰地呈现本文所提出的记忆机制所发挥的作用，我们在一个轻量级网络中进行了特征可视化，如图7所示。输入是一张模态照片，第一层先经过记忆增强模块，然后是卷积层。从图中可以看出，添加了记忆机制后的特征图更适合网络传播，而未添加记忆机制的原始特征图明显与原始图像存在偏差。

表4 在ISTD数据集上的阴影去除消融实验。

图7 记忆增强层的特征可视化，左侧是未经任何处理的第一层的特征图，右侧是经过记忆增强后的第一层的特征图。

对比学习的定量比较

与生成式学习不同，对比学习无需关注实例的繁琐细节，而只需在抽象语义层面学习如何在特征空间中区分数据，因此模型及其优化过程变得更简单，并且具有更强的泛化能力。在本研究中，我们在残差网络后面添加对比学习模块，将其转变为一个判别器，目的是学习一个编码器，该编码器对同一类别的相似数据进行编码，同时使各类数据的编码结果尽可能不同。如表4所示，本文中的对比学习在诸如去阴影这样的任务中是有帮助的，并且能够更好地提升各项评估指标。在本文中，将记忆网络和对比学习相结合，明显可以看到峰值信噪比（PSNR）提升了 1 个点，达到 33.44，结构相似性指数（SSIM）达到 0.986，均方根误差（RMSE）达到 6.03 。

5 结论

在本文中，我们提出了一种用于图像修复的通用网络，名为记忆网络（MemoryNet），它由一个记忆增强模块和对比学习网络组成。该网络能够修复退化图像，包括有阴影、雨痕以及模糊的图像。这三项实验表明，这些方法在提升图像修复性能方面是有效的。此外，本文的模型在三个具有不同退化类型的数据集上取得了显著的峰值信噪比（PSNR）和结构相似性指数（SSIM）增益，这有力地证明了修复后的图像在视觉观感上是逼真的。在未来，我们将尝试更多不同的修复任务，例如图像增强、去除条纹等。