Color Correction Meets Cross-Spectral Refinement: A Distribution-Aware Diffusion for Underwater Image Restoration
Laibin Chang, Yunke Wang, Bo Du, Chang Xu
Abstract
Underwater imaging is often plagued by significant degradation in visual quality, primarily due to the effects of light absorption and scattering in water. Although recent underwater image enhancement (UIE) methods rely on the current advances in deep neural network architecture designs, there is still considerable room for improvement in cross-scene robustness and computational efficiency. Diffusion models have shown great success in image generation, prompting us to explore their application to UIE tasks. However, directly applying them to UIE tasks will pose two challenges, <italic xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">i.e.</i>, high computational budget and color unbalanced perturbations. To tackle these issues, we propose DiffColor, a distribution-aware diffusion and cross-spectral refinement model for efficient UIE. Unlike single-noise image restoration tasks, underwater imaging exhibits unbalanced channel distributions due to the selective absorption of light by water. To address this, we design the Global Color Correction to balance the diverse color shifts, thereby avoiding potential global degradation disturbances during the denoising process. Instead of diffusing in the raw pixel space, we transform the image into the wavelet domain to obtain such low-frequency and high-frequency spectra. For the sacrificed image details caused by underwater scattering, we further present the Cross-Spectral Detail Refinement to enhance the high-frequency details, which are then integrated with the low-frequency signal as a dual-condition for guiding the diffusion. This strategy ensures the high-fidelity of sampled content and compensates for the sacrificed details. Extensive experiments demonstrate the superior performance of DiffColor over state-of-the-art methods in both quantitative and qualitative evaluations.