分享自:

UDAFormer: 基于双注意力变换器的水下图像增强

期刊:computers & graphics

水下图像增强领域的研究正面临着色偏与低对比度等长期挑战,这些问题源于水下环境中光波长依赖性的散射与吸收特性,严重阻碍了高级计算机视觉任务(如目标跟踪、辅助驾驶系统和水下机器人)的实际应用。虽然基于卷积神经网络(CNN)的方法已取得显著进展,但其固有的局限性,例如有限的感受野导致难以建模长距离像素依赖关系,以及推理时卷积滤波器权重静态不变而无法灵活适应输入内容,制约了其在处理复杂、不均匀的水下图像退化时的性能。与此同时,在自然语言处理领域取得巨大成功的Transformer架构,凭借其强大的自注意力机制(Self-Attention Mechanism),为计算机视觉任务开辟了新路径。然而,现有的视觉Transformer模型多侧重于计算图像块之间的注意力权重,这对于处理退化程度严重且分布不均的水下图像来说仍显不足。

在此背景下,研究人员 Zhen Shen, Haiyong Xu*, Ting Luo, Yang Song, Zhouyan He(来自宁波大学数学与统计学院和科学技术学院)提出了一种名为UDAformer的新型水下图像增强方法,该方法基于双注意力Transformer。该研究成果以题为《UDAformer: Underwater Image Enhancement Based on Dual Attention Transformer》发表于Computers & Graphics期刊第111卷(2023年,第77-88页)。UDAformer的核心在于其创新的双注意力Transformer块(DATB),它融合了通道自注意力Transformer(CSAT)与像素自注意力Transformer(PSAT),旨在更高效地编码和解码水下图像特征,以应对水下成像的独特挑战。

研究背景与目标 水下成像受到水体对光的选择性吸收和散射效应影响,导致图像普遍存在蓝/绿色调、细节模糊、对比度低等问题。这些不均匀的退化以及特定颜色通道(尤其是红色通道)的信息严重损失,使得传统的单一图像处理或标准CNN模型难以获得理想效果。Transformer的自注意力机制虽然能捕获全局依赖关系,但标准的视觉Transformer在处理水下图像时,无法兼顾局部细节恢复与通道间重要信息的差异化处理。因此,本研究的目标是设计一个能够同时建模局部像素关系全局通道依赖性的网络架构,以更精准地恢复水下图像的色彩、对比度和纹理细节。

详细研究方法与工作流程 UDAformer的整体架构采用了类似U-Net的编码器-解码器结构,但其核心构建块并非传统卷积,而是专门设计的DATB模块。整个工作流程可以概括为以下步骤:

  1. 整体架构设计:模型接收一个水下图像I ∈ R^(3×H×W)作为输入。首先,通过一个卷积层提取浅层特征F_l ∈ R^(C×H×W)。随后,特征经过三个编码器层逐步下采样,得到深层特征F_d ∈ R^(8C×H/8×W/8)。编码过程结束后,特征进入一个瓶颈层。解码器则通过三个解码器层将低分辨率特征逐步上采样恢复至高分辨率,并通过跳跃连接融合来自编码器对应层级的特征,以辅助恢复精细结构。最终,通过一个输出投影卷积层将特征解码为残差图K,并与原始输入图像相加,得到增强后的清晰图像J。这种设计灵感来源于水下成像物理模型(J(x) = K(x) + I(x)),将增强问题转化为学习一个残差映射,避免了直接估计复杂的介质传输图t和背景光A,减少了累积误差。

  2. 核心模块:双注意力Transformer块(DATB) DATB是UDAformer的创新核心,它并非简单地串联CSAT和PSAT,而是通过一个可学习的系数α将二者自适应融合。DATB根据其在网络中的位置(浅层/瓶颈层)有两种形式:W-DATB(用于第一个编码器层和最后一个解码器层)和G-DATB(用于瓶颈层)。

    • 通道自注意力Transformer(CSAT):针对水下图像中不同颜色通道(如红、绿、蓝)信息损失程度不同的问题而设计。在CSAT中,首先将空间维度的全局像素信息视为通道特征,然后计算通道特征之间的自注意力权重。具体而言,输入特征经过层归一化(LN)后,分别通过并行的1×1点卷积和3×3深度可分离卷积生成查询(Query)、键(Key)和值(Value)向量。注意力权重通过Query和Key的点积计算,再与Value加权求和。CSAT能够以线性计算复杂度捕获全局通道间依赖关系,使网络更关注信息丰富的通道特征,从而有效校正色偏。
    • 像素自注意力Transformer(PSAT)与移位窗口PSAT(SW-PSAT):为了建模局部像素关系以恢复纹理细节,并应对水下图像退化不均匀的特点,设计了PSAT。然而,直接在原始分辨率上计算全局像素自注意力的计算复杂度会随空间分辨率呈二次方增长。为此,论文引入了移位窗口(Shifted Window)技术,形成了SW-PSAT。在常规窗口中,特征图被均匀划分为不重叠的局部窗口,自注意力计算仅限于每个窗口内部,这大大降低了计算成本。为了引入跨窗口的连接,网络交替使用规则窗口划分和经过位移的窗口划分策略,使得相邻窗口的像素在计算中能够交互。在网络的瓶颈层,由于特征图尺寸已缩减至输入的1/8,计算复杂度可控,因此直接使用全局PSAT(即无窗口划分)来捕获全局像素依赖,以防止图像出现过增强或欠增强现象。
  3. 数据与实验设置

    • 训练与测试数据集:研究使用UIEB数据集进行训练和测试,该数据集包含890对真实水下图像及其对应的参考图像。其中800张用于训练,90张用于测试。训练时,所有图像被调整为256×256像素。
    • 对比方法与评估指标:研究将UDAformer与10种先进方法进行了全面比较,包括5种传统方法(如UDCP、UIBLA)、4种基于CNN的方法(如UWCNN、Water-Net、Ucolor、FUnIE-GAN)以及1种基于Transformer的方法(U-Shape)。评估指标分为两类:
      • 全参考评价:在包含参考图像的Test-L504(选自LUSI数据集)和Test-U90(选自UIEB测试集)上,使用峰值信噪比(PSNR)和结构相似性指数(SSIM)。
      • 无参考评价:在无对应参考图像的Test-C60SQUIDU45数据集上,使用水下图像质量度量(UIQM)、水下颜色图像质量评估(UCIQE)和自然图像质量评估器(NIQE)。
    • 实现细节:模型使用PyTorch框架实现,在NVIDIA RTX 3090 GPU上使用Adam优化器(学习率1e-4)和L2损失函数训练了500个epoch,耗时约20小时。

主要研究结果 实验结果表明,UDAformer在定性和定量评估上均超越了之前的先进方法。 1. 全参考评估结果:在Test-L504和Test-U90数据集上,UDAformer取得了最优的PSNR和SSIM值。在Test-L504上,PSNR达到23.01,SSIM达到0.86,优于当时最佳方法Ucolor(PSNR 21.54, SSIM 0.86)。在Test-U90上,PSNR达到23.48,SSIM达到0.92,同样显著优于Ucolor(PSNR 21.29, SSIM 0.88)。这证明了UDAformer生成的图像在像素级和结构上与参考图像高度一致。 2. 无参考评估结果:在更具挑战性的无参考数据集上,UDAformer也表现出色。虽然在UCIQE指标上HE方法得分最高,但论文指出UCIQE倾向于给过度饱和、出现伪色的图像打高分,因此不能准确反映视觉效果。综合UIQM和NIQE指标来看,UDAformer在SQUID、U45和Test-C60上都取得了具有竞争力的分数,表明其增强图像具有更好的自然感知质量。 3. 视觉对比:视觉对比图清晰地展示了UDAformer的优势。相比于其他方法,UDAformer能更有效地校正颜色(如将蓝/绿色调恢复至自然色彩),更显著地去除雾状模糊,提升整体对比度,并恢复出更清晰的细节纹理。而其他方法或存在色偏校正不足(如传统方法),或存在去雾效果不佳、图像模糊(如部分CNN方法),或产生过度增强与人工伪影(如HE)。 4. 消融实验(Ablation Study):为验证各个模块的有效性,研究进行了消融实验。结果表明: * 仅使用SW-PSAT的网络无法有效校正色偏。 * 仅使用CSAT的网络虽然能校正颜色,但去雾效果和SSIM值不如完整模型,因为缺乏对局部像素细节的建模。 * 在瓶颈层不使用全局PSAT(即W/O G-Fusion)的模型,其PSNR和SSIM值低于完整模型,视觉效果也不够理想。 * 完整的UDAformer模型取得了最佳性能,证明了CSAT与SW-PSAT融合以及瓶颈层全局PSAT设计的必要性。 5. 扩展实验与泛化能力: * 跨数据集训练:当使用LUSI数据集训练并在不同测试集上评估时,UDAformer同样表现出优于对比方法U-Shape的性能,证明了其鲁棒性。 * 水下视频增强:将UDAformer应用于真实水下视频序列,结果表明其能有效恢复物体颜色、提升视频对比度,并且在相邻帧间保持了增强效果的一致性。 * 泛化至其他场景:研究还将UDAformer扩展到雾天图像和低光照图像增强任务。初步实验显示,在处理这些具有类似退化(模糊、低对比度)的图像时,UDAformer也取得了一定的积极效果,展现了其良好的泛化潜力。 6. 失败案例分析:论文也坦诚地指出了当前模型的局限性。对于前景明亮、背景大面积黑暗的图像,增强效果不佳,黑暗背景区域会出现严重失真。作者分析认为,主要原因在于训练数据集中此类图像样本较少,导致监督学习模型难以处理。

研究结论与价值 本研究提出并验证了UDAformer,一种基于双注意力Transformer的水下图像增强网络。其主要结论是:通过巧妙地融合通道自注意力(CSAT)和移位窗口像素自注意力(SW-PSAT)构成的DATB模块,能够有效解决水下图像不均匀退化和颜色通道信息丢失的核心问题。CSAT使网络聚焦于信息丰富的颜色通道,SW-PSAT则专注于恢复局部纹理细节并抑制不均匀退化,而瓶颈层的全局PSAT进一步确保了整体增强的协调性。这种设计使UDAformer在恢复图像色彩、对比度和细节方面均达到了领先水平。

研究亮点 1. 创新的双注意力机制:针对水下图像的特性,首次提出将通道自注意力像素自注意力融合于统一的Transformer块中,兼顾了全局色彩校正与局部细节恢复。 2. 高效的移位窗口设计:将Swin Transformer中的移位窗口思想应用于像素级自注意力计算(SW-PSAT),在保持强大局部建模能力的同时,显著降低了计算复杂度,使其适用于高分辨率图像增强任务。 3. 针对性的架构设计:根据网络不同层次的特征分辨率,差异化地使用W-DATB(局部窗口)和G-DATB(全局注意力),实现了计算效率与模型性能的平衡。 4. 基于物理模型的残差学习:将网络输出设计为学习水下成像模型的残差项,简化了学习目标,避免了直接估计复杂物理参数的困难,提升了模型的稳定性和泛化能力。 5. 全面的实验验证:不仅在多个标准数据集上进行了充分的定量与定性对比,还通过消融实验、跨数据集测试、视频应用和跨任务泛化实验,系统性地验证了模型各个组件的有效性和整体性能的优越性。

其他有价值内容 论文在引言和实验部分对相关研究工作进行了梳理,涵盖了从传统方法、基于CNN的方法到视觉Transformer方法的发展脉络,为读者提供了清晰的领域背景。此外,作者公开了代码(GitHub链接),促进了研究的可复现性和后续发展。同时,论文对失败案例的讨论体现了研究的严谨性,并为未来工作指明了方向,例如需要收集更多样化的训练数据以处理极端光照条件下的水下图像。

上述解读依据用户上传的学术文献,如有不准确或可能侵权之处请联系本站站长:admin@fmread.com