Adaptive Double-Branch Fusion Conditional Diffusion Model for Underwater Image Restoration

Abstract

Underwater images suffer from light absorption and scattering, impairs their visibility and applications. Existing underwater image restoration (UIR) methods based on generative models struggle are difficult to adapt to the complex and dynamic underwater environments characterized by illumination interference, low-light conditions, and non-uniform turbidity. To address these issues, we propose Water-CDM, a novel Adaptive Double-Branch Fusion Conditional Diffusion Model for underwater image restoration. Specifically, an adaptive double-branch fusion conditional diffusion model is presented utilizing a U-shaped full-attention network and Guided Multi-Scale Retinex with Brightness Correction (GMSRBC) to restore the challenging regions within underwater images. More precisely, to correct color casts and enhance the sharpness of underwater images, a U-shaped full-attention network incorporating Attention Blocks is designed for noise estimation during the reverse process of the conditional diffusion model. Concurrently, to mitigate overexposure during the enhancement of low-light underwater images under illumination interference, the GMSRBC method, featuring an Adaptive Brightness Correction Module, is proposed to efficiently adjust the brightness of underwater images. Experimental results demonstrate that the proposed Water-CDM significantly improves the quality of underwater images in challenging scenarios. Encouragingly, our proposed Water-CDM yields superior restoration outcomes compared to current state-of-the-art methods on three challenging publicly available datasets. Our codes will be released at: https://github.com/HKandWJJ/Water-CDM.