WWE-UIE: A Wavelet & White Balance Efficient Network for Underwater Image Enhancement

TL;DR

Underwater image enhancement must contend with color casts and scattering that degrade visibility. The authors propose WWE-UIE, a compact CNN that integrates adaptive white balance, a Haar-wavelet–based multi-band enhancement block (WEB), and a Sobel-gradient fusion block (SGFB), all trained with a composite loss in the HVI color space. Key contributions include a channel-wise adaptive white balance prior, a fixed-filter wavelet decomposition for global–local restoration, gradient-guided edge refinement, and a loss combining Charbonnier, SSIM, perceptual, edge, and HVI terms for stable optimization. Empirical results on synthetic and real underwater datasets show competitive restoration quality with substantially fewer parameters and FLOPs, enabling real-time deployment on limited hardware with code available publicly.

Abstract

Underwater Image Enhancement (UIE) aims to restore visibility and correct color distortions caused by wavelength-dependent absorption and scattering. Recent hybrid approaches, which couple domain priors with modern deep neural architectures, have achieved strong performance but incur high computational cost, limiting their practicality in real-time scenarios. In this work, we propose WWE-UIE, a compact and efficient enhancement network that integrates three interpretable priors. First, adaptive white balance alleviates the strong wavelength-dependent color attenuation, particularly the dominance of blue-green tones. Second, a wavelet-based enhancement block (WEB) performs multi-band decomposition, enabling the network to capture both global structures and fine textures, which are critical for underwater restoration. Third, a gradient-aware module (SGFB) leverages Sobel operators with learnable gating to explicitly preserve edge structures degraded by scattering. Extensive experiments on benchmark datasets demonstrate that WWE-UIE achieves competitive restoration quality with substantially fewer parameters and FLOPs, enabling real-time inference on resource-limited platforms. Ablation studies and visualizations further validate the contribution of each component. The source code is available at https://github.com/chingheng0808/WWE-UIE.

Figures (9)

• Figure 1: Performance comparison between our method and its variant, and other state-of-the-art UIE methods on UIEB dataset waternet. FLOPs is measured through $256{\times}256$ input image.
• Figure 2: Overall architecture of WWE-UIE. The pipeline applies white balance correction, followed by a U-Net backbone with WGSRBs. Each block integrates a WEB for multi-band decomposition and an SGFB for edge refinement, enabling joint restoration of structure, detail, and color with lightweight efficiency.
• Figure 3: Visualization of feature maps within WEB and SGFB. The WEB highlights multi-band decomposition into structural and textural components, while the SGFB emphasizes gradient-based enhancement of edges and fine details.
• Figure 4: Visual comparison on full‐reference datasets. Each row corresponds to the datasets UIEB, LSUI, UFO, and EUVP-Scene.
• Figure 5: Visualization of our results in the CIE xyY color space, showing improved alignment of chromaticity and luminance distributions with the reference.