TSUIE：高效的两阶段水下图像增强框架

水下图像增强效率优化方法研究采用两阶段设计，第一阶段全局颜色恢复网络（GCRN）通过扩散模型高效生成低分辨率色彩校正图像，第二阶段细节重建网络（DRN）融合跨频特征融合模块（CF3）与原始高分辨率图像纹理，实现高质量增强。实验表明该方法在PSNR、SSIM等指标上显著优于SOTA方法，同时推理延迟降低约60%。

　　

陈凌峰|马天恒|徐元鑫|徐志涵|卢世波|王云农

中国浙江省大学信息科学与电子工程学院，杭州

摘要

基于扩散的水下图像增强（UIE）方法在全局颜色恢复方面表现出强大的能力，但当采样步骤大幅减少时，它们通常会遭受过高的推理延迟和细粒度细节的明显退化。为了解决这些限制，我们提出了一种有效的基于扩散的UIE方法，即TSUIE。该框架采用两阶段设计，包括两个核心网络：全局颜色恢复网络（GCRN）和细节重建网络（DRN）。具体来说，GCRN利用扩散模型生成具有恢复颜色信息的低分辨率增强图像，并结合优化的采样过程以提高实时性能。然后，DRN将这些低分辨率的颜色增强输出与原始的高分辨率图像（富含详细纹理）融合，以产生最终的高质量增强结果。为了实现高效的融合，DRN集成了新颖的跨频率特征融合（CF3）模块，从而能够有效整合来自两个输入源的高频和低频特征。得益于这种两阶段设计，TSUIE在保持相对较低推理延迟的同时，显著提高了水下图像增强的质量。在主流UIE数据集上的全面评估表明，与最先进（SOTA）方法相比，它在多个指标上都具有很强的竞争力：它在PSNR和SSIM方面取得了显著提升，并在其他关键指标上也达到了SOTA水平。

引言

水下光学成像在水下考古学、海洋探索和水下机器人等领域发挥着至关重要的作用。由于恶劣的水下环境，水下图像经常会出现严重的退化，包括颜色失真、模糊和低对比度[30]。这些问题对水下光学设备的正常运行构成了挑战，并严重阻碍了对水下世界的进一步探索[1]。

为了应对这些挑战并提高水下图像的质量，研究人员提出了一系列有效的方法，包括传统方法和基于深度学习的技术。传统的水下图像增强（UIE）方法[40]、[49]、[56]，如对比度限制自适应直方图均衡化（CLAHE）[35]、白平衡（WB）[25]和Retinex[34]，通常使用手动设计的先验，并依赖于图像和水下环境的统计特性来纠正图像质量退化。然而，这些方法很难适应不同的水下环境。近年来，深度学习的发展有助于解决UIE中的这些问题。得益于水下图像增强数据集[19]、[30]的建设，基于深度学习的UIE方法迅速涌现，通常能够实现更优的增强质量，并具有更高的鲁棒性和泛化能力。特别是，卷积神经网络（CNNs）、生成对抗网络（GANs）和Transformer的使用正在逐步提高增强图像的质量[21]、[30]、[42]。然而，CNNs受限于其较小的感受野，使得基于CNN的方法难以适应各种水下环境。基于GAN的方法需要复杂的损失函数设计和仔细的训练策略，这带来了显著的挑战。虽然Transformer在计算机视觉任务中表现出强大的性能，但其在序列长度上的二次复杂性限制了其处理高分辨率水下图像的能力。

利用扩散模型[7]、[38]的强大拟合能力，基于扩散模型的UIE方法在各种评估指标[3]、[24]、[39]、[52]上的增强性能超过了之前的基于深度学习的方法。然而，这些方法仍然存在两个关键缺点，阻碍了它们在现实世界场景中的实际应用。首先，它们依赖于数百到数千次迭代扩散步骤来逐步消除噪声。为了保证高质量的增强，扩散步骤的数量通常设置为数千次，这导致了过长的推理延迟，从而限制了它们在实时系统中的应用。其次，尽管去噪扩散隐式模型（DDIMs）[38]在保持图像质量的同时显著减少了推理时间，但极少的采样步骤（例如$$10步）仍然无法达到预期的增强效果。具体来说，细粒度细节容易丢失，颜色失真也无法完全纠正，因为模型缺乏足够的迭代过程来同时优化全局颜色平衡和局部纹理特征。

为了解决这两个挑战，我们提出了TSUIE，这是一个专为高效和高性能UIE设计的两阶段网络。具体而言，第一阶段专注于全局颜色恢复，由全局颜色恢复网络（GCRN）处理；第二阶段专门用于细节重建，通过细节重建网络（DRN）实现。受[39]的启发，GCRN利用DDIMs高效生成低分辨率的增强水下图像，从而实现所需的增强效果。在SwinFusion[28]的基础上，DRN将GCRN的输出与从原始高分辨率图像中提取的纹理信息融合，得到最终输出。

所提出的TSUIE的主要贡献如下：

•

我们提出了TSUIE，这是一种高效的水下图像增强架构，它通过解耦退化过程将增强任务分解为两个连续阶段：全局颜色恢复和细节重建。具体来说，第一阶段使用轻量级模型进行全局颜色校正，而第二阶段通过跨频率特征融合（CF3）进行细节补偿，实现颜色和细节的协同增强。

•

我们设计了一个高效的细节重建网络（DRN），具有多分支特征提取和自适应融合机制。特别是，我们在该框架中引入了跨频率特征注意力（CFFA）模块，该模块能够充分融合多频率信息。这个模块集成到自适应融合机制中，动态地对齐并融合来自原始图像的颜色恢复结果和纹理细节，有效避免了传统方法中全局优化导致的边缘模糊问题。

•

与SOTA方法的实验结果表明，TSUIE在推理时间缩短方面表现出色，如图1所示。即使GCRN的采样步骤较少，DRN也能实现高质量的增强。广泛的消融实验验证了我们方法的有效性。

本文的其余部分组织如下：第2节回顾相关工作；第3节详细描述所提出的TSUIE方法；第4节展示实验结果和分析；第5节总结本文。

相关研究

相关研究

当前的水下图像增强（UIE）方法大致可以分为两类：非深度学习方法和基于深度学习的方法。

概述

与其他基于深度学习的UIE方法相比，基于扩散模型的UIE方法通常在全局颜色一致性和视觉自然性方面提供更优的增强性能。然而，在[3]、[24]、[52]中采用的DDPMs由于在全分辨率输入图像上执行迭代去噪过程，导致了过高的计算延迟。虽然引入DDIMs在一定程度上缓解了这一瓶颈[39]，但仍然存在

实现细节

TSUIE使用PyTorch实现，并在配备NVIDIA GeForce GTX 4090 GPU和Intel Xeon Gold 6154 CPU的系统上进行了评估。网络使用Adam优化器进行优化，超参数为$$，初始学习率为$$。对于GCRN，扩散模型配置了$$次扩散步骤，其中方程（1）中的噪声调度$$的范围从$$到$$，采用均匀划分。在训练过程中，GCRN处理$$分辨率的输入，批量大小为$$百万

结论

在本文中，我们提出了一种创新的UIE方法TSUIE，它提供了一个简单而有效的解决方案。我们采用了两阶段策略：全局颜色恢复网络（GCRN）专注于恢复自然颜色效果，而细节重建网络（DRN）将这些颜色增强特征与结构细节结合起来。与多个数据集上的最先进（SOTA）方法进行全面比较后，表明TSUIE在PSNR等方面取得了显著改进

CRediT作者贡献声明

陈凌峰：本研究的概念化、方法论和软件设计。马天恒：撰写——原始草稿、正式分析、数据管理。徐元鑫：撰写——审阅与编辑、验证、监督、项目管理、资金获取、概念化。徐志涵：撰写——审阅与编辑。卢世波：撰写——审阅与编辑、概念化。王云农：撰写——审阅与编辑。

利益冲突声明

作者声明他们没有已知的财务利益或个人关系可能影响本文报告的工作。

致谢

本研究得到了中国浙江省自然科学基金（项目编号LMS26F010003）和中国国家重点研发计划（2021YFC2800101）的支持。

陈凌峰目前是中国浙江大学信息科学与电子工程学院的博士生。他于2023年获得中国东华科技大学信息科学与工程学院的信息工程学士学位。他的研究兴趣包括计算机视觉、图像处理和水下机器人。