论文阅读笔记 | Transformer-based No-Reference Image Quality Assessment via Supervised Contrastive Learning

文章链接：https://doi.org/10.48550/arXiv.2312.06995
代码链接：https://github.com/I2-Multimedia-Lab/SaTQA

文章题目

Transformer-based No-Reference Image Quality Assessment via Supervised Contrastive Learning

发表年限

2023

期刊/会议名称

Association for the Advancement of Artificial Intelligence, AAAI24

论文简要

• 在本文中，作者提出了 SaTQA，这是一个专门针对无参考图像质量评估（NR-IQA）的模型，它结合了监督对比学习和 Transformer。
• 最初，$SCL$ 被用来解决现有模型面对真实数据集进行质量预测时的不足。
• 为了增强对失真图像的特征提取能力并克服 CNN 全局建模的限制，作者引入了MSB模块。
• 此外，作者设计了 $PAB$ 模块以有效融合图像退化特征和失真信息。
• 在七个标准IQA数据集上的实验结果证明了所提出模型的卓越性能。

动机

• 由于现有的卷积神经网络（CNN）和基于 Transformer 的NR-IQA方法具有强大的特征提取能力，这些方法已经取得了相当的进展。然而，当面对未知的真实失真数据集时，它们仍然显示出有限的能力。
• 大多数现有的基于深度学习的 NR-IQA 模型都是使用整个训练数据集进行训练的，没有区分不同的失真类型和等级，这可能使它们在面对如 LIVE-FB 这样的大规模真实数据集时表现出有限的能力。
• 如图1所示，即使是最近最强大的基于 Transformer 的 NR-IQA 模型 MANIQA，在 LIVE-FB 上的PLCC结果也降至不到0.6。
• 这是因为对于人眼而言，即使是相同等级的失真，但由两种不同类型的噪声引起的，可能会导致非常不同的视觉感知。
• 因此，如果在训练过程中不区分不同的噪声类型和等级，模型可能在面对包含过多失真的大规模数据集时表现不佳。
  
• 一些基于CNN的方法，如 CONTRIQUE，尝试通过在训练中指定失真类型和等级或使用其他类型的内容作为失真分类的标签来解决这个问题。然而，这种策略不适用于大规模真实数据集，其中失真类型和等级不可能明确指定。
• 人类对图像质量的感知通常是通过观察整个图像来进行的，即通过全局整合整个图像的内容失真。由于 CNN 的接受域有限，基于 CNN 的 IQA 模型的性能可能会受到限制。与 CNN 相反，Transformer 由于其长距离依赖性建模能力，可以整合图像的全局失真信息，因此自然适合于无参考图像质量评估（NR-IQA）。
• 纯 Transformer 在提取图像边缘和纹理等局部细节方面能力较弱，这可能会降低最终的预测性能。
• 从图2中的 Grad-CAM 可以看出，纯 $ViT$ 模型更多地关注图像的语义区域，而忽略了图像的失真区域。

主要思想或方法架构

• 提出的方法是一个基于监督对比学习和Transformer的无参考图像质量评估（NR-IQA）方法，名为SaTQA，图3展示了所提出的模型架构。

• 给定一个失真图像 $I\in R^{H\times M\times 3}$，其中 $H$ 和 $W$ 分别表示高度和宽度，NR-IQA 的目标是预测其客观质量分数。
• 首先，作者使用ViT提取图像的原始特征，$ViT$ 是 Vision Transformer 的缩写，是一种基于 Transformer 架构的视觉处理模型，其中每一层提取的特征定义为 $F_i\in R^{(h\times w)\times C_i}$， i ∈ { 1 , 2 , . . . , 12 } i \in \{1,2,...,12\} i∈{1,2,...,12}。
• 在本文中，作者提取其中四层并通过通道维度将它们连接起来，表示为 $\hat{F}$ $\in R^{{\sum }_i(h\times w)\times C_i}$， i ∈ { 1 , 2 , 5 , 6 } i \in \{1,2,5,6\} i∈{1,2,5,6}，其中 $C_i$ 表示 ViT 层的维度，$h$ 和 $w$ 分别表示相应输出的高度和宽度。然后作者沿通道维度缩小 $\hat{F}$ 以获得 $\widehat{F^{{}^{\prime }}}=R^{h\times w\times D}$。
• 接下来作者使用一个三阶段的 $MSB$ 模块进行图像失真提取，其中前两个阶段在 $MSB$ 处理后进行下采样。
• 失真信息由骨干网络提取，输出为 $R\in R^{h_4\times w_4\times D}$，其中 $D$ 表示特征的通道维度，$h_4$ 和 $w_4$ 表示最后阶段输出的高度和宽度。同时，失真图像通过经过 $SCL$ 训练的 ResNet-50，输出为 $P$，$P$ $\in R^{(h_4\times w_4)\times C_i\times {\sum }_i{c}_i}$。
• 作者使 $P$ 和 $R$ 具有相同的形状，因此 $D={\sum }_i{c}_i$。接下来，$PAB$ 模块融合 $P$ 和 $R$ 的特征，以获得图像质量的特征 $S$，并将最终的 $S$ 回归到与图像对应的客观质量分数。

---

• 作者需要模型能够区分不同图像的失真类型和等级，因此这个辅助任务的目标类似于图像分类。
• 然而，不同于传统的监督分类，作者在一个大规模的合成数据集上使用监督对比学习（$SCL$）训练模型，以获得比分类更稳健和通用的退化特征，使得相同失真类型（等级）的退化比不同失真类型（等级）的退化更加接近。
• 然后作者将预训练的模型推广到含有潜在退化的其他数据集，例如真实失真数据集。

• 如图所示，在开始时，各种失真图像在特征空间上的分布是随机和无序的。作者的目标是，使用监督对比学习后，具有相同失真和等级的图像可以聚集在一起。
  

• 在作者提出的用于失真图像聚类的监督对比学习中，首先，定义 $I_u^v$ 表示失真图像， u ∈ { 1 , . . . , U } u\in\{1, ..., U\} u∈{1,...,U} 表示与 $I$ 对应的不同失真类型， v ∈ { 1 , . . . , V } v\in\{1, ..., V\} v∈{1,...,V} 表示每个 $u$ 下的失真等级，其中 $U$ 和 $V$ 分别表示失真类型和等级的总数。

• 因此，连同参考图像的类别，总共有 $U\times V+1$ 个类别。也就是说，这个辅助任务被转换为区分不同类型和等级的失真图像的问题。

• 受到前人工作的启发，作者使用编码器网络 $f$(·) 和投影器 $g$(·) 结构来学习图像的退化特征。在本文中，编码器使用 ResNet-50，投影器使用多层感知机（$MLP$）。

• 对于给定的图像 $x$，$r=f(x)$，$z=g(r)$（$z\in R^D$），其中 $z$ 表示输出的 $D$ 维特征，$r$ 表示编码器输出的特征表示。与前人工作一样，编码器输出在传递给投影器之前进行 $L_2$ 规范化。对比损失函数使用监督归一化温度缩放的交叉熵（NT-Xent），其可以表示为：
  

• 其中 $N$ 表示一个 mini-batch 中所有失真和参考图像的数量，$I$ 表示指示函数，$\tau$ 表示温度参数，P(i)表示属于同一类i的集合，$|P(i)|$ 是其基数。$\varphi (m,n)$ 表示 $\varphi (m,n)=m^{T}n/\Vert m{\Vert }_2\Vert n{\Vert }_2$。

---

• 作者提出的多流块（$MSB$）如图所示，它结合了 CNN 用于提取图像低级边缘特征和 Transformer 用于提取图像高级内容特征。

• 首先，作者将 $ViT$ 提取后的特征 $\hat{F}\in R^{D\times h\times w}$ 重塑为特征图 $\widehat{F^{{}^{\prime }}}$；

• 然后作者将其沿通道维度分为 $3$ 个部分，${\chi }_1\in R^{D_1\times h\times w}$，${\chi }_2\in R^{D_2\times h\times w}$，${\chi }_3\in R^{D_3\times h\times w}$，其中 ${\chi }_1$ 和 ${\chi }_2$ 被送入 CNN，而 ${\chi }_3$ 被送入 Transformer。

• ${\chi }_1$ 将通过可变形卷积（$DeformConv$）和线性层，得到 ${\chi }_1^{{}^{\prime }}$。

• ${\chi }_2$ 首先通过深度卷积（$DwConv$）和最大池化层，然后上采样得到 ${\chi }_2^{{}^{\prime }}$。

• ${\chi }_3$ 将通过多头自注意力（$MHSA$）和线性层，得到 ${\chi }_3^{{}^{\prime }}$。

• 这些可以被公式化为：
  

• 最后，${\chi }_1^{{}^{\prime }}$、${\chi }_2^{{}^{\prime }}$ 和 ${\chi }_3^{{}^{\prime }}$ 被连结在一起并重塑，表示为 ${\hat{F}}^{{}^{\prime }}\in R^{D\times h\times w}$。此外，作者通过 $CBAM$ 模块（即通道注意力和空间注意力）将 $\hat{F}$ 与 ${\hat{F}}^{{}^{\prime }}$ 进行残差学习连接。

• 在本文中，${\chi }_1$、${\chi }_2$ 和 ${\chi }_3$ 在通道维度的百分比会随着不同阶段而变化。一开始，所有三者在通道中分享相同的百分比，然后 ${\chi }_3$ 会逐渐增加其在总通道中的百分比。

• 注意，在 CNN 分支中，可变形卷积可以增加空间采样位置，而深度卷积可以减少计算数量并补充在 Transformer 分支中未建模的通道注意力。

---

• 根据上面所述，作者可以获得扭曲图像特征 $P$ 和 $R$。
• 通过在大规模图像质量评估数据集 $KADIS$ 上进行自监督预训练后，退化 $P$ 可以用来区分待评估图像的扭曲级别和类型。如图所示，作者将 $P$ 映射到查询（$Q$），将 $R$ 映射到键（$K$）和值（$V$），然后通过公式计算它们，以获得最终的图像质量特征 $S$。
• 通过这种计算，模型可以在评估中表现良好，并且由于 $P$ 辅助训练，即使在小规模的图像质量评估数据集上也能更快地收敛。

---

• 对于得分预测，作者首先根据上面将补丁注意力块的 tokens 重塑为特征图，然后通过全局平均池化（$GAP$）获得的质量向量回归质量得分。

• 在每一批图像中，作者通过最小化回归损失 $L$ 来训练质量预测模型，如下所示：
  

• 其中 $N$ 表示批次大小，$p_i$ 表示模型对第i张图像预测的图像质量得分，$y_i$ 表示相应的客观质量得分。

实验结果

• 个体数据集评估。如表2所示，SaTQA 在所有七个数据集上的 $SROCC$ 和 $PLCC$ 表现优于当前最先进的方法。
• SaTQA 在真实数据集上也表现良好，不仅在小规模数据集 LIVE Challenge 上表现出色，还在当前的大规模 LIVE-FB 数据集上表现出色。

---

• 单独失真评估。由于图像中存在多种类型的失真，尤其是在真实数据集中，一张图像可能包含多种类型的失真（例如，模糊、噪声等）。
• 为了验证 SaTQA 在失真类型上的通用性能，作者在合成数据集 LIVE 和 CSIQ 上进行了实验，并比较了 $SROCC$ 表现。
• 从表3可以看出，所提出模型在 CSIQ 和 LIVE 数据集上具有强大的泛化性能。

---

• 跨数据集评估。为了验证模型在不同数据集上的泛化性能，作者进行了跨数据集实验。在这个验证中，训练过程在特定数据集上进行，测试则在另一个不同的数据集上执行，测试期间不进行任何参数调整。
• 如表4所示，所提出的模型具有强大的泛化能力。

---

• 消融研究。在表5中，作者提供了消融研究以验证所提出的模型的每个组件的效果。
• 作者使用在 ImageNet 上预训练的 $ViT$ 作为基线（#1），其 patch 大小为 8，并使用原始的 CLS token 进行质量回归。
• 监督对比学习。为了有效利用大量未标记数据并增强模型在真实 IQA 数据集上的预测性能，作者提出了 $SCL$ 模块。
• 在 #2 中的结果表明，使用 $SCL$ 可以显著提高模型在 LIVE Challenge上 的性能（$+0.044$ $PLCC$，$+0.037$ $SROCC$）和在 CSIQ 数据集上（$+0.041$ $PLCC$，$+0.04$ $SROCC$）。
• 多流块。为了结合 CNN 和 Transformer 在 IQA 任务上的特征提取能力，作者提出了 $MSB$ 模块。在 #3 中，该模型在 LIVE Challenge 和 CSIQ 数据集上的 $SROCC$ 和 $PLCC$ 有 $0.01$ 的提升。
• 补丁注意力块。提出 $PAB$ 是为了使用 $SCL$ 学到的退化作为查询，与 $MSB$ 提取的感知失真计算交叉注意力。
• 从 #4 中可以看出，通过添加$PAB$，模型在 LIVE Challenge 数据集上的 $SROCC$ 和 $PLCC$ 提高了 $0.02$，在 CSIQ 数据集上增加了 $0.002$。

---

• 视觉结果分析。作者提供了 MANIQA 和 SaTQA 在 $CSIQ$ 测试集上对 $JPEG$ 失真图像的 $DMOS$ 预测结果作为示例。

• 如图所示，可以看出 SaTQA 在图像失真等级排名和整体 $DMOS$ 值预测方面都优于 MANIQA，其中 MANIQA 在预测第四和第五失真等级时产生了错误。