UniUIR: Considering Underwater Image Restoration as An All-in-One Learner

TL;DR

Underwater images suffer from mixed degradations such as color cast, haze, low contrast, blur, and low illumination. The paper proposes UniUIR, an all-in-one restoration framework built on a Mamba Mixture-of-Experts backbone (MMoE-UIR), a spatial-frequency prior generator (SFPG), depth cues from Depth Anything V2, and a Latent Conditional Diffusion Model (LCDM) to refine priors. Stage I optimizes $\mathcal{L}_{\mathrm{stageI}} = \mathcal{L}_1(X_{gt}, X_{hq}) + \lambda_1 \mathcal{L}_{\mathrm{depth}}$ and Stage II optimizes $\mathcal{L}_{\mathrm{stageII}} = \mathcal{L}_1(X_{gt}, X_{hq}) + \mathcal{L}_{\mathrm{diff}}$, with a forward diffusion $q(\mathbf{Z}_T|\mathbf{Z}) = \mathcal{N}(\mathbf{Z}_T; \sqrt{\bar{\alpha}_T}\mathbf{Z}, (1-\bar{\alpha}_T)\mathbf{I})$ and a denoiser $\epsilon_{\theta}(\mathbf{Z}_t, \mathbf{C}, t)$. Experiments show strong improvements over state-of-the-art across reference and non-reference metrics and better generalization to downstream tasks like depth estimation and segmentation, indicating practical benefits for marine imaging.

Abstract

Existing underwater image restoration (UIR) methods generally only handle color distortion or jointly address color and haze issues, but they often overlook the more complex degradations that can occur in underwater scenes. To address this limitation, we propose a Universal Underwater Image Restoration method, termed as UniUIR, considering the complex scenario of real-world underwater mixed distortions as an all-in-one manner. To decouple degradation-specific issues and explore the inter-correlations among various degradations in UIR task, we designed the Mamba Mixture-of-Experts module. This module enables each expert to identify distinct types of degradation and collaboratively extract task-specific priors while maintaining global feature representation based on linear complexity. Building upon this foundation, to enhance degradation representation and address the task conflicts that arise when handling multiple types of degradation, we introduce the spatial-frequency prior generator. This module extracts degradation prior information in both spatial and frequency domains, and adaptively selects the most appropriate task-specific prompts based on image content, thereby improving the accuracy of image restoration. Finally, to more effectively address complex, region-dependent distortions in UIR task, we incorporate depth information derived from a large-scale pre-trained depth prediction model, thereby enabling the network to perceive and leverage depth variations across different image regions to handle localized degradation. Extensive experiments demonstrate that UniUIR can produce more attractive results across qualitative and quantitative comparisons, and shows strong generalization than state-of-the-art methods.

Figures (12)

• Figure 1: The figures present a subjective statistical analysis of the predominant distortions in the UIEB UIEB dataset. Although each image may exhibit multiple distortions, for simplified classification, each is categorized by its most visually prominent distortion. For example, while Sample 1 shows both color distortion and foreground blurring (blue box) as well as regional blurring in the background (orange box), it is primarily classified as blurred because the blurring is more pronounced than the color distortion, which is less severe compared to other samples like Sample 3. Reference images from the UIEB dataset are provided alongside each sample for comparison.
• Figure 2: Comparison between our proposed UniUIR and previous deep learning-based approaches to underwater image restoration. (a) Previous depth estimation methods rely on costly paired ground-truth depth data, while physical model-based approaches fail in complex distortions due to idealized assumptions, leading to poor generalization. (b) Our method leverages a depth prediction network trained on large-scale datasets and a reliable degradation prior extractor, demonstrating improved generalization in complex distortion scenarios through a mixture-of-experts (MoE)-based network.
• Figure 3: Overview of the proposed UniUIR. (a) In Stage I, the depth map $\mathbf{D}_{lq}$, degradation prior $\mathbf{Z}$, and image $\mathbf{X}_{lq}$ are processed by the MMoE-UIR network to reconstruct the high-quality image $\mathbf{X}_{hq}$. The $\mathcal{L}_{1}$ loss and the edge-aware depth loss $\mathcal{L}_{\text{depth}}$ encourage the network to focus on overall pixel-level reconstruction and the accurate recovery of edges and structures, respectively. (b) In Stage II, the prior $\mathbf{Z}$ extracted by the SFPG and conditional embedding $\mathbf{C}$ will undergoes Latent Condition Diffusion Model to iteratively denoise, yielding $\mathbf{\hat{Z}}$. Then, $\mathbf{D}_{lq}$, $\mathbf{X}_{lq}$, and $\mathbf{\hat{Z}}$ are processed by MMoE-UIR to restore the clean image. The network is fine-tuned with $\mathcal{L}_{1}$ and $\mathcal{L}_{\text{diff}}$ to optimize image restoration and the diffusion process. (c) During Inference, SFPG* first produces the conditional embedding $\mathbf{C}$. Subsequently, a random noise $\mathbf{\hat{Z}_{t}}$ is sampled from a Gaussian distribution. Through the reverse process, $\mathbf{C}$ and $\mathbf{\hat{Z}_{t}}$ undergo iterative denoising to generate $\mathbf{\hat{Z}}$. Lastly, the fine-tuned MMoE-UIR then leverages $\mathbf{X}_{lq}$, $\hat{\mathbf{Z}}$, and $\mathbf{D}_{lq}$ to reconstruct the final high-quality image $\mathbf{X}_{hq}$.
• Figure 4: The detailed structure of the proposed MMoE-UIR. (a) Mamba Mixture-of-Experts Underwater Image Restoration (MMoE-UIR); (b) Mamba Mixture-of-Experts Block (MMoEB); (c) Water Mixture-of-Experts (W-MoE).
• Figure 5: Visual comparison of restored results for the T90UIEB test set.