OAPT：用于双JPEG伪影去除的偏移感知分区的Transformer

OAPT: Offset-Aware Partition Transformer for Double JPEG Artifacts Removal

https://github.com/QMoQ/OAPT

2408.11480 (arxiv.org)

基于深度学习的方法在去除单个JPEG伪影任务中表现出了显著的性能。然而，现有方法在处理双重JPEG图像时往往会退化，而这类图像在真实场景中普遍存在。

本文提出了一种用于去除双重JPEG伪影的偏移感知分区Transformer，称为OAPT。本文对双重JPEG压缩进行了分析，发现每个8x8块内最多会出现四种模式，并设计了我们的模型以聚类相似模式来克服恢复难度。

OAPT由两部分组成：压缩偏移预测器和图像重建器。具体来说，预测器估计第一次和第二次压缩之间的像素偏移量，然后利用这些偏移量来区分不同的模式。重建器主要基于多个混合分区注意力块（HPAB），该块结合了基于窗口的自注意力和稀疏注意力来处理聚类模式特征。

大量实验表明，在双重JPEG图像恢复任务中，OAPT的性能优于最先进的方法，高出0.16dB以上。此外，在不增加任何计算成本的情况下，HPAB中的模式聚类模块可以作为插件来增强其他基于Transformer的图像恢复方法。

问题描述

与单次压缩相比，双JPEG压缩在现实场景中更为常见，因为互联网上的图像通常会经历多个压缩周期[14, 28, 40, 62]。例如，如图1(a)所示，一张新拍摄的照片在首次压缩后存储在相机的磁盘中，之后在图像编辑（如裁剪）或上传到社交媒体云服务器时，会进行第二次压缩，从而成为双压缩图像。随着上述过程的重复发生，自然会产生多压缩图像。

然而，大多数以前的方法都是在单次压缩数据上进行训练的，并且只考虑了质量因子（QF）的范围，这导致在双JPEG图像恢复中性能大幅下降。

FBCNN[28]作为双JPEG伪影去除的开创性工作，估计了双压缩中的主导QF或在合成的双JPEG图像上进行微调。尽管它在减少单次和双次压缩伪影方面都是有效的，但它并没有专门针对双JPEG压缩的特性进行设计。

通过分析双JPEG图像的出现，本文发现除了QF的各种组合外，双JPEG图像中还存在明显的压缩偏移。由于像素级偏移，非对齐压缩导致第二次压缩中每个8x8块内不超过四种模式，如图1(b)所示。这会在8x8块中产生明显的边界和不同的模式特征，如图1(c)所示。通过在各种双压缩场景下测试DnCNN[73]发现，非对齐压缩的图像恢复任务要困难得多。与对齐压缩相比，非对齐压缩的峰值信噪比（PSNR）明显下降。

Motivation

首先，本文先证明了上面问题的存在。证明方法如下：

通过后面介绍的的数据退化方法，从LIVE1[55]数据集中生成了两组灰度图像，分别是偏移量为(0,0)的对齐双重JPEG图像和偏移量为(4,4)的非对齐双重JPEG图像。

计算了在(QF1, QF2) = (30, 50)和(QF1, QF2) = (50, 30)条件下，低质量图像与DnCNN处理后的相应增强图像之间的ΔPSNR值。

对于对齐的双重压缩图像，ΔPSNR值分别为2.06 dB和3.24 dB，而对于非对齐的双重压缩图像，ΔPSNR值分别为1.66 dB和1.64 dB。

发现，非对齐双重压缩图像的恢复性能不如对齐双重压缩图像。由此可以得出，非对齐双重JPEG图像的恢复更为复杂和困难。

在图1(b)中，描绘了双重压缩中的偏移量在每个8×8编码单元中产生了四种模式。这四个部分在第二次压缩中位于同一个8×8块中，但在第一次压缩中分别属于四个不同的8×8块。因此，在第二次压缩后，同一个8×8块中的四个部分会受到不同种类的块效应和压缩模式的影响而退化。

这启发了双重JPEG图像恢复的潜在简化方法可能是分别处理具有相同模式的聚类块。

方法描述

Overview

为了解决非对齐压缩图像的恢复问题，本文尝试获取压缩偏移并将相同的模式聚合在一起以进行更好的恢复。

受Swin Transformer中分区操作的启发，提出了一种针对8x8块中相同位置模式的新型分区策略。

基于此，设计了偏移感知分区Transformer（Offset-Aware Partition Transformer，简称OAPT）。OAPT由两个主要部分组成：一个基于CNN的压缩偏移预测器和一个基于Transformer的图像重建器。

压缩偏移预测器估计两次JPEG压缩之间的像素偏移，从而在整个图像中每个8x8块的相同位置对表现出相似压缩效果的模式进行有效聚类。

图像重建器由混合分区注意力块（Hybrid Partition Attention Blocks，简称HPAB）组成，结合了基于窗口的自注意力机制和稀疏注意力机制来处理聚类后的模式特征。交替的注意力和模式聚类机制可以增强双压缩伪影去除的鲁棒性。

压缩偏移预测器

偏移体现在均匀且周期性的8x8块中，表明行和列的偏移范围从0到7。在图像退化过程中，在第二次压缩中随机移除前i行和前j列，以模拟非对齐压缩，并将得到的偏移标记为(i, j)。

关于合成双重压缩图像的详细信息见后面。如图2所示，压缩偏移预测器采用ResNet-18。它以低质量图像为输入，并估计行和列的偏移。

预测器的输入仅是输入图像左上角的一个44x44的区块，JPEG从此处开始将图像分割成块并进行压缩。

为了最小化参数数量，将残差块中的所有卷积层替换为深度可分离卷积[13]层，并将这些修改后的块表示为D-Resblocks。

此外，引入了一个额外的线性层、sigmoid激活函数和取整操作，以生成一个由两个介于0和7之间的整数组成的输出，如公式所示：

其中，r′ 和 c′ 是最后一个线性层的输出，而 r^ 和 c^ 是行和列的最终预测偏移。

利用紧凑的CNN网络和缩小的输入尺寸，预测器的计算成本远低于重构器。

通过最小化以下损失函数来优化偏移预测器：

其中，r 和 c 是行和列的真实偏移。

基于混合分区注意力块的图像重构器

本文设计了一种新的分区操作，以匹配混合注意力机制，用于双重JPEG图像恢复。

利用预测器估计的偏移量，可以将每个8x8块分割成四部分，并对相同位置的模式进行聚类，以便进行后续的自注意力处理。

图像重构器主要由这种新的混合注意力机制组成。图像重构器包含三个部分：一个用于浅层特征提取的卷积层、一系列基于Transformer的模块用于深层特征提取，以及一个用于高质量图像重构的卷积层。

深层特征提取部分是OAPT的核心，由混合分区注意力块（HPAB）和一个卷积层组成。如图2所示，HPAB是一个残差块，包含一个卷积层、四个Swin Transformer层（STL）和两个基于模式聚类的Swin Transformer层（PC-STL）。HPAB可以提供基于窗口的普通自注意力[39]和稀疏注意力。

STL专注于每个窗口内局部连续特征的密集注意力。

而PC-STL则配备了基于模式聚类的多头自注意力（PC-MSA）。模式聚类模块利用估计的偏移量将每个8x8块分成四种模式，并对相同位置的模式进行聚类以进行自注意力处理，如图3所示。