[论文解读] U-shape Transformer for Underwater Image Enhancement

Lintao Peng, Chunli Zhu+1 more|arXiv (Cornell University)|Nov 23, 2021

Image Enhancement Techniques参考文献 46被引用 24

一句话总结

本文提出一种U形Transformer网络，结合通道注意力与空间注意力机制，实现最先进的水下图像增强效果。该工作构建了一个大规模真实世界水下数据集（LSUI），包含4,279张图像及其对应的参考图像、分割图和透射率图，并采用结合RGB、LAB和LCH色彩空间的多色彩空间损失函数，以提升对比度与饱和度，相比先前方法实现超过2 dB的PSNR增益。

ABSTRACT

The light absorption and scattering of underwater impurities lead to poor underwater imaging quality. The existing data-driven based underwater image enhancement (UIE) techniques suffer from the lack of a large-scale dataset containing various underwater scenes and high-fidelity reference images. Besides, the inconsistent attenuation in different color channels and space areas is not fully considered for boosted enhancement. In this work, we constructed a large-scale underwater image (LSUI) dataset including 5004 image pairs, and reported an U-shape Transformer network where the transformer model is for the first time introduced to the UIE task. The U-shape Transformer is integrated with a channel-wise multi-scale feature fusion transformer (CMSFFT) module and a spatial-wise global feature modeling transformer (SGFMT) module, which reinforce the network's attention to the color channels and space areas with more serious attenuation. Meanwhile, in order to further improve the contrast and saturation, a novel loss function combining RGB, LAB and LCH color spaces is designed following the human vision principle. The extensive experiments on available datasets validate the state-of-the-art performance of the reported technique with more than 2dB superiority.

研究动机与目标

• 解决缺乏大规模、高保真度、场景多样的水下图像数据集及对应参考图像的问题。
• 通过专用注意力机制克服水下图像中颜色与空间衰减不一致的问题。
• 通过增强对比度、饱和度与色彩保真度，提升水下图像的视觉质量。
• 开发一种深度学习框架，超越现有的数据驱动与物理模型驱动的水下图像增强方法。

提出的方法

• 提出一种U形Transformer架构，集成通道注意力多尺度特征融合（CMSFFT）与空间注意力全局特征建模（SGFMT），以应对不同颜色通道与空间区域的非均匀衰减问题。
• 设计一种新型多色彩空间损失函数，结合RGB、LAB与LCH色彩空间，以契合人类视觉感知并提升色彩准确性。
• 构建了一个大规模水下图像（LSUI）数据集，包含4,279组真实世界水下图像，附带高质量参考图像、语义分割图与中等透射率图。
• 采用[33]中的U形架构作为主干网络，并通过引入任务特定的注意力模块，提升特征表示能力。
• 通过感知损失与多色彩空间损失的组合进行模型训练，以保留结构细节并增强色彩保真度。
• 采用多尺度梯度流机制（MSG），以改善局部细节重建并减少伪影。

实验结果

研究问题

• RQ1 大规模、真实世界水下图像数据集，若包含高保真参考图像，是否能显著提升数据驱动型UIE方法的性能？
• RQ2 在深度学习框架中，如何有效建模颜色通道与空间区域之间不一致的衰减特性？
• RQ3 在损失函数中整合RGB、LAB与LCH多种色彩空间，是否能带来更优的色彩校正与更佳的感知质量？
• RQ4 在基于Transformer的架构中，通道注意力与空间注意力机制在多大程度上提升了水下图像的色彩保真度与细节恢复能力？

主要发现

• 所提出的U形Transformer在Test-L400数据集上达到24.16的PSNR，在Test-U90数据集上达到22.91，优于所有基线方法超过2 dB。
• 包含全部组件（CMSFFT、SGFMT、MSG、MCSL）的完整模型在PSNR与SSIM上均达到最高值，证实了集成设计的有效性。
• 多色彩空间损失函数显著提升了色彩准确性，表现为所有对比方法中最低的CIEDE2000评分（最佳色彩校正）。
• 消融实验表明，每个组件（CMSFFT、SGFMT、MSG、MCSL）均对性能有独特贡献，完整模型在视觉质量与定量指标上均表现最佳。
• LSUI数据集在水体类型、光照条件与目标类别多样性方面优于现有数据集，支持更鲁棒的模型训练。
• 该方法在色彩校正方面表现最佳，与真实Macbeth Color Checker的色彩差异最小，尤其在专业相机拍摄图像上表现突出。