[论文解读] Underwater Image Enhancement by Transformer-based Diffusion Model with Non-uniform Sampling for Skip Strategy

Yi Tang, Takafumi Iwaguchi+1 more|arXiv (Cornell University)|Sep 7, 2023

Image Enhancement Techniques被引用 9

一句话总结

本文提出一种用于水下图像增强的条件扩散模型，使用基于轻量变换器的去噪网络和非均匀跳跃采样以提升效率和质量。

ABSTRACT

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}}.

研究动机与目标

• 以稳定高质量输出为目标，推动鲁棒的水下图像增强（UIE）。
• 设计一个条件扩散框架，以水下图像作为条件引导增强。
• 设计一个轻量级基于变换器的去噪网络以降低每次迭代成本。
• 引入具有非均匀时间步采样的跳跃采样，以在不牺牲质量的前提下加速推理。

提出的方法

• 采用一个条件去噪扩散概率模型，其输入为有噪声的图像、作为条件 c 的低质量水下输入，以及时间步 t。
• 使用一个具通道注意力的轻量级基于变换器的去噪网络来预测噪声 epsilon_theta。
• 如同 DDPM 一样引入条件均值 mu_theta，使其在 c 条件下进行迭代逆扩散以生成 x0。
• 通过受 DDIM 启发的非马尔可夫性逆过程实现跳跃采样以减少迭代次数。
• 提出两种非均匀采样策略：分段采样（时间步区间的两段步长）和一个用于优化采样序列的进化算法。
• 在 UIEB 和 LSUI 数据集上进行训练与评估，使用 PSNR 和 SSIM 指标。

实验结果

研究问题

• RQ1 一个条件扩散模型能否在低质量输入的引导下产生高质量的水下图像增强？
• RQ2 轻量级基于变换器的去噪网络是否在不牺牲增强质量的前提下提升效率？
• RQ3 与均匀采样相比，非均匀采样策略（分段采样与进化采样）是否能提高推理速度和/或恢复性能？
• RQ4 与标准数据集上的最先进 UIE 方法相比，所提方法在 PSNR 和 SSIM 上有何表现？

主要发现

• 所提方法在 LSUI 上获得最佳 PSNR，在 LSUI 和 UIEB 上在 SSIM 方面具有竞争力，相较于八个最先进方法。
• 基于变换器的去噪网络在参数量相似或更少、每次迭代运行时间相同的情况下，优于 UNet 基线。
• 跳跃采样显著加速推理；所提出的非均匀采样策略比均匀采样进一步提升性能。
• 在 LSUI 上，该方法在 10 次迭代时达到 PSNR 27.65、SSIM 0.8867；在 UIEB 上达到 PSNR 28.20、SSIM 0.9429。
• 该方法在视觉色彩矫正与修复方面具有竞争力，尤其是在具有挑战性的水下场景中显著降低色偏。
• 运行时间在 RTX 3080 上约为每张图 0.13 秒，兼顾性能与效率。