[论文解读] Spatial-Temporal Transformer for Video Snapshot Compressive Imaging

Lishun Wang, Miao Cao+2 more|arXiv (Cornell University)|Sep 4, 2022

Sparse and Compressive Sensing Techniques被引用 10

一句话总结

该论文提出 STFormer，一种用于视频快照压缩感知成像（SCI）重建的时空 Transformer 网络，通过在空间和时间域中采用分解的自注意力机制来建模长程相关性。其在复杂高速场景下实现了超过 30 dB 的 PSNR，达到当前最优性能，并可在仅 2.7 秒内完成快速、端到端的推理，同时支持可变掩码和输入尺寸而无需微调。

ABSTRACT

Video snapshot compressive imaging (SCI) captures multiple sequential video frames by a single measurement using the idea of computational imaging. The underlying principle is to modulate high-speed frames through different masks and these modulated frames are summed to a single measurement captured by a low-speed 2D sensor (dubbed optical encoder); following this, algorithms are employed to reconstruct the desired high-speed frames (dubbed software decoder) if needed. In this paper, we consider the reconstruction algorithm in video SCI, i.e., recovering a series of video frames from a compressed measurement. Specifically, we propose a Spatial-Temporal transFormer (STFormer) to exploit the correlation in both spatial and temporal domains. STFormer network is composed of a token generation block, a video reconstruction block, and these two blocks are connected by a series of STFormer blocks. Each STFormer block consists of a spatial self-attention branch, a temporal self-attention branch and the outputs of these two branches are integrated by a fusion network. Extensive results on both simulated and real data demonstrate the state-of-the-art performance of STFormer. The code and models are publicly available at https://github.com/ucaswangls/STFormer.git

研究动机与目标

• 为解决传统感受野受限的 CNN 在视频 SCI 重建中难以有效捕捉长期时间相关性的问题。
• 克服现有深度学习模型在输入尺寸或掩码发生变化时需重新训练的僵化缺陷。
• 开发一种快速、端到端的重建网络，在不同压缩率和复杂运动场景下均保持高保真度。
• 提升真实彩色视频 SCI 数据的重建质量，该任务仍具挑战性，原因在于噪声和复杂的光谱串扰。

提出的方法

• STFormer 网络由标记生成模块、视频重建模块以及多个通过融合网络集成空间与时间自注意力的 STFormer 模块构成。
• 每个 STFormer 模块分别在空间窗口和时间序列上应用局部自注意力，实现高效建模长程依赖，同时降低计算量。
• 空间与时间注意力分支通过可学习的融合网络进行分解与结合，以保持空间与时间的一致性。
• 该架构对输入尺寸和掩码模式具有高度灵活性，当这些参数变化时无需重新训练。
• 模型通过端到端训练，使用重建损失最小化预测帧与真实帧之间的像素级误差。
• 该方法支持灰度与彩色视频 SCI，使用拜耳滤波测量数据，并与真实世界系统（如 CACTI）兼容。

实验结果

研究问题

• RQ1 基于 Transformer 的架构能否有效建模视频 SCI 重建中的长程时空相关性，超越基于 CNN 的方法？
• RQ2 所提出的 STFormer 网络能否在不重新训练的前提下，保持不同压缩率和输入尺寸下的高重建质量？
• RQ3 与现有 SOTA 方法相比，STFormer 在具有复杂运动与噪声的真实彩色视频 SCI 数据上的表现如何？
• RQ4 STFormer 架构能否在保持高 PSNR 和结构保真度的同时实现快速推理速度？

主要发现

• STFormer 在复杂高速场景（如 Duomino 和 Water Balloon）中实现了 SOTA 重建性能，PSNR 超过 30 dB，显著优于先前方法。
• 在真实彩色视频数据（如 Hammer、Ball Rotate、Duomino）上，STFormer 生成的图像边缘更清晰，伪影更少，尤其在高运动区域优于 GAP-TV、DeSCI 和 PnP-FastDVDnet。
• 该方法仅用 2.7 秒即可完成视频帧重建，远快于 DeSCI（超过 3 小时）和 PnP-FastDVDnet（3.4 分钟），展现出卓越的推理速度。
• STFormer 是首个无需重新训练即可支持可变输入尺寸和掩码的端到端深度学习视频 SCI 模型，显著提升了实际部署的可行性。
• 该模型在压缩率从 B=10 到 B=50 的范围内均保持高性能，即使在高斯压缩下仍能清晰保留运动细节。
• 在新搭建的西湖大学视频 SCI 系统上的补充实验结果证实，STFormer 在不同硬件配置和掩码模式下均表现出强鲁棒性与泛化能力。