分享自:

WaterDiffusion:一种用于联合水下显著性检测与视觉复原的先验展开扩散模型

期刊:Proceedings of the AAAI Conference on Artificial Intelligence

一、 研究概览与发表信息

本文的题目为“WaterDiffusion: Learning a Prior-Involved Unrolling Diffusion for Joint Underwater Saliency Detection and Visual Restoration”。该研究由Laibin Chang1, Yunke Wang2, Longxiang Deng1, Bo Du1*, Chang Xu2* 共同完成。主要作者来自两个机构:1 武汉大学计算机学院、国家多媒体软件工程技术研究中心、人工智能研究院、多媒体与网络通信工程湖北省重点实验室(中国);2 悉尼大学计算机学院(澳大利亚)。此项研究成果发表于第三十九届人工智能促进协会人工智能会议(AAAI-25),作为该会议录用的学术论文之一。

二、 学术背景与研究目的

本研究隶属于计算机视觉领域,具体聚焦于水下视觉任务。水下显著性目标检测(Underwater Salient Object Detection, USOD)在基于视觉的海洋探索任务中扮演着关键角色,其目标是从复杂的水下场景中快速识别出视觉上最引人注目的物体。然而,由于水下环境的特殊性——例如光的选择性吸收和介质散射——导致捕获的原始水下图像普遍存在颜色偏差、低对比度、细节模糊等严重的质量退化问题。这种退化使得直接从原始图像中精确定位和分割显著物体变得极具挑战性。

传统的解决思路是一种两阶段(串联式)方法:首先使用水下图像增强(Underwater Image Enhancement, UIE)技术恢复图像视觉质量,然后再对增强后的图像进行显著性检测。然而,这种分离的流程存在一个关键缺陷:它忽略了视觉复原过程本身可能对显著性检测带来的潜在积极影响,两者被割裂处理,未能形成协同效应。

与此同时,扩散模型(Diffusion Models)在图像生成与修复任务中展现出卓越性能,尤其在噪声鲁棒性方面表现优异。尽管已有一些工作尝试利用扩散模型进行水下图像增强,但它们通常仅专注于全局的图像去噪与增强,并未与显著性检测这一局部感知任务相结合。迄今为止,尚未有研究探索基于扩散模型的、能协同完成水下图像复原与显著性检测的联合框架。

基于以上分析,本研究的主要目标是:提出一个统一的、能够协同执行水下显著性目标检测与视觉复原的生成模型。研究旨在解决现有两阶段方法的局限性,通过一个单一的网络同时提升图像质量和精准定位显著物体,实现两个任务的相互促进。

三、 详细研究流程与方法论

本研究提出了一种名为 WaterDiffusion 的、融合了先验知识的展开扩散模型。整个框架的核心思想是,将显著性检测重新定义为一个联合任务:在复原清晰图像的同时,迭代地优化显著性掩码。模型以原始退化水下图像为输入,通过一个精心设计的扩散过程,并行输出增强后的清晰图像和精确的显著性物体掩码。

研究流程包含以下几个核心模块和步骤:

  1. 动态显著性感知联合扩散(Dynamic Saliency-Aware Joint Diffusion)

    • 这是模型的基线框架,对传统扩散模型进行了修订。与标准扩散模型仅以退化图像y为条件不同,WaterDiffusion 引入了一个动态的显著性掩码m_t作为额外的条件。在反向扩散过程中,模型不仅学习去噪以恢复清晰图像x_{t-1},同时还学习预测并细化当前时间步的显著性掩码。其联合反向过程定义为 pθ(x_{t-1} | x_t, y, m_t),这意味着在每一步去噪时,模型都同时参考了原始输入图像和当前估计的显著性掩码。
    • 该框架由多个自注意力联合扩散(Self-Attention Joint Diffusion, SAJD)模块堆叠而成。SAJD 模块(如图3所示)的核心作用是将编码的时间步信息t与前一层的特征f_in进行级联,并利用自注意力机制来确保生成内容的高保真度。这些模块构成了噪声与掩码预测器(类似于U-Net)的主干。
  2. 介质传输先验辅助显著性定位(Medium Transmission Prior, MTP)

    • 为了应对严重退化图像中显著物体难以定位的问题,研究者创新性地将水下光学传输图(Optical Transmission Map)的先验知识集成到解码器中。水下成像物理模型表明,传输图描述了光在水下传播过程中的衰减情况,能够反映图像中显著区域与非显著区域在退化程度上的差异。该模块首先利用暗通道先验等方法从输入图像y中估算出传输图t_m(x)
    • 随后,MTP模块(如图4a所示)利用计算出的传输图t_m(x)作为一个特征选择器,对SAJD模块提取的中间特征f_in进行加权。具体操作如公式 f_out(x) = f_in(x) + f_in(x) × t_m(x) 所示,传输图中值较高的像素(通常对应衰减较小的前景/显著区域)会获得更高的权重,从而引导网络更加关注潜在的显著物体区域,辅助其精确定位。
  3. 门控引导的特征选择以提升泛化能力(Gate-Guided Feature Selection, GFS)

    • 水下图像存在明显的域间差异,即不同类型的退化(如蓝绿色偏、雾状模糊等)混杂。为了提升模型对不同退化类型的泛化能力,研究者设计了GFS模块(如图5所示)。
    • 该模块的核心思想是自适应地选择使用实例归一化(Instance Normalization)后的特征还是原始特征。实例归一化通过减去通道均值并除以标准差,有助于提取对特定风格(退化类型)不变的特征表示,增强泛化性,但可能牺牲部分细节信息。
    • GFS模块首先通过全局平均池化和多层感知机生成一个概率向量γ,然后通过设定阈值(如与γ的均值比较)将其转换为二值开关指示器θ。对于每个特征通道,若θ(i)=1,则选择原始特征;若θ(i)=0,则选择归一化后的特征。最终输出特征f_out(x)是这两种特征根据θ进行通道级选择的加权和。这种设计使模型能够根据输入特征自身的情况,灵活地平衡特征的辨别力与泛化能力。
  4. 基于半二次分裂的展开式采样优化(Unrolling Sampling with Half-Quadratic Splitting)

    • 在联合采样的过程中,模型高度依赖于精炼的掩码信息和传输图特性。研究者设计了一种迭代优化策略,将掩码引导的修改与扩散采样步骤结合起来。
    • 他们首先定义了一个掩码引导的修改函数 ψ(x_t) = t_m × m_t × x_t + (1 - m_t) × x_t。该函数利用当前估计的传输图t_m和显著性掩码m_t对生成的中间样本x_t进行加权,突出显著区域。
    • 接着,通过最小化一个包含图像-掩码正则项的能量函数,将清晰图像x和掩码m的估计转化为一个约束优化问题。采用半二次分裂(Half-Quadratic Splitting)算法,该问题被分解为三个子问题交替求解:更新辅助变量φ(与清晰图像相关)、更新辅助变量ϕ(与掩码相关)、以及通过WaterDiffusion模型的反向采样步骤更新联合状态<x|m>。这种展开式迭代采样过程使得掩码信息能够被逐步细化,并与图像复原过程深度耦合,相互促进。
  5. 损失函数与训练

    • 模型的训练目标是同时预测噪声ϵ̂_t和显著性掩码m̂_t。总损失函数L_total由两部分组成:扩散噪声预测损失 L_noise = E ||ϵ̂_t - ϵ_t||^2 和显著性掩码细化损失 L_mask = E ||m̂_t - m_gt||^2,其中m_gt是真实的显著性掩码。两项损失通过权重ζ=0.1进行加权求和。
  6. 实验设置与评估

    • 数据集:研究使用了三个公开的UIE数据集(UIEB, UWScene, LSUI)进行视觉复原评估,以及三个USOD数据集(UFO-120, SUIM, USOD10K)进行显著性检测评估。其中,来自SUIM和USOD10K的3019张带掩码的图像用于训练,其余用于测试。
    • 评估指标:视觉复原采用UIF、SSIM、PSNR指标;显著性检测采用S-measure、E-measure、F-measure、MAE指标以及PR曲线。
    • 对比方法:与5种基于深度学习的UIE方法和8种基于深度学习的USOD方法进行了全面对比。
    • 消融实验:系统地评估了迭代掩码细化(IMR)介质传输先验(MTP)门控引导特征选择(GFS) 三个核心模块的必要性和有效性。

四、 主要研究结果与分析

  1. 定性结果

    • 视觉复原方面(图6):在UIEB、UWScene和LSUI数据集上的可视化对比表明,尽管多数对比方法都能取得一定的增强效果,但WaterDiffusion复原出的图像在色彩饱和度、内容结构相似性和全局对比度上最接近真实参考图像,视觉效果更为自然、清晰。
    • 显著性检测方面(图7):在UFO-120、SUIM和USOD10K数据集上的结果显示,相较于其他USOD方法,WaterDiffusion在显著物体的定位准确性和细节分割的精细度上表现出更优的鲁棒性。其生成的掩码边界更清晰,对复杂背景的抑制能力更强。
  2. 定量结果

    • 视觉复原定量评估(表1):在三个UIE数据集上,WaterDiffusion在UIF、SSIM、PSNR指标上取得了最佳或接近最佳的综合性能。特别是在UWScene数据集上,PSNR达到了25.751 dB,显著优于多数对比方法。这表明模型在像素级保真度和结构相似性方面都达到了先进水平。研究者分析,这得益于扩散模型强大的生成能力,以及GFS和MTP模块对特征通道注意力及水下传输先验的有效利用。
    • 显著性检测定量评估(表2):在三个USOD数据集上,WaterDiffusion在F-measure、E-measure、S-measure和MAE等核心指标上均表现优异。例如,在T-SUIM数据集上,其F-measure达到0.862,E-measure达到0.915,MAE低至0.056,全面领先于所列的对比方法。这证明了联合学习框架能够有效利用复原信息提升检测精度。
    • PR曲线分析(图8):在UFO-120和T-SUIM数据集上绘制的精确率-召回率曲线显示,WaterDiffusion的曲线下面积(AUC)最大,进一步从整体性能上验证了其优越性。
  3. 模型效率评估

    • 推理时间(图9):研究将WaterDiffusion的推理时间(0.196秒)与任何一组串联的UIE+USOD方法的推理时间之和进行了对比。结果显示,WaterDiffusion的单一模型推理时间小于最快的两阶段方法组合(如MLLE + TC-USOD的0.201秒),证明了其协同设计在效率上的优势——用一个模型的时间完成了两个模型的任务。
    • 参数量与计算量(表3):虽然WaterDiffusion需要同时处理两个任务,但其参数量(55.49M)和浮点运算量(169.56G)与部分单一任务的先进模型相比仍具有竞争力,甚至更少,体现了模型设计的效率。
  4. 消融实验结果分析

    • 迭代掩码细化(IMR):表4显示,在UIE任务中,不使用IMR(即“-w/o” IMR)会导致UIF、SSIM、PSNR指标全面下降,证实了在扩散过程中迭代优化掩码对于提升图像复原质量是必要的。而对于USOD任务,没有掩码细化则无法进行。
    • 介质传输先验(MTP):表5和表6分别显示,移除MTP模块后,无论是UIE任务还是USOD任务的各项指标均出现显著下滑。这强有力地证明了将水下光学传输先验知识融入扩散过程,对于协同提升图像复原和显著性检测的性能具有关键作用。
    • 门控引导特征选择(GFS):图10展示了训练过程中,带有GFS模块的模型最终在SSIM和PSNR指标上能够超越不带GFS模块的版本。尽管在训练初期收敛速度较慢且有波动,但GFS最终帮助模型学习到了更具泛化能力的特征表示,从而获得更高的性能上限。

五、 研究结论与价值

本研究成功提出并验证了WaterDiffusion,这是一个开创性的、融合先验知识的展开扩散模型,用于协同完成水下显著性目标检测与视觉复原。研究的主要结论是:通过将动态显著性掩码嵌入扩散过程作为潜在特征,并引入水下介质传输先验来辅助显著性定位,可以构建一个统一的生成框架。该框架利用门控机制自适应选择特征以提升泛化能力,并采用基于半二次分裂的展开式采样来迭代优化输出结果。

此项研究的科学价值在于: 1. 方法创新:首次将扩散模型应用于水下显著性检测任务,并创新性地提出了一个联合学习范式,打破了传统两阶段串行处理的壁垒,实现了两个任务的端到端协同优化与相互促进。 2. 理论贡献:提出了动态显著性感知扩散、介质传输先验集成、门控特征选择等新颖模块,为如何将领域先验知识(水下物理模型)与前沿生成模型(扩散模型)相结合,以解决特定领域的联合视觉任务提供了新的思路和技术路径。 3. 性能突破:在多个公开基准数据集上的实验表明,WaterDiffusion在视觉复原和显著性检测两个任务上均达到了最先进的性能,且推理效率优于串联的两阶段方法。

其应用价值显著,可广泛应用于需要同时进行水下图像质量提升和目标识别/分割的场合,如自主水下航行器(AUV)导航、海洋生物监测、水下设施巡检、海底考古等,为海洋视觉探索提供了更强大、更高效的智能工具。

六、 研究亮点

  1. 首创性:这是首个基于扩散模型进行水下显著性检测的研究,也是首个将水下显著性检测与视觉复原在一个统一的扩散模型框架内进行联合求解的工作。
  2. 巧妙的先验融合:创新性地将水下光学传输图这一物理先验引入解码网络,作为特征加权器,有效地利用了水下退化在空间上的不均匀性来引导显著性定位,是领域知识与深度学习模型深度融合的成功范例。
  3. 动态协同机制:提出的动态显著性感知联合扩散模型,使显著性掩码不再是固定的输入或单独网络的输出,而是在扩散过程中被迭代估计和细化的动态条件,与图像复原过程形成了紧密的闭环反馈。
  4. 高效的泛化设计:门控引导的特征选择模块通过简单的二值开关,自适应地在保持特征辨别力和提升模型泛化性之间做出权衡,有效应对了水下图像退化类型多样性的挑战。
  5. 全面的验证:研究不仅进行了充分的定量与定性对比实验,证明了性能优势,还通过详尽的消融实验和效率分析,清晰地揭示了各个核心组件的贡献以及模型整体的高效性。
上述解读依据用户上传的学术文献,如有不准确或可能侵权之处请联系本站站长:admin@fmread.com