这篇文章把(早年张同学讲过的)超像素和 Transformer 进行了融合,也算一个创新。
原文链接:ICCV2023 | SPIN:超像素聚类与Transformer结合的轻量级图像超分辨率网络,性能SOTA
项目地址:https://github.com/ArcticHare105/SPIN
00 | 导言
基于transformer的方法在单图像超分辨率(SISR)任务中显示了令人印象深刻的结果。然而,当应用于整个图像时,自注意机制的计算成本很高。
现状:
目前的方法是将低分辨率的输入图像分割成小块,这些小块分别进行处理,然后融合生成高分辨率图像。然而,这种常规的patch划分过于粗糙,缺乏可解释性,导致在注意操作过程中产生伪影和非相似结构干扰。
解决:
提出了一种新的超级token交互网络(SPIN)。该方法利用超像素对局部相似像素进行聚类,形成可解释(因为只有相似的区域相互补充,不相似的区域被排除在外。)的局部区域,并利用超像素内关注实现局部信息交互。
设计了一个超像素交叉关注模块,通过超像素的替代实现信息的传播。大量的实验表明,所提出的自旋模型在精度和轻量化方面优于最先进的SR方法。
创新点:
-
提出了一种新的超分辨率模型,将超像素聚类与transformer结构相结合,产生了一个更可解释的框架。
-
提出了超像素内注意(ISPA)和超像素交叉注意(SPCA)模块,它们在超像素内和超像素之间运行,在保持捕获远程依赖关系的能力的同时,能够在不规则区域进行计算。
01 | 方法
图2 SPIN方法架构
所提出的SPIN模型架构如图2所示,主要由所提出的超像素交互(SPI)块组成。在SPI块之前,利用一个3 × 3的卷积编码器,将低分辨率图像ILR嵌入到高维特征空间。给定编码器,可以得到浅特征xemb为:
然后,将K个 SPI块堆叠在编码器的顶部,以提取包含输入图像丰富的低级和高级信息的更深的特征。每个SPI块包括四个部分:超像素聚合(SPA)、超像素交叉注意(SPCA)、超像素内注意(ISPA)和局部注意。每个块的输入特征首先通过SPA模块聚合成超像素。然后,ISPA模块捕获每个超像素内像素的依赖关系和相互作用,SPCA模块捕获远程像素之间的依赖关系和相互作用。为了增强局部区域内像素之间的交互作用,在ISPA和SPCA模块之后使用了一个局部关注模块,该模块使用基于窗口的关注,用于加强局部区域内的互动。我们使用重叠的patches来加强特征交互。形式上,对于第i个SPI块,整个过程可以表示为:
式中,Si表示第i个SPI块中超像素的特征,f(·)表示每个独立分量的函数。在之前的工作中,残差连接被用来简化整个训练过程。在K个SPI分块后,采用3×3卷积层和像素shuffle运算获得全局残差信息,并将残差信息加入到ILR的上采样图像中,以求解高分辨率图像ISR。
1)SPA 模块
与以往将输入图像或特征划分为规则块的方法不同,我们提出将输入特征划分为超像素。与常规patch容易将连接区域裁剪成不同的patch相比,超像素分区可以感知地将相似的像素分组在一起,可以描绘更精确的边界,降低生成模糊和不准确边界的风险。
具体来说,在超像素聚合过程中,我们在SSN中使用了基于软k均值的超像素算法。给定视觉标记(其中N = H × W为视觉标记的数量),假设每个标记属于M个超像素中的一个,因此有必要计算视觉标记与超像素标记之间的关联。
从形式上讲,超像素聚合过程是一个类似于期望最大化的过程,它包含了总共T次迭代。首先,如图3所示,我们通过在规则网格中平均标记来采样初始超级标记,称为Patchify。假设网格大小为,则超级token的数量为。对于第t次迭代,计算关联映射为:
其中为关联映射,为第i行第j列的值。请注意,超像素聚合只计算从每个token到周围超像素的关联映射,这保证了超像素的局部性,使其在计算和内存方面都很高效。之后可以得到超像素作为视觉标记的加权和,定义为:
其中表示沿列的归一化项。经过T次迭代,可以得到最终的关联映射AT。
2)SPCA 模块
由于超分辨率只捕获局部区域像素的局部性和互连性,这可能缺乏捕获超分辨率的远程依赖关系的能力。在这里,我们利用自关注范式通过超像素的替代来增强远程通信,这有助于利用特征之间的互补性来生成高质量的超分辨率图像。由于像素特征与所属的超像素特征高度相似,使得超像素成为在像素之间尽可能多地传播信息的有前途的替代品。
如图4所示,给定超像素特征,其中M表示超像素个数,平面像素特征。采用注意机制首先将像素信息传播到超像素。具体来说,使用线性投影来计算查询:, key:, value:为:
其中
分别为根据query、key、value的权值矩阵。可以通过首先计算查询和键之间的相似度并将其用作权重来聚合值来获得输出,该值可以表示为
其中为避免梯度消失的缩放因子,Su为更新后的超像素特征。请注意,与超像素聚合不同,该过程不考虑邻居限制,从而确保远程信息的传播。
一旦信息从像素传播到超像素,就需要将聚合的信息再分发回像素,以实现像素之间的信息传播。在这里进一步运用注意机制。具体来说,利用另一个权重矩阵从像素特征中获取查询。为了减少参数的数量,直接使用超像素特征作为键,将更新的超像素特征作为值,利用交叉关注将更新的超像素特征映射回像素级。
与Transformer块类似,经过上述过程,我们也采用了前馈网络(FFN)。FFN包含一个层归一化层,之后利用特征门控来调制输入特征和通道关注来提取全局信息。最后使用两个全连接层和GELU激活函数。
3)ISPA模块
有了关联图,提高超分辨率图像质量的一种直观方法是利用相同超像素内相似像素的互补性。为了实现这一点需要获得每个超像素对应的像素。然而,不同的超像素可能包含不同数量的像素,这使得并行处理变得困难,也会导致意想不到的内存消耗,因为总是有一些超像素包含大量的像素。
为了解决这个问题,如图5所示,我们采用关联图并选择与每个超像素最相似的top-N像素。假设一个超像素的附属像素为,其中N为所选像素的个数。我们按照标准的自注意机制,即Eq. 5和Eq. 6进行超像素内注意,其中包括权矩阵,用于查询、键和值投影。在超像素交互之后,利用top-N选择过程中生成的索引,将精细的像素特征分散回图像内各自的位置。
top-N选择可能会导致一些“忽略”像素,即这些像素不包含在任何超像素中。对于那些“被忽略”的像素,我们利用值投影来投影它们以获得更新的特征,然后将其与那些通过超像素内交互更新的像素集成在一起。与SPCA模块类似,在ISPA模块之后采用了相同的FFN。
02 | 实验结果
与其他先进的基于cnn和基于transformer的SISR模型的平均PSNR/SSIM比较。最好和次好的结果分别被突出显示和下划线,可以看出SPIN获得了SOTA结果。
针对×4升级任务的最先进的经典和轻量级变压器SR模型的定性比较。与其他型号相比,our (SPIN)可以还原更准确、更清晰的细节。
参考
https://openaccess.thecvf.com/content/ICCV2023/papers/Zhang_Lightweight_Image_SuperResolution_with_Superpixel_Token_Interaction_ICCV_2023_paper.pdf