Underwater Image Enhancement by Diffusion Model with Customized CLIP-Classifier

TL;DR

This paper tackles the lack of real reference images in underwater image enhancement by introducing CLIP-UIE, a two-stage diffusion-based framework. The first stage trains on a large synthetic UIE-air dataset generated via color transfer to learn mappings from underwater degradation to the in-air natural domain, while the second stage employs a CLIP-Classifier to preserve prior in-air domain knowledge during joint fine-tuning with UIE benchmarks, bridging synthetic and real data gaps. A fast fine-tuning strategy concentrates updates in the high-frequency content, and a prompt-learning CLIP-Classifier provides robust, learnable prompts to stabilize guidance. Across multiple underwater image datasets, CLIP-UIE achieves state-of-the-art or competitive results with improved naturalness and color rendition, while offering up to 10× faster optimization than traditional fine-tuning. The work advances practical UIE by reducing dependence on matched references and leveraging language-visual priors to guide domain-adaptive enhancement.

Abstract

Underwater Image Enhancement (UIE) aims to improve the visual quality from a low-quality input. Unlike other image enhancement tasks, underwater images suffer from the unavailability of real reference images. Although existing works exploit synthetic images and manually select well-enhanced images as reference images to train enhancement networks, their upper performance bound is limited by the reference domain. To address this challenge, we propose CLIP-UIE, a novel framework that leverages the potential of Contrastive Language-Image Pretraining (CLIP) for the UIE task. Specifically, we propose employing color transfer to yield synthetic images by degrading in-air natural images into corresponding underwater images, guided by the real underwater domain. This approach enables the diffusion model to capture the prior knowledge of mapping transitions from the underwater degradation domain to the real in-air natural domain. Still, fine-tuning the diffusion model for specific downstream tasks is inevitable and may result in the loss of this prior knowledge. To migrate this drawback, we combine the prior knowledge of the in-air natural domain with CLIP to train a CLIP-Classifier. Subsequently, we integrate this CLIP-Classifier with UIE benchmark datasets to jointly fine-tune the diffusion model, guiding the enhancement results towards the in-air natural domain. Additionally, for image enhancement tasks, we observe that both the image-to-image diffusion model and CLIP-Classifier primarily focus on the high-frequency region during fine-tuning. Therefore, we propose a new fine-tuning strategy that specifically targets the high-frequency region, which can be up to 10 times faster than traditional strategies. Extensive experiments demonstrate that our method exhibits a more natural appearance.

Figures (12)

• Figure 1: CLIP with proper prompts can serve as the classifier to evaluate the image quality and distinguish between in-air natural and underwater images. However, as a prevailing phenomenon, a simple approximate rephrasing of the prompt results in a significant change in CLIP score. On the contrary, our CLIP-Classifier with learned prompt (the last row of each example) shows more robust results for different types of in-air natural/underwater scenes.
• Figure 2: The preparation for the pre-trained model. (a) Randomly select a Template from the Template Pool (underwater domain). Then, the Color Transfer module, guided by the template, degrades an in-air natural image from INaturalist van2018inaturalist into underwater domain, constructing paired datasets for training image-to-image diffusion model. (b) The image-to-image diffusion model SR3saharia2022image is trained to learn the prior knowledge of mapping transitions from the real underwater degradation domain to the real in-air natural domain and to generate the corresponding enhancement results based on input synthetic underwater images produced by Color Transfer.
• Figure 3: New fine-tuning strategy. From state $x_{T}$, according to Eq. \ref{['eq:leaning_objective_y']}, we first adopt a single classifier guidance strategy, using condition $y_{1}$---the input source image---to guide the reverse diffusion process until state $x_{t}$. Then, we switch to the multi-classifier guidance strategy according to Eq. \ref{['eq:multi-classifiers_learning_objective']}. With multi-condition guidance, the intermediate results from $x_{t}$ to $x_{0}$ move towards to the in-air natural domain, mitigating the damage of fine-tuning to the prior knowledge of the pre-trained model.
• Figure 4: Illustration of the prompt learning for CLIP-Classifier. (A) Prompt Initialization. Given two text prompts describing the in-air natural image $I_{n}$ and underwater image $I_{u}$. We encode each text to get the initial in-air natural image prompt $T_{n}\in\mathbb{R}^{N\times512}$ and the initial underwater image prompt $T_{u}\in\mathbb{R}^{N\times512}$. (B) Prompt Training. We use the cross-entropy loss to constrain the learnable prompts, aligning learnable prompts and images in the CLIP latent space by maximizing the cosine similarity for matched pairs $\left \{T_{n}, I_{n}\right \}$ and $\left \{T_{u}, I_{u}\right \}$. The base model of CLIP is frozen throughout. The further use of learned prompt for CLIP-Classifier can be seen in Fig. \ref{['Fig:clip_guidance']}.
• Figure 5: Qualitative analysis of the effectiveness of the CLIP-Classifier. (a) CLIP score of the intermediate variable $x_{t}$ of the image versus time $t$. The time step is set to 2000. The CLIP score curves of the raw and reference images are plotted separately. (b) We calculate the difference in CLIP score between the raw and reference images, then plot the difference curve and label several key time points on the curve.