【图像增强】U-shape Transformerfor Underwater Image Enhancement

本文讲述了一种水下图像增强的U形Transformer,本文发布在arXiv、SpringerLink、IEEE Xplore 等多个平台上，本文作者为 Lintao Peng（彭林涛）、Chun Li Zhu（朱春丽）和 Liheng Bian（边丽衡）。其中，彭林涛是北京理工大学的博士生，他的研究兴趣主要集中在计算成像和传感方面，特别是基于深度学习的复杂环境成像和传感技术，同时也在研究各种新型自注意力机制和损失函数以提高成像和传感网络的性能。

文献地址：U-Shape Transformer for Underwater Image Enhancement | IEEE Journals & Magazine | IEEE Xplore

代码地址：https://github.com/LintaoPeng/U-shape_Transformer_for_Underwater_Image_Enhancement

Abstract-摘要

翻译

水下杂质的光吸收和散射导致水下成像质量差。现有的基于数据驱动的水下图像增强（UIE）技术缺乏包含各种水下场景和高保真参考图像的大规模数据集。此外，在不同的颜色通道和空间区域的不一致的衰减没有充分考虑到升压增强。在这项工作中，我们建立了一个大规模的水下图像（LSUI）数据集，它涵盖了更丰富的水下场景和更好的视觉质量的参考图像比现有的水下数据集。该数据集包含4279个真实水下图像组，其中每个原始图像的清晰参考图像、语义分割图和介质透射图对应配对。我们还提出了一个U形Transformer网络，其中Transformer模型首次引入UIE任务。在U型Transformer中集成了通道级多尺度特征融合Transformer（CMSFFT）模块和空间级全局特征建模Transformer（SGFMT）模块，增强了网络对衰减较严重的颜色通道和空间区域的关注。同时，为了进一步提高对比度和饱和度，结合人眼视觉原理，设计了RGB、LAB和LCH三种颜色空间的损失函数。在可用数据集上进行的大量实验验证了所报告技术的最新性能，具有超过2dB的优越性。

一、Introduction-引言

水下图像增强（UIE）技术对于获取水下图像和研究水下环境至关重要，但水下图像常因水中杂质的光吸收和散射而质量不佳。现有 UIE 技术存在诸多问题，本文构建了大规模水下图像（LSUI）数据集，提出了 U 形 Transformer 网络及相关方法，在水下图像增强领域取得了新的成果。

二、Relative Work-相关工作

2.1相关工作分析

UIE 方法分类及问题：现有 UIE 方法分为视觉先验、物理模型和数据驱动三类。视觉先验方法通过修改像素值改善图像视觉质量，但忽略物理退化过程，在复杂水下环境效果不佳；物理模型方法基于物理成像模型假设，通过估计参数反转退化过程，但模型假设在复杂水下环境不一定合理，且评估多参数困难，未考虑人眼感知特性；数据驱动方法虽有良好性能，但现有水下数据集存在不足，且未充分考虑水下图像不同颜色通道和空间区域的不一致衰减问题。
水下图像数据集现状：现有数据集分为非参考数据集和全参考数据集。非参考数据集缺乏匹配的清晰参考图像，不能用于端到端训练；全参考数据集存在生成样本不真实、数量和场景有限等问题。
Transformers 优势：与基于 CNN 的 UIE 方法相比，Transformer 在处理水下图像时具有优势，其内容与注意力权重的交互可视为空间变化卷积，自注意力机制能有效建模长距离依赖和全局特征，本文基于 Transformer 设计的模块可引导网络关注衰减严重的区域，结合 CNN 与 Transformer 能以较少参数取得更好性能。

2.2本文贡献

提出新颖网络架构
- 提出了一种用于处理水下图像增强（UIE）任务的 U 形 Transformer。其设计的基于 Transformer 的通道式和空间式注意力机制，能够有效去除颜色伪影和偏差，解决了水下图像在不同颜色通道和空间区域不一致衰减的问题，这是现有方法未充分考虑的。
设计多颜色空间损失函数
- 设计了一种结合 RGB、LCH 和 LAB 颜色空间特征的新型多颜色空间损失函数。此函数进一步提高了输出图像的对比度和饱和度，使增强后的图像在视觉质量上更优，更符合人眼对颜色的感知。
构建大规模水下图像数据集
- 发布了一个大规模的 LSUI 数据集，包含 4279 张真实水下图像以及相应的高质量参考图像、语义分割图和介质传输图。该数据集涵盖了更丰富的水下场景（如水类型、光照条件和目标类别），相比现有水下数据集，能更好地促进 UIE 技术的进一步发展，为相关研究提供了更充足的数据支持。

三、Method-方法

本文提出的模型为 U 形 Transformer，主要用于水下图像增强（UIE）任务，包含基于 CMSFFT（通道式多尺度特征融合 Transformer）和 SGFMT（空间式全局特征建模 Transformer）的生成器和判别器。生成器中，原始图像先多次下采样再经卷积输入卷积块，CMSFFT 和 SGFMT 的输出用于解码；判别器中，多个卷积块输入来自不同层的特征图，多尺度连接使梯度流自由，丰富图像细节。

3.1SGFMT（空间式全局特征建模Transformer）

3.1.1整体功能

基于水下图像在不同空间区域非均匀退化的先验知识，SGFMT 旨在辅助网络对全局信息进行建模，增强网络对图像中严重退化部分的关注，从而实现更均匀有效的水下图像增强效果。

3.1.2工作流程

1、输入特征处理：输入特征图大小为 $F_{in} \in \Re ^{\frac{H}{16}*\frac{W}{16}*C}$ ，为了适应 Transformer 期望的一维序列输入，首先使用线性投影将二维特征图拉伸为特征序列 $S_{in} \in \Re^{\frac{HW}{256}*C}$ 。同时，为了保留每个区域的位置信息，直接合并可学习的位置嵌入（PE），其计算方式为 $S_{in}=W*F_{in}+PE$ ，其中 $W*F_i$ 表示线性投影操作。

2、Transformer层计算：将处理后的特征序列输入到包含 4 个标准 Transformer 层的模块中。每个 Transformer 层包含一个多头注意力块（MHA）和一个前馈网络（FFN）。MHA 用于计算特征序列中不同位置之间的注意力关系，从而捕捉全局信息；FFN 则包括一个归一化层和一个全连接层，用于对特征进行进一步的变换和处理。具体计算过程如下：

3、特征映射恢复：经过 4 个 Transformer 层的计算后，输出特征序列 $S_l\in\Re^\frac{HW}{256}*C$ ，最后通过特征重映射将其恢复为特征图 $F_{out} \in \Re ^{\frac{H}{16}*\frac{W}{16}*C}$ 。

3.2CMSFFT（通道式多尺度特征融合Transformer）

3.2.1整体功能

用于替代原始生成器编码 - 解码架构中的跳跃连接，目的是强化网络对水下图像中更严重衰减颜色通道的关注，以解决水下图像在不同颜色通道上不一致衰减的问题，从而提升水下图像增强的效果。

3.2.2工作流程

1、多尺度特征编码

输入是不同尺度的特征图 $F_i \in \Re^{\frac{H}{2^i}*\frac{W}{2^i}*C_i}(i=0,1,2,3)$ .采用卷积核大小 $\frac{P}{2^i}*\frac{P}{2^i}*C_i(i=0,1,2,3)$ ,在本次工作中P设置为32和步长为 $\frac{P}{2^i}(i=0,1,2,3)$ 的卷积对不同尺度特征图进行线性投影，得到四个特征序列 $S_i\in\Re^{d*C_i}(i=1,2,3,4)$ ,其中 $d\in\frac{HW}{P^2}$ 。这些卷积核将特征图划分为相同数量的块，且通道数 $C_i(i=1,2,3,4)$ 保持不变（在本工作中 $C_1=64,C_2=128,C_3=256,C_4=512$ ）.然后通过特定公式计算查询向量 $Q_i \in \Re^{d*C_i}(i=1,2,3,4)$ 、 $K\in\Re^{d*C_i}$ 和 $V\in\Re^{d*C}$ 。具体为 $Q_i=S_iW_{Q_i}$ ， $K=SW_K$ ， $V=SW_V$ ，其中、和是可学习的权重矩阵，是由通过通道维度拼接而成，且 $C=C_1+C_2+C_3+C_4$ 。

2、通道多头注意力

该模块有六个输入，分别是 $K\in\Re^{d*C}$ 、 $V\in\Re^{d*C}$ 和 $Q_i\in\Re^{d*C_i}(i=1,2,3,4)$ 。
首先计算通道式注意力 $CA_i\in\Re^{C_i*d}(i-1,2,3,4)$ ，计算方式为 $CA_i=SoftMax(IN(\frac{Q^{T}_iK}{\sqrt[2]{C}}))V^T$ ，其中IN表示实例归一化操作，此注意力操作沿通道轴进行，可引导网络关注图像质量退化更严重的通道，同时IN有助于梯度流平稳传播。
然后计算第个 CMHA 层的输出 $CMHA_i=(CA^1_i+CA^2_i+...CA^N_i)/N+Q_i$ （在本实现中N=4）。

3、前馈网络

其输出计算方式与其他 Transformer 类似， $O_i=CMHA_i+MLP(LN(CMHA_i))$ ，其中 $O_i\in\Re^{d*C_i}(i=1,2,3,4)$ ,MLP表示多层感知机。此操作需要重复l次（在本工作中l=4）以构建层 Transformer。

4、特征重映射

最后对四个不同输出特征序列进行特征重映射 $O_i \in \Re^{C_i*d}(i=1,2,3,4)$ ，将它们重组为四个特征图 $F_i\in\Re^{\frac{H}{2^i}*\frac{W}{2^i}*C_i}(i=0,1,2,3)$ ，这些特征图作为生成器解码部分卷积块的输入。

3.3 Loss Function-损失函数

为提升水下图像增强效果，充分利用 LAB 和 LCH 颜色空间特性，本文设计了结合 RGB、LAB 和 LCH 颜色空间的多颜色空间损失函数训练网络。将 RGB 图像转换至 LAB 和 LCH 空间：

其中、x,y,G(x)和分别为原始输入、参考图像、生成器输出的清晰图像。LAB 和 LCH 空间损失函数如下：

其中Q为量化算子，量化不同色彩空间某通道是为计算增强图像与参考图像在此通道的交叉熵损失。RGB色彩空间 $L_2$ 的损失 $Loss_{RGB}$ 、感知损失 $Loss_{pre}$ 及 $Loss_{LAB}$ 和 $Loss_{LCH}$ 为生成器的四个损失函数。

此外引入标准 GAN 损失函数最小化生成与参考图片差异

最终损失函数为：

我们在文章给定的数据集的200张图片上训练了100个epoch得到了如上结果psnr结果为19.25

总结

提出了一个U形Transformer网络的最先进的UIE性能。该网络的CMSFFT和SGFMT模块可以解决水下图像在不同颜色通道和空间区域衰减不一致的问题，这是现有方法中没有考虑的。大量的实验验证了网络的上级能力，去除颜色伪影和投影。结合图像空间损失函数，进一步提高了输出图像的对比度和饱和度。然而，要收集所有复杂场景的图像是不可能的，例如深海低光场景。因此，我们将介绍其他一般增强技术，如低光增强用于未来的工作。