AAAI 2024 | 利用通道模态冗余实现端到端 RGB-D 图像压缩

题目： End-to-End RGB-D Image Compression via Exploiting Channel-Modality Redundancy 作者： Huiming Zheng, Wei Gao 论文链接： https://underline.io/lecture/92675-end-to-end-rgb-d-image-compression-via-exploiting-channel-modality-redundancy 来源：AAAI 2024 内容整理：周楚骎 RGB-D图像作为一种3D数据，已被广泛应用于目标跟踪、三维重建、遥感测绘等任务。在计算机视觉领域，RGB-D图像的重要性正在逐渐增长。然而，现有的基于学习的图像压缩方法通常分别处理RGB图像和深度图像，不能完全利用模态之间的冗余信息，限制了码率失真性能的进一步提高。为了克服这一缺陷，本文提出了一种基于学习的双分支RGB-D图像压缩框架。与传统的RGB域压缩方案相比，本文提出一种YUV域的压缩方案，用于空间冗余去除。此外，还引入了模态内注意力 （IMA） 和跨模态注意力 （CMA） 来消除模态冗余。为了从跨模态先验信息中获益，在条件熵模型中提出了上下文预测模块（CPM）和上下文融合模块（CFM），使上下文概率预测更加准确。实验结果表明，该方法在两个RGB-D图像数据集中均优于现有的图像压缩方法。与BPG相比，本文提出的框架可以为RGB图像实现高达15%的比特率节省。

简介

RGB-D 图像是一种重要的 3D 数据格式。它已被广泛用于 3D 场景重建、突出目标检测、机器人与自主导航、医学影像与健康监测、环境监测等领域。与 RGB 图像不同，深度图像包含有关从视点到场景对象表面的距离的信息，该视点提供了 3D 场景之间的深度信息。因此，RGB-D联合分析方法在计算机视觉任务中很受欢迎。然而，这些方法使用额外的模态，这将带来多余的存储和传输成本。因此，设计一个高效的RGB-D图像压缩方法是一项重要且具有挑战性的工作。

基于深度学习的图像压缩已经发展了好几年。众多工作提出了提高码率失真性能和优化编码框架的方法。此外，一些开源算法库也有效地促进了该领域的繁荣。然而，现有的方法侧重于单图像压缩，而忽略了RGB和深度模态的直接交互性。没有充分考虑模态冗余，限制了码率失真性能的改进。此外，知识引导压缩是最相关的主题之一。编码框架可以使用数据源本身的其他信息，也可以分析其自身模块中的其他信息，以更好地消除冗余。立体图像压缩框架采用单调变换来消除视图冗余。光场图像压缩框架利用光场图像固有的相似性来消除不同视角的冗余。360° 图像压缩框架利用纬度自适应编码方案为不同区域分配不同的比特数。虽然这些方法在一定程度上探索了模态冗余去除，但由于RGB图像与深度图像的分布存在显著差异，因此在RGB-D图像压缩中无法实现更高的压缩比。因此，有必要开发一个专用于RGB-D图像的压缩框架。

本文利用模态和通道之间的冗余信息，提出了一种基于学习的高效RGB-D图像压缩网络。大多数基于学习的方法通常在RGB域中对图像进行采样和压缩，而本文的方法选择在YUV域中对图像进行采样，以消除深度图像变换域中的空间冗余。此外，本文在特征提取模块中设计了模态内注意力（IMA），在主编码器模块中设计了跨模态注意力（CMA），以分别消除通道冗余和模态冗余。本文采用上下文预测模块（CPM）和上下文融合模块（CFM）在条件熵模型中充分挖掘了两种模态之间的相干性，并利用了跨模态先验信息，为熵编码器提供了更准确的概率预测信息。实验结果证明，与单图像压缩方法相比，所提网络在多个广泛使用的RGB-D数据集上具有更好的码率失真性能。本文提出的方法的贡献可以总结如下：

• 本文提出了一种基于学习的RGB-D图像压缩框架，利用通道和模态之间的冗余信息。该框架在YUV域而非RGB域进行，有利于消除深度图像的空间冗余。
• 模态内注意力（Intra-Modality Attention）和 模态间注意力（Cross-Modality Attention）旨在消除跨通道冗余和跨模态冗余，以获得更高的压缩比。具体来说，多头自注意力和多头交叉注意力被集成到模块中，以实现更高效的跨通道和跨模态信息交互。
• 基于条件上下文的熵模型被改编为揭示模态之间的依赖关系。此外，上下文预测模块和上下文融合模块经过精心设计，可实现高效的概率预测。
• 实验结果表明，在两个RGB-D图像数据集中，与现有的图像压缩方法相比，所提框架实现了SOTA性能。

方法

概述

本文的RGB-D图像压缩框架的整体架构如下图所示。该网络基于Transformer架构。输入 RGB 和深度图像在 YCbCr 子采样色彩空间中转换为 4 个通道。在RGB图像中，U和V通道的重量和高度是Y的权重和高度的一半。深度图像仅保留 Y 通道信息。本文将

作为输入通道。

图 1 所提方法的整体网络架构。

首先，将输入通道馈入特征提取模块，以消除通道冗余。特征提取后分别从

得到特征图

。然后连接

进行下一阶段输入

。在编码器阶段（分析变换），为输入

和

提供了一个双分支网络。所提出的跨模态注意力允许潜在表征相互学习跨模态信息。在编码器阶段之后，潜在表示

和

被发送到量化器。然后将量化的潜在表示

和

发送到条件熵模型中，以进行准确的符号概率预测。在解码器端（合成变换），

和

被馈送到对偶分支解码器框架中，用于特征恢复和上采样。在解码过程之后获得特征图

和

。最后，在特征重构模块中，将特征图yuv分为Y、U、V通道

。细节恢复和重建在特征重建模块中进行。本文将特征提取模块、编码器、量化器、解码器、特征重构模块分别命名为

。除超先验外，主要的编解码过程可以表述为：

其中

代表

输入之一

模态内注意力

在本文提出的框架中，本文在特征提取模块和特征重构模块中使用模态内注意力来减少通道冗余。IMA的框架如图2所示。主要框架基于两个连续的 Swin Transformer 块。

图 2 模态内注意力 （IMA） 的架构

给定尺寸为

的输入特征图，基于窗口的注意力首先将特征图重塑为

的大小，而

表示窗口大小。

窗口是从操作中获得的。然后，对每个窗口都采用自注意力机制。三个可学习的共享相同的参数的权重矩阵

乘到局部特征图

，以便分别得到查询

、键

和值

。该过程可以描述为：

然后，注意力函数使用每个键计算查询的点积。结果包括相对位置偏差，以获得更好的计算复杂性。采用softmax算子对注意力分数的结果进行归一化。上述过程可以定义为：

其中

是维度，

是相对位置偏差。模态内注意力的主要过程可以概括为：

其中

和

分别是 (S)W-MSA 和 MLP 模块的输出特征。

是输入特征图。

是 LayerNorm 函数。

表示基于窗口的多头自注意力，

表示基于移位窗口的多头自注意力。

跨模态注意力

图 3 跨模态注意力 （CMA） 的架构

继模态内注意力之后，本文还设计了跨模态注意力。网络架构如上图所示。与消除通道冗余的IMA不同，CMA致力于消除模态冗余。此外，CMA可以进一步整合不同模态之间的查询。IMA和CMA的框架相似，主要区别在于CMA采用多头交叉注意力而不是多头自注意力来实现模态之间的信息交互。给定本地窗口中的输入RGB特征图

和深度特征图

，适应z

的跨模态注意力的完整过程可以定义为：

而适应

的跨模态注意力的完整过程可以描述为：

其中

和

表示 (S)W-MCA 的输出 RGB 特征图和深度特征图。

和

是 MLP 模块的输出特性。

表示基于窗口的多头交叉注意力，

表示基于移位窗口的多头交叉注意力。

基于条件上下文的熵模型

图 4 基于条件上下文的熵模型的架构

传统的单图像压缩方法通常利用超先验信息作为条件先验。一个空间位置的概率密度可以通过其他位置的已知概率密度来估计。但对于具有跨模态信息的RGBD图像，超先验提供附加信息是不够的。在本文提出的方法中，本文采用基于条件上下文的熵模型进行更准确的概率估计。基于条件上下文的熵模型的架构如上所示。在编码器阶段之后，潜在表示被发送到超编码器和超解码器以获取空间分布信息。此外，它还被馈送到上下文预测模块（CPM）以获取上下文先验信息。然后，将 CPM 的输出特征图发送到上下文融合模块 （CFM） 进行跨模态信息聚合。对于深度潜在表示，本文使用上下文和空间先验来估计熵参数。对于更复杂的RGB潜在表示，除了前者之外，本文还使用多余的跨模态信息来提高概率预测的准确性。具体来说，本文将

作为深度潜在表示的相似性，

作为 RGB 潜在表示的相似性。

和

表示

和

的第 i 个元素。估计概率质量函数 （PMF）

和

如下式所示。

上下文预测模块和上下文融合模块

为了进一步模拟 PMF，对上下文预测模块进行了调整，以准确估计上下文先验信息。在上下文预测模块中采用了掩模标度余弦注意力（MSCA）。此外，本文提出了上下文融合模块而不是concat操作，以更好地聚合跨模态信息。Mask Scaled Cross Cosine Attention （MSCCA） 集成到上下文融合模块中，以实现模态之间的信息交互。为了保证串行编码解码顺序，本文在 transformer 架构中使用了look ahead mask mechanism。本文没有采用缩放的点自注意力，而是采用缩放余弦注意力，这使得模型的训练更加稳定。此外，使用对数空间连续相对位置偏置代替线性空间相对位置偏差，以获得针对高分辨率图像的更好的重建质量。

损失函数

在训练阶段，损失函数

描述如下：

其中

和

是 YUV 通道和深度通道的加权均方误差 （MSE）。他们可以是如下形式：

图像收集

模型参数量的缩放需要对训练数据进行相应的缩放。但是目前还没有可用于 IR 的大规模高质量图像数据集。尽管DIV2K和LSDIR提供了高图像质量，但它们的数量有限。较大的数据集，如ImageNet（IN）、LAION-5B和SA-1B，包含更多的图像，但它们的图像质量不符合本文的高标准需求。为此，本文收集了一个新的高分辨率图像大数据集，其中包括 2000 万张 1024×1024 高质量、纹理丰富且内容清晰的图像。本文还从FFHQ-raw数据集中加入了额外的70K未对齐高分辨率面部图像，以改善模型的面部修复性能。

多模态语言指导

扩散模型以其基于文本提示生成图像的能力而闻名。本文认为，文本提示也可以显着帮助 IR，原因如下：

（1） 理解图像内容对 IR 至关重要。现有的框架经常忽略或隐含地处理这种理解。通过结合文本提示，本文明确地将对LQ图像的理解传达给IR模型，促进有针对性地恢复缺失的信息。 （2） 在严重退化的情况下，即使是最好的 IR 模型也可能难以恢复完全丢失的信息。在这种情况下，文本提示可以作为一种控制机制，根据用户偏好有针对性地完成缺失的信息。 （3）本文还可以通过文本描述所需的图像质量，进一步增强输出的感知质量。

为此，本文进行了两个主要修改。首先，本文修改了整体框架，将LLaVA多模态大型语言模型纳入本文的设计中。LLaVA 将经过降解的鲁棒性 LQ 图像

作为输入，并明确理解图像中的内容，输出以文字描述的形式。然后，这些描述将用作指导恢复的提示。此过程可以在测试期间自动执行，无需人工干预。其次，遵循PixART的方法，收集所有训练图像的文本注释，以加强文本控制在模型训练中的作用。这两个变化使SUPIR能够理解图像内容，并根据文本提示恢复图像。

负面提示词提示与采样

无分类器引导（CFG）提供了另一种控制方式，即使用负面提示词来指定模型不需要的内容。本文可以使用此功能来指定模型不生成低质量的图像。具体来说，在扩散的每一步，本文将使用正面提示词

和负面提示词

进行两次预测，并将这两个结果的融合作为最终输出

:

其中

是带适配器的扩散模型，

是噪声在时间步长 t 处的方差，

是超参数。在本文的框架中，

可以是正面提示词的图像描述，

是负面提示词，例如油画、卡通、模糊、肮脏、凌乱、低质量、变形、低分辨率、过度平滑。预测的正方向和负方向的准确性对于CFG技术至关重要。然而，本文的训练数据中缺少负面质量的样本和提示可能会导致微调的 SUPIR 无法理解负面提示。因此，在采样期间使用负质量提示可能会引入伪影。为了解决这个问题，本文使用 SDXL 生成了对应负质量提示的 100K 图像。本文反直觉地将这些低质量的图像添加到训练数据中，以确保所提出的SUPIR模型可以学习负质量概念。

恢复引导采样

强大的生成先验是一把双刃剑，因为过多的生成容量反过来会影响恢复图像的保真度。这突出了 IR 任务和生成任务之间的根本区别。需要限制生成的方法，以确保图像恢复忠实于 LQ 图像。本文改进了EDM采样方法，并提出了一种恢复引导采样方法来解决该问题。本文希望在每个扩散步骤中选择性地引导预测结果

接近LQ图像

。

算法1 恢复引导采样算法

具体算法如上图所示，其中T为总步数，

为T步的噪声方差，c为附加文本提示条件。

是五个超参数，但只有

是与恢复有关指导，与原始EDM方法相比，其他方法保持不变。本文在预测输出

和 LQ 潜在

之间执行加权插值，作为恢复引导的输出

。由于图像的低频信息主要产生于扩散预测的早期阶段（其中t和

比较大，权重

也很大），因此预测结果更接近

，以增强保真度。在扩散预测的后期阶段，主要产生高频细节。此时不应有太多约束，以确保可以充分生成细节和纹理。此时，t和

相对较小，权重k也较小。因此，预测结果不会受到太大影响 通过这种方法，可以在扩散采样过程中控制生成，以保证保真度。

实验

数据集

SUN-RGBD： SUN-RGBD 数据集是一个广泛使用的计算机视觉研究数据集，用于室内场景理解和深度感知任务。该数据集提供室内环境中的RGB图像、深度图像和语义分割标签等数据，适用于许多不同的计算机视觉任务。该数据集包含 10,000 张 RGB-D 图像。对于训练，随机选择了 8,000 个图像对，同时选择了 1,000 个图像对进行验证，另外 1,000 个图像对被保留用于测试。

NYU-Depth V2： NYU-Depth V2数据集包含视频序列，捕获由Microsoft Kinect的RGB和深度摄像头记录的各种室内场景。它包括 1,449 张带注释的 RGB 图像和深度图像。这些图像来自三个城市的 464 个场景。本文将整个数据集分为三部分，1,159 个图像对用于训练，145 个图像对用于验证，145 个图像对用于测试。

实验细节

训练策略

本文共同训练整个网络。建议的网络基于支持 CUDA 的 PyTorch 实现。本文为超参数

设置不同的值来控制比特率。

配置参考 CompressAI 。在训练过程中采用了 Adam 优化器。本文将学习率初始化为 1e − 4。在训练过程中，它随着模型的更新而逐渐减少，最终下降到 1e−5。批大小设置为 4。本文为每个模型训练大约 1000 个 epoch。根据Tesla V100，训练阶段大约需要十天。输入训练数据被修剪为 256×256 的大小，便于模型推理。训练数据主要基于SUN-RGBD数据集。当模型在 NYU-Depth V2 数据集上测试时，本文使用 NYU-Depth V2 数据集中的训练数据集对预训练模型进行了大约 100 个 epoch 的微调。

评估指标

本文采用PSNR作为评估指标。PSNR是评估图像质量的客观指标，它反映了图像的信号保真度。此外，本文比较了 Bjontegaard delta 码率 （BD-Rate），以获得定量码率失真性能。请注意，PSNR 和 BD-Rate 指标是在 YUV420 域中评估的。

基线

本文将本文的方法与若干性能良好的单图像方法，立体图像压缩方法和一些经典的基于学习的方法，超先验， Mbt2018 ， Cheng2020attention。此外，传统的单模态图像压缩方法BPG也与本文提出的框架进行了比较。

实验结果

量化结果

表1 SUN-RGBD 数据集和 NYU-Depth V2 数据集与 BPG 的 BD-Rate （%） 比较。

表 1 显示了两个数据集中针对 BPG 的方法的编码性能。BD-Rate值为负，表示该算法的编码性能优于基准算法。否则，它比基准算法更差。为了确保公平的比较，本文采用与该模型相同的训练数据集和训练方法来训练其他基于学习的方法。很明显，本文提出的方法获得了最佳的RD性能。相比之下，采用单图像压缩方法，显著提高了所提方法的RD性能。具体来说，本文的方法在两个数据集中RGB图像的BD-Rate指标上比BPG提高了10%以上。此外，本文还绘制了RD曲线，以进一步可视化各种方法之间的性能差距。图5显示了SUN-RGBD数据集中RGB图像和深度图像的YUV-RSNR结果。这表明我们提出的框架超越了其他框架，展示了最佳的研发性能。此外，从下图可以看出，该模型对深度图像的压缩效果明显优于RGB图像。

图5 在SUN-RGBD数据集中测试的RGB图像（左）和深度图像（右）的码率失真曲线。

定性结果

为了更直观地展示每个模型的压缩效果，我们将每个模型的压缩图像可视化到下图中。需要注意的是，为了公平起见，我们尽量保持所有模型以相同的比特率进行压缩。如下图所示，本文的方法在使用较少比特率的前提下表现出优越的主观视觉质量。局部细节放大后，我们的方法仍能保留原始图像的语义信息（如图中的字母）。

图6 RGB图像压缩结果的视觉质量比较

运行时间和复杂性

我们提出的模型参数数为 69.03 M。对于输入分辨率为 256×256 的 RGB-D 图像对，FLOP 达到 6.93 Mil/pixel。当本文在Tesla V100的两个数据集中测试本文提出的模型时，平均编码时间和解码时间分别为 11.696 秒和 8.582 秒。与其他基于学习的模型相比，我们的方法引入了额外的计算成本，但获得了显着的速率失真性能提升。

消融研究与分析

案例1：条件熵模型的有效性 如下表所示，我们通过代入来验证熵模型中每个模块的有效性。它是在SUN-RGBD数据集上进行的。我们可以发现每个模块都有助于提高整体编码性能。此外，值得注意的是，基于条件上下文的熵模型对RD性能的贡献最大。

表2 条件熵模型中每个分量的消融研究。本文的熵模型基于 Mbt2018。

案例 2：YUV 域压缩的有效性 为了验证对于RGB-D图像，YUV域的压缩效率更高，与所提出的框架相比，本文呢设计了一个原始输入为RGB图像和深度图像的框架，而不是四个通道。为了确保比较实验的公平性，本文保留了IMA和CMA。消融实验表明，在YUV域中测试时，YUV域压缩方法与RGB域压缩算法相比，性能提升明显。

案例 3：IMA 和 CMA 的有效性 我们评估了IMA和CMA的疗效，结果如下表所示。结果表明，每个模块都提高了整个RD性能。值得注意的是，单独使用 CMA 时的结果比单独使用 IMA 时效果更好。结果表明，不同模态信息交互和跨模态冗余消除在RGB-D图像压缩中具有重要意义。

表3 IMA和CMA的消融研究

总结

本文提出了一种基于学习的RGB-D图像压缩框架，显著提高了RGB-D图像的压缩效率。首先，本文将输入图像对从RGB域转换为YUV420域，以消除空间冗余。模态内注意力（IMA）在特征提取和特征重建阶段进行设计，以减少跨通道冗余。然后，在编码器和解码器中调整跨模态注意力（CMA），以消除跨模态冗余。为了有效地利用模态之间的先验信息，采用基于条件上下文的熵模型进行更好的符号概率估计。在熵模型中，本文用掩码缩放余弦注意力 （MSCA） 更改上下文预测模块 （CPM）。还提出了上下文融合模块（CFM）来聚合跨模态信息。对比实验结果和消融研究验证了所提方法的有效性。