实时4K图像修复:TSFormer 以 3.38M 参数实现 40fps高性能运行,推动超高分辨率图像修复技术落地!

大模型向量数据库机器学习

点击下方卡片,关注「集智书童」公众号

点击加入👉「集智书童」交流群

picture.image

picture.image

picture.image

picture.image

picture.image

想要了解更多:

前沿AI视觉感知全栈知识👉「分类、检测、分割、关键点、车道线检测、3D视觉(分割、检测)、多模态、目标跟踪、NerF」

行业技术方案 👉「AI安防、AI医疗、AI自动驾驶」

AI模型部署落地实战 👉「CUDA、TensorRT、NCNN、OpenVINO、MNN、ONNXRuntime以及地平线框架」

欢迎扫描上方二维码,加入「 集智书童-知识星球 」,日常分享论文、学习笔记、问题解决方案、部署方案以及全栈式答疑,期待交流!

免责声明

凡本公众号注明“来源:XXX(非集智书童)”的作品,均转载自其它媒体,版权归原作者所有,如有侵权请联系我们删除,谢谢。

picture.image

超高分辨率(UHD)图像修复对于需要出色视觉保真度的应用至关重要,然而现有方法往往在修复质量和效率之间存在权衡,限制了其实际部署。

在这篇论文中,作者提出了一种名为TSFormer 的全功能框架,该框架将可信 学习与S 化简相结合,以提高UHD图像修复的泛化能力和计算效率。关键在于,模型中允许的 Token 移动量非常有限。

为了有效地过滤 Token ,作者使用基于随机矩阵理论的Min-方法来量化 Token 的不确定性,从而提高模型的鲁棒性。作者的模型可以在实时(40fps)运行4K图像(3.38 M参数)。

大量实验表明,TSFormer在实现最先进的修复质量的同时,提高了泛化能力和降低了计算需求。

此外,作者的 Token 过滤方法可以应用于其他图像修复模型,以有效地加速推理并保持性能。

1 Introduction

超高分辨率(UHD)图像修复对于需要高分辨率图像质量的各种应用至关重要,包括医学影像、视频流媒体和数字监控。由于UHD图像拥有数百万像素,在资源有限的情况下进行处理是一项巨大的挑战。

目前,有一些深度学习框架[34, 39]可以在消费级GPU上处理UHD图像。然而,这些方法直接或间接地降低输入图像的分辨率,可能导致图像修复过程中丢失一些重要信息。因此,这会导致图像细节丢失,这对UHD图像尤为重要。

为了解决这个问题,作者提出了一种名为TSFormer 的轻量级和可信框架,它结合稀疏化和随机矩阵理论。TSFormer旨在保留最具信息量的特征,从而在降低计算开销的同时保持高修复质量和鲁棒性。

确实,TSFormer的第一个关键组成部分是Min-采样 ,这是一种基于概率稀疏化理论的稀疏表示技术,灵感来源于[1]中最近的发展。与传统的Top-过滤方法[18]不同,Min-采样根据概率阈值选择性地保留高置信度的特征,允许动态特征选择。这种"少就是多"的方法认识到,并非所有特征对最终输出做出相同程度的贡献,尤其是在高维数据(如UHD图像)中。然而,尽管Min-采样有效降低了模型的计算需求,但在大量 Token 中可能引入不稳定性,因为存在噪声。为了解决这个问题,TSFormer引入了一个基于随机矩阵理论的可信 机制[7, 23]。可信过滤涉及分析特征矩阵的特征值,以确保仅保留稳健、高置信度的特征。如图1(a)所示,Min-采样应用基于概率的阈值保留高置信度的特征,同时丢弃较低重要性的特征,实现有效的稀疏化。在图1(b)中,基于随机矩阵理论的可信过滤进一步细化剩余特征,通过排除位于可信阈值(虚心圆)之外的点。这种可信驱动的特征选择改进了泛化能力,使模型在处理各种退化图像时能够可靠运行。

到目前为止,作者通过强制使用Min-并引入可信机制构建了一个Min-稀疏注意力(MSA),它能够高效、可靠地生成注意力图。MSA被集成到TSFormer的每个块中,每个块还包含频域学习和多尺度学习组件。TSFormer在多个UHD图像任务上表现出令人鼓舞的性能,并且可以在资源受限的设备(单个3090 GPU着色器,24G RAM)上实时运行UHD图像。

picture.image总之,作者的贡献如下:

  1. 作者开发了一种名为TSFormer的UHD图像恢复模型,它可以在单个GPU上实时运行4K分辨率图像,并具有强大的泛化能力。
  2. 作者设计了一种基于可信机制的 Token 过滤方法,将其集成到TSFormer中,以生成高质量的关注图。
  3. 基于可信机制的 Token 过滤方法可以在任何基于Transformer的图像恢复框架中使用,以提高模型的效率。大量实验结果证明了我国方法的有效性。

2 Related Work

UHD Image Restoration

超高分辨率(UHD)图像修复对于医学影像、视频 Stream 和数字监控等应用至关重要 [31, 36]。最近,深度学习方法通过采样和并行化技术有效地重构了UHD图像的细节和颜色。郑等人在[40]中提出了一个多引导双边上采样模型用于UHD图像去雾,通过多个引导输入增强清晰度。邓等人[6]开发了一个可分块整合网络用于UHD视频去模糊,采用多尺度整合方案来减轻运动和模糊伪影。王等人[26]提出了LLFormer,这是一个基于 Transformer 的低光增强方法,利用轴向多头自注意力以及跨层注意力融合块来提高照明和对比度。此外,有一些方法[34, 39]可以通过构建轻量级模型和查找表实时重构清晰UHD图像。

尽管这些方法可以在实时中增强UHD图像,但取样和表格查找方法并未受到可信机制的监督,这限制了模型的泛化能力。相比之下,作者提出的TSFormer采用了Min-采样进行自适应稀疏化,并基于随机矩阵理论采用可信过滤[2, 9],同时提高了特征可靠性和恢复质量,保持了较低的计算成本。

Token Sampling Technology

当前,大型语言模型(LLM)使用一些 Token 采样技术来加快推理。传统方法,如Top-过滤[8, 18],根据大小优先权特征,但使用固定阈值,这些阈值可能无法适应不同的数据分布,可能丢弃有价值的信息。概率稀疏化方法,如Min-采样[1, 42],引入了基于概率的动态阈值,更好地适应数据分布,允许灵活地保留特征。受此启发,作者引入了 Token 技术来构建高效模型。在此基础上,确保实时性能的同时,引入了一种可信机制(随机矩阵理论)来准确采样 Token [4]。

Random Matrix Theory

信任过滤 Token (trusted filtering of tokens)的难度较大,这可能导致模型推理速度显著降低。与其他可信建模方法(如贝叶斯、变量推理和 Token 分布)相比,随机矩阵理论(RMT)在速度和准确性上进行权衡。随机矩阵理论(RMT)提供了一个分析并增强高维数据特征稳定性的框架 [7, 23]。通过分析特征值分布,基于RMT的方法可以识别和保留最稳定且显著的特征,从而提高鲁棒性和泛化能力 [10, 35]。然而,将RMT集成到深度学习以进行UHD修复的研究仍然不充分,这为在不显著增加计算成本的情况下增强特征可靠性提供了机会。

3 Method

在本节中,作者介绍TSFormer,一种用于超高清图像恢复的高效Transformer网络。它是一种对称编码器和解码器结构,如图2所示。

picture.image### Preliminaries

Min-采样是一种专为大型语言模型任务设计的新颖采样方法,它动态地调整采样阈值以保留高置信度的 Token [20]。具体来说,给定一个特征块内的一个标量,Min-采样应用一个定义为特征块内最大值比例的阈值:

picture.image

其中是一个超参数,用于控制稀疏性。这个阈值通过将中低于阈值的元素设置为零来过滤低置信特征:picture.image

该方法通过作用于注意力图的稀疏化来过滤掉不重要的特征和噪声。

为了进一步优化作者的注意力机制的计算效率,作者在频域中使用快速傅里叶变换(FFT)进行注意力计算。FFT显著降低了计算注意力分数的复杂性,特别是对于高分辨率图像。

1D信号的前向傅里叶变换为:

picture.image

逆傅里叶变换的公式为:

picture.image

通过利用FFT,作者可以更高效地执行注意力操作,尤其是在处理高分辨率图像时,通过在频域处理数据。

Overall Pipeline

给定退化图像 , 其中 表示空间分辨率, 首先通过一个 卷积进行分词。 然后, 这个分词特征图被输入到一个编码-解码结构中。

在编码器中,作者使用一系列可信稀疏块(TSB)对特征图逐步下采样,每个层次都具有特定数量的TSB,用表示。每个TSB集成频域变换和最小稀疏注意力(MSA)。此外,每个TSB还包含一个 FFN (FFN),该网络通过一系列的深度卷积()和点卷积()来丰富特征表示,以跨越多个尺度。

解码器是解码器对高分辨率图像的细节和颜色进行重建的镜像操作。在相应的编码器和解码器 Level 之间引入特征融合块(FFB),以合并不同分辨率的特征。每个FFB应用了一个轻量级卷积块,带有PReLU激活,有助于在跨尺度上平滑地转换和细化特征,最终提高模型性能。

最后,作者在网络周围使用残差连接来获取恢复后的图像,如下所示:

picture.image

其中 表示网络的变换。该模型通过最小化恢复输出与真实图像 之间的 损失进行训练:

在论文中, 表示 范数。

Trusted Sparse Block (TSB)

自注意力机制不仅计算成本高,而且容易受到噪声影响,特别是在高频、细节丰富的UHD图像中。为了解决这个问题,作者设计了一个可信稀疏块(TSB)作为特征提取单元,它集成了Min-采样和可信学习。

具体而言,对于来自第l-1个块的输入特征,TSB中的编码过程可以定义如下:

picture.image

在此,LN 表示层归一化,MSA 代表1000in- 稀疏注意力机制,FFN 是 FFN 。在这里, 和 分别表示注意力层和前馈层的输出。

基于 Min- 稀疏关注的 MSA (Min- Sparse Attention)。 作者的 MSA 利用 Min- 采样动态地保留只有高概率特征,使用傅里叶变换有效地在频域计算注意力。

给定输入特征图 ,其中 、、 和 分别表示批处理、通道、高度和宽度维度,作者首先通过卷积变换获得 Query 、键 和值 的表示:

为了捕捉局部特征交互,作者将和划分为 patches,并对每个 patches 应用快速傅里叶变换(FFT),将其转换为频域:

picture.image

在频域中,注意力分数通过将和进行元素乘法计算得到,这显著降低了与空间域操作相比的计算复杂度。然后,应用逆FFT(IFFT)将结果转换回空间域:

接下来,Min-采样应用于向量以过滤低置信度的注意力分数。定义一个阈值为每行最大值的比例:

picture.image

最小操作保留超过此阈值的元素在中,通过设置较低的分数为零来稀疏化注意力图:

picture.image

随机矩阵理论下的可信学习

为了为建立一个可信元素过滤方法,作者将视为一个随机矩阵进行研究。首先,作者计算随机矩阵的谱密度。然后,通过平均池化层和 sigmoid 函数使用谱密度生成一个概率权重标量。最后,将乘以阈值动态调整阈值。这可以表示如下:

picture.image

如果 更加随机,它可能包含更多的噪声,因此需要更高的阈值 。需要注意的是,可信机制可能导致更多的计算开销。为了缓解这个问题, 在计算谱密度之前进行了下采样(双线性插值)。最终的稀疏注意力输出计算如下:

picture.image

FFN 遵循稀疏注意力机制之后, FFN (FFN)用于进一步优化特征。这个模块包含一系列卷积来捕捉和增强多尺度信息。对于输入 ,定义 FFN (FFN)如下:

picture.image

是 GELU 激活, 代表逐点卷积, 是逐个卷积。这种 gating 机制选择性地优化特征表示,确保只有最具有信息性的元素被传播。

Feature Fusion Blocks (FFB)

特征融合块(FFB)旨在将编码器和解码器不同层级的多种尺度特征进行集成。通过动态调整每个特征图的贡献,FFB使模型能够有效地捕捉到细粒度细节和全局上下文。

给定两个来自网络不同阶段的输入特征图和,FFB执行以下操作:

picture.image

其中表示和的拼接。

picture.image

其中, 和 分别表示对 和 进行通道加权。这种选择性融合允许 FFB 根据任务的时空需求,优先从编码器或解码器中获取信息。

4 Experiments and Analysis

作者在五个超高清图像恢复任务上(包括低光增强、去雾、去模糊、除雪和脱雨)与最先进的算法进行了性能比较。

Experimental settings.

对于UHD低光图像增强,作者使用UHD-LL数据集[13]和UHD-LOL4K [29]。为了评估去模糊能力,作者使用UHD-Blur数据集[25]。对于去雾评估,作者采用UHD-Haze数据集[25]。这些选择与前期研究[13, 25]所采用的方法一致。此外,为了评估作者的UHD图像除雪和去雨能力,作者利用了王等人[31]提出的UHD-Snow和UHD-Rain数据集。作者将PSNR [11]和SSIM [32]作为上述基准测试的评估指标。

对比方法。在作者研究中,作者将作者的方法与八种通用图像恢复(IR)技术进行了比较:SwinIR [16],Uformer [33],Restormer [38],DehazeFormer [22],Stripformer [24],FFTformer [12],和SFNet [30]。此外,作者还包括了四种超高清图像恢复(UHDIR)方法:LLFormer [26],UHD-Four [15],UHD [41],和UHDformer [27],以及UHDDIP [28]。为了进行公平比较,作者使用这些模型的官方实现进行再训练,并使用与作者的提出方法相同数量的迭代进行评估。

训练细节 在作者的模型中,初始通道 为 ,扩展比设置为 ,FFN 中通道扩展因子 r 设置为 。在训练过程中,作者使用 AdamW 优化器,批量大小为 , patch 尺寸为 ,总共进行 次迭代。初始学习率固定为 。对于数据增强,随机应用水平和垂直翻转。整个框架在 PyTorch 上运行,使用 2 个 NVIDIA RTX 3090 GPU。

Main Results

低光图像增强结果。作者在UHD低光图像增强上评估了两个训练数据集,包括UHD-LOL4k [26]和UHD-LL [14]。在表1中,TSFormer在UHD-LOL4k和UHD-LL数据集上都取得了最先进的低光图像增强结果,同时在这两个数据集上实现了最高的PSNR和SSIM值。尽管其性能强大,但TSFormer保持了轻量级的架构,只有3.38M参数,比其他高性能模型如Restormer和UHDFour更高效。这种准确性和效率的平衡表明TSFormer适合实时应用。图3说明了视觉改进,TSFormer在低光条件下有效地减少了噪声并增强了细节,超越了先前的方法。

picture.image图像去模糊结果。 作者在UHD-Blur数据集上评估图像去模糊。如表2所示,UHD-Blur数据集上的图像去模糊定量结果。TSFormer在关键指标上取得了显著改进,性能优异。具体而言,TSFormer实现了29.52 dB的PSNR,超过了UHDformer的28.82 dB,并优于其他通用模型,如Restormer、Uformer和Stripformer,它们在25-25.4 dB之间。这突显了TSFormer在UHD图像中恢复高分辨率细节的能力。作者还评估了GoPro [19]数据集的性能,并报告了表3中的结果。

picture.image图像去雾结果 。表4展示了在UHD-Haze数据集上使用训练好的模型得到的UHD-Haze的定量结果。TSFormer在PSNR和SSIM指标上取得了最高分,相较于UHDformer在UHD-Haze数据集上实现了0.77 dB的提升,同时在各视觉质量方面始终优于其他方法。TSFormer在保持与先驱模型相当的可比参数数量的同时,显著降低了LPIPS,展示了其在效率和性能方面的良好平衡。图5说明了定性结果,其中TSFormer在所有方法中产生的输出最为清晰,展示了其有效去除雾气并恢复其他方法通常被雾气掩盖的细微细节的能力。

picture.image图像去雨和除雪结果。 作者使用构建的UHD-Rain数据集评估UHD图像去雨。结果报告在表5中。TSFormer实现了最先进的性能,显著优于所有现有方法的关键指标。与先前的模型如UHDformer、Restormer和UHD-DIP相比,TSFormer在实现显著更高的保真度和感知质量的同时,还能有效处理UHD分辨率。图6说明了视觉比较,表明TSFormer在去除雨条纹的同时保留更细微的细节方面更有效。

picture.image作者在UHD-Snow数据集上实现了超高清降雪实验,结果汇总于表8,TSFormer取得了优越性能。

picture.image5 Ablation studies

在本节中,作者评估了不同采样策略对UHD图像去雾的影响。具体来说,作者比较了三种采样技术:Top-k采样,无信任机制的Min-采样,以及带有信任机制的Min-采样。

Sampling Method Comparison

为了理解每种采样方法对注意力分布和特征保留的影响,作者在表6中提供了定量结果,并在图7中给出了累积概率分布的视觉化。

picture.image表6显示,Min-采样与可信机制在去雾质量方面优于其他采样技术,实现了较高的PSNR和SSIM分数,这些分数反映了更优秀的结构保真度。图9显示了每种方法如何捕捉细节,Min-采样与可信机制主要关注高置信区域,同时有效过滤噪声。

picture.image表格2:在UHD-Blur数据集上的图像去模糊结果。TSFormer在PSNR,SSIM和LPIPS指标上实现了最先进的性能,这证明了其在UHD去模糊任务中的有效性和效率。

方法会议/期刊PSNR(向上)SSIM(向上)LPIPS(向下)参数
RestormerCVPR’2225.210.7520.37026.10M
UformerCVPR’2225.270.7520.38520.60M
StripformerECCV’2225.050.7500.37419.70M
FFTformerCVPR’2325.410.7570.37116.60M
UHDformerAAAI’2428.820.8440.2350.34M
UHDDIParxiv’2429.510.8590.2130.81M
TSFormer( ours)-29.520.8610.2033.38M

表格2展示了在UHD-Blur数据集上的图像去模糊结果。TSFormer在PSNR,SSIM和LPIPS指标上实现了最先进的性能,这证明了其在UHD去模糊任务中的有效性和效率。

Trusted Strategies Comparison

作者评估了将可信学习集成到作者的模型中以提高稳定性和鲁棒性的不同策略。

完整特征值分解(FED) 。这种策略涉及对每个注意力块进行完整的特征值分解,以识别并丢弃不稳定的块。给定一个稳定性阈值τ,只有满足最大特征值λ_{max} < τ的块被保留。尽管这种方法在隔离稳定特征方面有效,但计算成本高,导致推理时间较长。

迭代稳定性调整(ISA)。在这种方法中,稳定性阈值是根据在迭代过程中稳定区域(即,特征值方差较小的区域)的方差动态调整的:

picture.image

然而,这种迭代调整虽然具有自适应性,但由于重复计算,会显著增加运行时间,且没有带来显著的性能提升。如表9所示,采用可信机制的Min-采样在保证最低峰值信噪比和SSIM值的同时,大幅降低了运行时间,相较于其他方法具有明显优势。

6 Running Time and Application

在不考虑 I/O 操作的情况下,作者推理 4K 图像需要 40 fps。作者的方法高效的原因有两个:一是 Token 过滤机制确保了模型的稀疏性,二是特征图在可信机制下进行下采样。

图10展示了在DarkFace数据集[37]上,使用各种图像增强方法作为预处理步骤的目标检测结果的定性比较。

picture.image7 MSA of Potential

表10展示了在各种UHD模型中,通过集成MSA所实现的效率提升。通过动态关注高置信度特征,MSA在LLFormer、UHDFormer和UHDFour模型中显著减少了FLOPs和运行时间。值得注意的是,MSA在UHDFormer中的集成将FLOPs降低20%,运行时间降低25%,同时对性能的影响很小。这些改进证明了MSA在实时UHD图像修复应用中的潜力。图11说明了MSA在增强多样场景下低光图像质量的影响。通过选择性保留高置信度特征,MSA使LLFormer、UHDFormer和UHDFour模型在应对挑战性低光条件时能更好地泛化。这种能力对于照明条件较差的实际应用环境至关重要。视觉结果表明,使用MSA可以提高图像的清晰度和细节保留,突显了其在实际场景中的有效性。

picture.image8 Conclusion

作者提出了一种高效且鲁棒性强的模型,称为TSFormer,它专注于处理超高清图像。

其优势源于一种可信 Token 过滤机制,通过动态阈值和随机矩阵理论实现。

TSFormer提出了一种可信 Token 过滤机制,可应用于其他Transformer框架以提高鲁棒性。

实验结果表明,与最先进的多种UHD恢复任务相比,TSFormer在多个UHD恢复任务中实现了优越性能。

参考文献

[0]. TSFormer: A Robust Framework for Efficient UHD Image Restoration.

picture.image

扫码加入👉「集智书童」交流群

(备注: 方向+学校/公司+昵称 )

picture.image

点击下方“ 阅读原文 ”,

了解更多AI学习路上的 「武功秘籍」

0
0
0
0
关于作者
关于作者

文章

0

获赞

0

收藏

0

相关资源
字节跳动 XR 技术的探索与实践
火山引擎开发者社区技术大讲堂第二期邀请到了火山引擎 XR 技术负责人和火山引擎创作 CV 技术负责人,为大家分享字节跳动积累的前沿视觉技术及内外部的应用实践,揭秘现代炫酷的视觉效果背后的技术实现。
相关产品
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论