Source-linked AI summary
Underwater Image Enhancement by Transformer-based Diffusion Model with Non-uniform Sampling for Skip Strategy
Yi Tang, Takafumi Iwaguchi, Hiroshi Kawasaki
TL;DR
Underwater enhancement must handle diverse degradations while diffusion models impose costly iterative inference. The paper uses conditional diffusion with a lightweight transformer denoiser and non-uniform skip sampling, achieving competitive enhancement performance with higher efficiency, while SSIM remains below state-of-the-art quality.
Problem
Underwater images contain diverse degradations, while diffusion-model reverse inference requires many costly iterations.
Method
The paper combines conditional diffusion, a lightweight transformer-based denoising network, and skip sampling with piecewise and evolutionary time-step selection.
Results
The method achieves the best PSNR and competitive SSIM, including a 3.49 dB gain over Ushape, while its lightweight network improves efficiency.
Takeaways & Limitations
Conditional diffusion with lightweight denoising and non-uniform sampling provides competitive underwater enhancement while reducing network-forward and reverse-process costs.
Takeaways & Limitations
The method's SSIM does not reach state-of-the-art because recovering structural information such as edges is harder than recovering color information.
Abstract
from arXiv · showhide
In this paper, we present an approach to image enhancement with diffusion model in underwater scenes. Our method adapts conditional denoising diffusion probabilistic models to generate the corresponding enhanced images by using the underwater images and the Gaussian noise as the inputs. Additionally, in order to improve the efficiency of the reverse process in the diffusion model, we adopt two different ways. We firstly propose a lightweight transformer-based denoising network, which can effectively promote the time of network forward per iteration. On the other hand, we introduce a skip sampling strategy to reduce the number of iterations. Besides, based on the skip sampling strategy, we propose two different non-uniform sampling methods for the sequence of the time step, namely piecewise sampling and searching with the evolutionary algorithm. Both of them are effective and can further improve performance by using the same steps against the previous uniform sampling. In the end, we conduct a relative evaluation of the widely used underwater enhancement datasets between the recent state-of-the-art methods and the proposed approach. The experimental results prove that our approach can achieve both competitive performance and high efficiency. Our code is available at \href{mailto:https://github.com/piggy2009/DM_underwater}{\color{blue}{https://github.com/piggy2009/DM\_underwater}}.
1 INTRODUCTION
Underwater image enhancement is important for analyzing ROV imagery, but underwater scenes combine diverse distortions that challenge conventional methods. The paper addresses diffusion-model efficiency and conditional image generation through a lightweight transformer denoiser and skip-based sampling.
- Underwater image enhancement supports analysis tasks such as object tracking, detection, and scene recognition from ROV imagery.
- Complex underwater scenes exhibit color distortion, biased illumination, and heavy non-uniform blurriness rather than a single noise type.
- Diffusion reverse processes are slow because generating a clean image requires many iterative denoising steps.
- The method uses conditional diffusion to generate an enhanced image corresponding to an underwater input while reducing reverse-process cost.It combines a lightweight transformer denoising network with skip sampling.
- Two non-uniform time-step sampling methods retain the iteration count while further improving enhancement performance over uniform sampling.
2 RELATED WORKS
Related work spans heuristic physical-model methods, learning-based enhancement, and diffusion-based generative frameworks. The paper builds on conditional diffusion and lightweight transformer processing to target underwater enhancement and inference efficiency.
- 2.1 Underwater image enhancement: Heuristic underwater enhancement methods use physical models such as Retinex and underwater optical imaging, but variable marine conditions limit robustness and efficiency.
- 2.1 Underwater image enhancement: Learning-based methods progressed from GAN approaches to complex models enabled by datasets such as UIEB and LSUI.
- 2.2 Diffusion model: DDPM uses a forward Markov process that progressively adds Gaussian noise and a reverse process that estimates prior states.
- 2.2 Diffusion model: The proposed framework combines diffusion, a transformer-based denoising network, and channel-wise attention for underwater enhancement.
- 2.2 Diffusion model: The reverse process starts from x_T and iteratively generates the predicted x_0, motivating methods that reduce inference computation.
- 2.2 Diffusion model: Conditional diffusion supplies an image condition so the model can produce an output corresponding to the input rather than an indeterminate result.
3 PROPOSED METHOD
The method combines conditional diffusion for underwater enhancement with a lightweight channel-attention transformer denoiser and accelerated reverse sampling. It reduces per-iteration cost and uses skip, piecewise, and evolutionary sampling strategies to improve inference efficiency and performance.
- 3.1 Overview: Conditional diffusion uses a noisy image, underwater conditioning image, and time step to predict noise and generate a specific enhanced image.The conditional image supplies context, avoiding the indeterminate outputs of an unconditional diffusion model.
- 3.2 Trasnformer-based denoising network: The lightweight denoising network concatenates noisy and conditional images, encodes the time step, and combines these representations before transformer processing.The two images form a 6-channel input, while the time-step embedding is reshaped into feature maps and added to image features.
- 3.2 Trasnformer-based denoising network: Channel-wise attention and eight transformer blocks reduce model scale and parameters while improving denoising efficiency.The design avoids spatial self-attention over image patches and uses channel-wise attention to lower computational complexity.
- 3.3 Skip sampling strategy: DDIM-based skip sampling reduces reverse-process iterations, while setting η to 0 removes the random term and makes sampling deterministic.The resulting inference sequence uses S time steps with |S| < |T|.
- 3.3 Skip sampling strategy: Piecewise sampling assigns different strides to the early and late time-step ranges because the early reverse process is more important.The sequence is split at c, using d1 over [a,c] and d2 over [c,b].
- 3.3 Skip sampling strategy: The proposed non-uniform sampling further improves diffusion-model performance while using the same number of steps as uniform sampling.Figure 2 contrasts uniform, piecewise, and searching sampling strategies.
- 3.3 Skip sampling strategy: The evolutionary sampling method treats time-step sequences as genes and retains better-performing offspring after mutation and crossover.The algorithm validates generated sequences and updates the gene queue when a child performs better.
4 EXPERIMENTS
Experiments evaluate the proposed underwater enhancement method on paired datasets against prior methods, then study its denoising network and non-uniform skip sampling. The method achieves strong PSNR and visual color-restoration results, while both proposed sampling strategies improve performance at fixed iterations.
- 4 EXPERIMENTS: The experiments use paired underwater datasets, PSNR and SSIM full-reference metrics, and comparisons with traditional and deep-learning methods.UIEB contains 890 paired images, while LSUI provides 4500 paired training images and diverse underwater scenes.
- 4.3 Comparison with the state-of-the-arts: The proposed method achieves the best PSNR and competitive SSIM, with a 3.49 dB gain over Ushape, while balancing model size and performance.Its runtime is only slightly slower than previous methods, and skip sampling accelerates the iterative diffusion process.
- 4.3 Comparison with the state-of-the-arts: Visual comparisons show stronger color correction and restoration than prior methods, avoiding green noise and reddish color artifacts observed in some alternatives.The comparison includes enhanced outputs from previous methods, the proposed method, and ground truths.
- 4.4 Ablation studies: The lightweight Transformer denoising network improves performance over traditional UNet while using a smaller model and achieving faster runtime.The comparison is conducted on the LSUI dataset.
- 4.4 Ablation studies: With 10 reverse-process iterations, both piecewise and evolutionary-search sampling improve performance, and evolutionary search slightly outperforms piecewise sampling.The evolutionary-search strategy reaches the PSNR associated with S=40 and surpasses the value for S=20.
5 CONCLUSION
The paper presents a conditional diffusion framework for underwater enhancement that combines a lightweight transformer denoiser with skip-based non-uniform sampling. It reports competitive enhancement performance while improving reverse-process efficiency.
- The conditional diffusion model generates corresponding enhanced images from underwater inputs.
- A lightweight transformer denoising network improves image quality and accelerates each reverse-process network forward.
- Piecewise and evolutionary-algorithm searching sampling further improve performance using the same number of time steps as uniform sampling.