[Paper Review] 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 Research10 citations

TL;DR

PA-Diff introduces a physics-guided diffusion framework for underwater image enhancement, combining a Physics Prior Generation branch with a Physics-aware Diffusion Transformer to leverage underwater imaging physics and diffusion modeling.

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.

Motivation & Objective

• Motivate improved UIE by incorporating underwater physical imaging mechanisms into diffusion models.
• Propose a dual-branch architecture that couples physics priors with diffusion-based restoration.
• Enable long-range dependency modeling and region-aware enhancement through a physics-aware transformer.
• Demonstrate state-of-the-art performance on real-world underwater datasets with ablations validating contributions.

Proposed method

• Two-branch PA-Diff architecture: Physics Prior Generation (PPG) branch and Physics-aware Diffusion Transformer (PDT) branch.
• PPG generates transmission maps and global background light via a modified Koschmieder light scanning model; physics priors guide diffusion via a reconstruction loss.
• PDT uses a diffusion process (DDPM framework) conditioned on the underwater input and physics priors; includes Physics-aware Self-Attention (PA-SA) and Multi-Scale Dynamic FFN (MS-FFN).
• PA-SA modulates features with transmission and light priors to produce transmission-aware and light-aware features for long-range diffusion modeling.
• MS-FFN employs multi-scale depth-wise convolutions to expand receptive fields and aggregate local features.
• Training optimizes a diffusion loss that matches estimated noise to the true noise in the conditional diffusion process.

Experimental results

Research questions

• RQ1 How can physical underwater imaging knowledge be integrated into a diffusion-based UIE framework?
• RQ2 Does incorporating transmission maps and background light priors improve diffusion-based restoration for underwater images?
• RQ3 Can a transformer-based diffusion module with physics-aware attention better model long-range dependencies in underwater scenes?
• RQ4 What is the contribution of each physics-derived component (transmission map, background light, MS-FFN) to UIE performance?

Key findings

• PA-Diff achieves state-of-the-art perceptual and fidelity metrics on UIE benchmarks (UIEBD and LSUI).
• Ablation studies show transmission prior (T) provides stronger gains than background light (B) alone, and full integration of both with MS-FFN and PA-SA yields best performance.
• On LSUI, the proposed method attains PSA improvements with lower FID and higher SSIM compared to several baselines.
• The ablations indicate that generated physics priors are beneficial but exact prior accuracy impacts final gains, suggesting room for improved physics prior estimation.