[论文解读] Learning A Physical-aware Diffusion Model Based on Transformer for Underwater Image Enhancement

Chen Zhao, Chenyu Dong+2 more|arXiv (Cornell University)|Mar 3, 2024

Underwater Acoustics Research被引用 10

一句话总结

PA-Diff 引入一个物理引导的扩散框架用于水下图像增强，结合 Physics Prior Generation 分支与 Physics-aware Diffusion Transformer，以充分利用水下成像物理与扩散建模。

ABSTRACT

Underwater visuals undergo various complex degradations, inevitably influencing the efficiency of underwater vision tasks. Recently, diffusion models were employed to underwater image enhancement (UIE) tasks, and gained SOTA performance. However, these methods fail to consider the physical properties and underwater imaging mechanisms in the diffusion process, limiting information completion capacity of diffusion models. In this paper, we introduce a novel UIE framework, named PA-Diff, designed to exploiting the knowledge of physics to guide the diffusion process. PA-Diff consists of Physics Prior Generation (PPG) Branch, Implicit Neural Reconstruction (INR) Branch, and Physics-aware Diffusion Transformer (PDT) Branch. Our designed PPG branch aims to produce the prior knowledge of physics. With utilizing the physics prior knowledge to guide the diffusion process, PDT branch can obtain underwater-aware ability and model the complex distribution in real-world underwater scenes. INR Branch can learn robust feature representations from diverse underwater image via implicit neural representation, which reduces the difficulty of restoration for PDT branch. Extensive experiments prove that our method achieves best performance on UIE tasks.

研究动机与目标

• 通过将水下物理成像机制纳入扩散模型，推动 UIE 的改进。
• 提出一个双分支架构，将物理先验与基于扩散的重建耦合。
• 通过一个物理感知变换器实现长期依赖建模和区域感知增强。
• 在真实水下数据集上展示最先进的性能，并通过消融实验验证贡献。

提出的方法

• 两分支 PA-Diff 架构：物理先验生成（PPG）分支和物理感知扩散变换器（PDT）分支。
• PPG 通过修改后的 Koschmieder 光扫描模型生成透射图和全局背景光；物理先验通过重构损失引导扩散。
• PDT 在水下输入和物理先验条件下使用扩散过程（DDPM 框架）；包含物理感知自注意力（PA-SA）与多尺度动态前馈网络（MS-FFN）。
• PA-SA 结合透射和光照先验调制特征，生成用于长程扩散建模的透射感知和光照感知特征。
• MS-FFN 采用多尺度深度卷积以扩展感受野并聚合局部特征。
• 训练通过使条件扩散过程中的估计噪声与真实噪声匹配的扩散损失来优化。

实验结果

研究问题

• RQ1 如何将水下物理成像知识整合到基于扩散的 UIE 框架？
• RQ2 引入透射图和背景光先验是否能提升水下图像的扩散重建？
• RQ3 具有物理感知注意力的基于变换器的扩散模块是否能更好地建模水下场景的长期依赖？
• RQ4 每个物理派生分量（透射图、背景光、MS-FFN）对 UIE 性能的贡献是什么？

主要发现

• PA-Diff 在 UIE 基准（UIEBD 和 LSUI）上实现了最先进的感知和保真指标。
• 消融研究表明透射先验（T）比单独的背景光（B）提供更强的提升，且两者结合并与 MS-FFN 和 PA-SA 完整集成时达到最佳性能。
• 在 LSUI 上，所提出的方法在 PSA 提升方面优于若干基线，且 FID 较低、SSIM 较高。
• 消融结果表明生成的物理先验有益，但确切的先验精度会影响最终增益，表明仍有改进物理先验估计的空间。