Illumination Aware Multiscale Attention Fusion Model for Underwater Image Enhancement

Due to the effects of light absorption and scattering in underwater environments, images captured underwater often suffer from reduced contrast, uneven lighting, blurring, color distortion, and other degradations. Many advanced methods struggle to balance detail preservation and enhancement, often over-enhancing the image. This paper introduces an illumination-aware, multi-scale, attention-based fusion network designed to address these challenges by mitigating degradations and retaining fine details. Our U-Net-based encoder-decoder model leverages features at multiple scales to extract both high-level spatial features, which aid in removing color casts and low-level texture details, restoring the texture, and eliminating blurriness, while incorporating illumination information through spatial feature transform. Furthermore, a fusion module utilizing multi-head attention combines features by detecting correlations among channels across all scales. We also introduce a multi-scale loss function that computes loss at all scales to optimize the preservation of global low-frequency contextual details and local high-frequency texture details. An additional color loss term aids in restoring color-accurate images by leveraging the LAB and LCH color spaces. We trained our architecture using 800 paired images from the UIEB dataset and 2,185 pairs from the EUVP dataset. Our architecture achieves PSNR of 31.83 dB, which is 4.11 dB higher than the present state-of-the-art method, and SSIM of 0.948. On the UIEB benchmark, the model achieves UIQM of 4.284 and UCIQE of 0.622 on reference images, while on non-reference images, it achieves UIQM of 4.138 and UCIQE of 0.586. The code is available at <uri xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">https://github.com/subudhi-badri/IMFEN</uri>.