海洋中含有丰富的资源，针对海洋的探索任务是水下资源开发的重要环节^[1]。水下机器人是探索海洋的重要工具，可用于水下图像采集的任务。水下机器人采集到的图像可用于后续的其他操作。然而，由于不同波长的光在海水中会出现不同程度的衰减，因此水下机器人拍摄到的图像存在颜色偏移和对比度降低等问题^[1]。水下图像增强算法的作用是对失真的图像进行质量增强，此项任务是计算机视觉和水下机器人领域的重要研究方向之一。

当前水下图像增强算法主要可分为非基于深度学习的增强算法与基于深度学习的增强算法。非基于深度学习的增强算法一般将统计观测假设作为参数估计的依据，并且通过显式建模的方式对水下图像进行增强与颜色校正。非基于深度学习的增强算法包括水下暗通道先验（underwater dark channel prior，UDCP）^[2-3]、图像模糊与亮度吸收（image blurriness and light absorption，IBLA）^[4]以及水下光衰减先验（underwater light attenuation prior， ULAP）^[5]等。由于水下场景的复杂性，非基于深度学习的增强算法存在一定的参数估计失准问题。基于深度学习的增强算法通常需要充足的数据用于增强过程，采用特定的损失对用于质量增强的深度神经网络进行优化。基于深度学习的增强算法包括水下残差网络（underwater residual network，UResNet）^[6]、浅层水下网络（shallow underwater network，UWNet）^[7]以及水下卷积神经网络（underwater convolution neural network，UWCNN）^[8]等。当前基于深度学习的增强算法已经在一定程度上解决了水下图像的失真问题，但是仍存在质量复原效果不好的问题。

针对水下拍摄的图像存在的颜色偏移、色彩不真实和对比度降低问题 ^[9]，本文结合视觉Transformer 网络^[10]设计了一种有效的水下图像质量增强模型。训练出可用于水下图像增强的Transformer是一个新兴研究领域，本文采用多种损失组合的方式进行训练，损失的选取过程充分考虑了图像的像素信息、结构信息、边缘信息与特征信息，该模型可以采用完全端到端的方式进行训练，充分利用水下图像数据存在的统计规律。

水下图像增强研究的主要目的是对失真的水下图像进行增强，复原色彩信息与对比度信息等^[1]。本节首先探讨水下图像增强的研究现状，然后介绍本文提出的增强模型的基础——视觉 Transformer网络^[10,11,12]的相关研究。

严浙平等人^[1]指出水下成像的机制要比陆地上更加复杂，光线在水中与陆地上的传播衰减不同，水体中存在的溶解物会影响成像过程，而且当水下深度较深时，拍摄过程中需要辅助光源。因此，水下图像的增强是具有挑战性的研究问题^[13]。当前的水下图像增强研究可分为非基于深度学习与基于深度学习两种。第一种是非基于深度学习的算法。Drews Jr P等人^[2-3]基于图像去雾领域的DCP提出了UDCP算法。Peng Y T等人^[4]基于图像模糊和亮度吸收，提出了一种针对水下图像的深度估计算法，并进一步用于水下图像的质量复原过程，该方法被称为IBLA。Song W等人^[5]同样从水下图像的深度估计角度出发，基于ULAP设计了一种快速的深度预测模型。常戬等人^[14]利用导向滤波和自适应算子设计了水下图像增强算法。

神经网络已经被应用于多种视觉任务中^[15,16,17]。第二种水下图像增强算法是基于深度学习的算法。Liu P 等人^[6]使用均方误差损失和边缘差分损失对卷积神经网络参数进行更新，提出了端到端的UResNet增强算法。Naik A等人^[7]利用常规的卷积设计了UWNet，结果表明，其在公开数据集上能获得有效的增强效果。Li C Y^[8]提出了采用残差学习策略的UWCNN算法，并利用像素损失和结构损失进行网络训练。陈龙等人^[18]将残差学习和注意力机制结合，提出了端到端的神经网络增强模型。范新南等人^[19]利用生成对抗网络的建模能力和对称型网络的图像重建能力设计了基于特征融合的增强网络。Li C Y 等人^[20]设计了水下图像增强网络（underwater image enhancement network， WaterNet），融合了白平衡、直方图均衡和伽马校正算法。Guo Y C等人^[21]利用密集连接和跳跃连接设计了基于生成对抗网络的增强模型，采用对抗损失和梯度损失对增强模型进行优化。Fabbri C等人^[22]利用生成对抗网络构建数据集，并设计了具有全卷积编码解码结构的水下生成对抗网络（underwater generative adversarial network， UGAN）。已有的增强算法主要采用卷积神经网络进行水下图像的增强，缺少对局部与全局特征的高质量建模以及对多种优化目标的考虑。Boudiaf A 等人^[23]采用预训练的 Transformer，初步将Transformer 用于水下图像增强任务。本文将含有多头自注意力（multi-head self-attention）的Transformer^[10]模块作为增强模型的基础组件构建具有编码-解码结构的去雾网络，并利用多种损失组合提高增强后的水下图像的质量。

深度学习领域包含众多有效的通用骨干网络， Transformer网络^[10,24]受到了研究人员的广泛关注。在自然语言处理领域中，Transformer^[24]是一种被广泛应用的架构，研究人员设计了有效的改进策略，将语言任务中的 Transformer 模型应用到计算机视觉任务中^[10,12]。Transformer网络采用了多头自注意力机制和多层感知机（multi-layer perceptron）等技术，可以用于多种视觉任务。Swin Transformer^[10]架构在图像分类、目标检测和语义分割任务上验证了层级式 Transformer 网络的性能，所设计的分阶段网络可以作为多种计算机视觉任务的基础网络，基于窗口的多头自注意力机制如式（1）所示：

其中，Q、K和V分别为查询、键和值，d为特征维度，B为相对位置偏差。Restormer 算法^[11]主要用于图像去雨和图像去噪等任务，采用了编码-解码的结构，能够获取局部和全局的特征表示。DehazeFormer算法^[25]用于图像去雾任务，该方法通过去掉部分网络层的注意力模块，构建了具有不同参数容量的去雾网络，适用于计算能力不同的平台。对 Transformer 架构的研究证明了它对计算机视觉任务的有效性，本文设计的编码-解码结构的水下图像增强模型将Swin Transformer架构中的基本Transformer模块作为基础网络块。

本文提出的水下图像增强网络由 Transformer模块^[10]和常规卷积模块组成，模型总体上采用编码与解码的结构。为了提升水下图像增强模型的效果，模型采用了4种损失组合的方式进行训练，不同的损失具有不同的作用。为了便于分析，将本文提出的模型记为UWT-Net（underwater Transformer network）。

UWT-Net 采用了编码-解码的结构方式，包括编码的降维过程和解码的升维过程。编码过程负责学习特征表示，解码过程负责增强图像的像素重建。下面给出 UWT-Net 的整体结构及前向计算的原理，并分别对网络的各个模块进行介绍。

（1）整体网络架构

编码-解码结构的 Transformer 网络如图1 所示。其中，阶段1、阶段2和阶段3是编码过程，阶段 4 和阶段 5 是解码过程。网络中的组件包括分块、Transformer 模块、下采样模块、上采样模块以及块投射。阶段1～阶段3是降维过程，输入数据的高度H、宽度W和通道数C。输入数据经过阶段1以后获得的特征图尺寸为H ×W × L，其中L代表网络前向过程的通道数，后面 4 个阶段的通道数均是L的整数倍。阶段 2 和阶段 3 采用具有残差连接的双路模式，输出特征维度分别是$\frac{H}{2}\times \frac{W}{2}\times 2L$和 $\frac{H}{4}\times \frac{W}{4}\times 4L$，这种设计的作用是降低特征图的尺寸并增加通道数。阶段4～阶段5是解码过程，分别将特征图的尺寸恢复到$\frac{H}{2}\times \frac{W}{2}\times 2L$和H ×W × L。UWT-Net阶段1输出通道数被设定为24，因此阶段2和阶段3的输出通道数分别为48和96。

（2）并行连接编码模块

图1的阶段2和阶段3作为编码过程，由并行连接的 Transformer 模块（parallel connection Transformer block，PCTB）构成，该编码过程的作用是提升增强网络的特征提取能力。PCTB 将前一层的输入特征分两个支路进行处理，然后对获得的特征进行求和，作为下一层的输入。

（3）分块与块投射

图1中输入端的分块与输出端的块投射均采用核尺寸为3× 3且步长为1的卷积实现。分块的作用是将输入的图像投射到特征空间，获取通道数为L的特征图。块投射用于将尺寸为H ×W × L的特征图投射为尺寸为H ×W × C的图像。

（4）下采样模块与上采样模块

在编码过程中，Transformer模块前使用了下采样模块降低特征图尺寸，下采样模块通过步长为 2且核尺寸为3× 3的卷积实现。在解码过程中， Transformer 模块前使用了上采样模块提升特征图尺寸，并且采用步长为1的卷积及像素重组（pixel shuffle）操作将特征图尺寸提升一倍。

（5）Transformer模块

基本的 Transformer 模块主要包含两个核心组件，第一个是多头自注意力模块，采用的是基于窗口的模式，第二个是多层感知机模块。为了实现信息的有效传递，Transformer模块一般采用层归一化（layer norm）模块对信息流进行处理，第i层的计算方式如式（2）～式（3）所示：

其中，z_i为当前层的输出，z_i-1为增强网络中前一层的输出，$z_i^{*}$为中间过程获得的特征，ϖ为层归一化，φ为式（1）所示的多头自注意力模块，φ为多层感知机。

Transformer 模块结构如图2 所示，重复使用Transformer模块可以构建UWT-Net中的5个阶段，每个阶段含有的模块数量见第3节。

为了对增强网络进行有效优化，设计了多种损失组合的方式对网络的参数进行更新。损失的设计从以下4个方面展开。

· 像素损失（pixel loss）：像素损失在像素层面对网络进行优化，使增强后图像的像素值趋近参考图像。

· 结构损失（structural loss）：为了精确地对水下图像的结构信息进行复原，根据结构相似性原理使用结构损失优化网络。

· 边缘损失（edge loss）：水下图像的边缘信息是图像质量的重要组成部分，因此将基于梯度的边缘损失作为总体损失中的一项。

· 特征损失（feature loss）：基于大规模图像数据集训练得到的模型可以表征图像的统计规律，因此使用特征损失对增强网络进行优化可以提高增强效果。

下面对4种损失进行介绍，给出每种损失的设计原理和计算过程，并说明最终损失的构成方式。τ代表失真图像，即增强网络的输入数据；y代表真实标签，即水下数据集中的参考图像；$\mathcal{H}(\cdot )$代表UWT-Net，则网络获得的增强结果为$\mathcal{H}(\tau )$。

（1）像素损失

逐像素计算的损失是图像复原任务的基础损失，对增强网络的优化采用 L1 损失，则像素损失计算如式（4）所示：

其中，w和h分别为图像的长和宽，i和 j分别为增强结果和参考图像的像素位置。

（2）结构损失

参考文献[26]证明了结构相似性度量可以作为图像复原任务的损失，为了提升水下图像增强网络对结构信息的复原能力，将结构相似性度量作为优化目标之一。将增强网络预测结果$\mathcal{H}(\tau )$记为η，则结构相似性（structural similarity，SSIM）^[18]的计算方式如式（5）所示：

其中，u_η和u_y分别为增强后图像和参考图像的均值， σ_η和 σ_y分别为增强后图像和参考图像的方差，σ_ηy为协方差，C₁与C₂为常数，P为像素。SSIM的值越大，代表两张图像的相似程度越高，所以结构损失的定义如式（6）所示：

其中，N为图像块P中的像素数量。

（3）边缘损失

水下的失真图像存在模糊现象，导致图像的边缘强度降低。边缘信息是图像的重要组成部分^[27]，为了提升增强后的图像的边缘质量，使用图像的梯度属性设计了边缘损失，通过计算预测结果$\mathcal{H}(\tau )$与参考图像 y的梯度信息，可以获得边缘损失如式（7）所示：

其中，G_h和G_v分别为水平和垂直方向的梯度值。

（4）特征损失

基于大规模图像数据集训练得到的预训练模型^[28]能够辅助基于 Transformer 的水下图像增强网络的训练，将这种在特征层面计算的损失称为特征损失。记预训练模型为Ψ(·)，则特征损失的计算方式如式（8）所示：

特征损失的计算需要选择预训练模型，在此将VGG19 网络^[29]作为特征提取网络，并将网络最后的全连接层去掉，以全连接层前面的池化层输出为特征输出。

（5）总体损失与训练流程

上述4种用于水下图像增强网络训练的损失需要在每次迭代中被同时计算，并用于指导网络参数的更新，损失的组合方式为线性加权，具体如式（9）所示：

其中，λ_s、λ_e与 λ_f分别为对应损失的权重值，具体的设置见第3节。对所设计的增强网络和组合型损失，采用梯度下降算法进行网络的参数更新。增强网络的训练流程伪代码见算法1。

算法1 增强网络的训练流程伪代码

网络的参数记为${\theta }_{\mathcal{H}}$，学习率记为lr，对数据集的总迭代轮数记为E，数据集中数据总量记为N，每次迭代的输入批数据量（batch size）记为b，则总迭代次数为$\left(\lfloor N/b\rfloor \times E\right)$ ，其中$\lfloor N/b\rfloor$代表向下取整。

初始化

for $\lfloor N/b\rfloor \times E$ steps do

利用参数$Y=\left\{y_1,y_2,y_3,\cdots y_b\right\}$ 采 样$U=\left\{{\tau }_1,{\tau }_2,{\tau }_3,\cdots {\tau }_b\right\}$

计算增强图像$\mathcal{H}(U)$

计算损失L $L\left((U),Y\right)=L_p+{\lambda }_s{L}_s+{\lambda }_e{L}_e+{\lambda }_f{L}_f$

更新参数

end for

为了公平地对比不同的水下图像增强算法，实验将两个大型公开的水下图像数据集作为训练和测试基准数据集。UIEB^[20]数据集含有890组水下图像与对应的参考图像，根据 UIEB 中的设置将800组图像作为训练数据，其余90组图像作为测试数据。此外，在实验中使用了 EUVP^[30]数据集中的水下场景（underwater-scenes）数据子集，将此水下场景数据子集记为 EUVP-US，EUVP-US 含有2 185组有参考数据，随机选取其中1 967组作为训练数据集，其余218组作为测试数据集。UIEB和EUVP-US的失真图像和参考图像如图3所示。

实验采用PyTorch 1.8深度学习框架。优化器选择Adam，学习率设置为0.0002，批数据量设置为4；在UIEB和EUVP-US数据集上的迭代轮数为400，在训练过程中将图像尺寸统一为256×256；损失的权重λ_s、λ_e与λ_f理论上可以选取大于0的全体实数，因此是无法枚举所有情况的，根据经验与实际调试将λ_s、λ_e与λ_f的比例设置为1:2:1。阶段1～阶段5的Transformer模块数量分别设置为4、8、8、4、4，每个阶段的注意力模块比例为$\frac{1}{4}$。

作为对比试验的评估算法包括 UDCP^[3]、IBLA^[4]、ULAP^[5]、绿色蓝色增强（green blue enhancement，GB）^[31]、UResNet^[6]、UWNet^[7]与WaterNet^[20]。非基于深度学习的增强算法实现来自参考文献[32]；基于深度学习的对比算法用到的批数据量、学习率和迭代轮数等设置均与 UWT-Net保持一致。

评估指标包括峰值信噪比（peak signal-to-noise ratio，PSNR）、SSIM 与无参考水下图像质量度量（underwater image quality measure，UIQM）^[33]。根据参考文献[18]及其中提供的均方误差（mean square error，MSE）计算方式，PSNR的计算方式如式（10）所示：

其中，k为比特数，对于当前数据集，k=8。SSIM是结构信息复原效果的评估指标。UIQM被广泛用于水下图像增强的质量评估任务，它主要包括3个角度的评估，第1是水下图像的色彩评估，第2是水下图像的锐度评估，第3是水下图像的对比度评估，将3种评估值进行线性求和可以获得最终的UIQM值。

在UIEB数据集上进行了定量实验与增强效果的视觉评估，定量实验的结果见表1。表1结果表明了UWT-Net在SSIM和UIQM上获得的指标都是最高的，分别为0.870与3.084，PSNR值为第二高，该结果证明了算法获得的增强图像与参考图像较为接近。

图4展示了不同算法在UIEB数据集上获得的增强效果，其中图4（a）为待增强的失真图，图4（j）为参考图。为了便于分析，将含有鱼的图像记为场景1，含有虾的图像记为场景2。UDCP导致场景整体偏绿色；IBLA 造成图像偏淡蓝色，场景 1 中的鱼和场景 2 中的虾颜色明显失真；ULAP 对场景 2的增强效果较好，但是处理后的场景 1 明显偏红；GB获取的场景1相对接近参考图像，但是场景2与参考图像具有一定的色彩差异；UResNet 处理后的场景颜色偏白，场景的细节变得模糊；UWNet对场景1和场景2的处理均出现一定的对比度降低现象；WaterNet处理后的场景1与场景2与参考图像存在一定的色彩差异；UWT-Net 获得的增强图像与参考图像在颜色和对比度方面最为接近，同时保持了图像细节的完整性。

EUVP-US数据集上的定量实验结果见表2，表2结果显示， UWT-Net能够获得最高的PSNR和SSIM值，分别为28.343和0.871，同时UIQM也相对较高。总体上看，UWT-Net的定量实验结果是最好的。

EUVP-US测试集的增强效果如图5所示，图5的上图和下图分别记为场景3和场景4。UDCP处理后的两个场景整体颜色偏绿，场景4中的海星与环境的对比度降低；IBLA 获得的增强后的场景 3明显偏绿且亮度降低；ULAP处理后的图像相比参考图像存在一定的色差；GB 获得的增强结果明显偏红，颜色失真明显；UResNet与UWT-Net获得的视觉效果较为接近，在色彩方面与参考图像保持一致；UWNet处理后的场景3上半部分颜色偏绿，与参考图像差异明显；WaterNet处理后的图像整体偏绿；UWT-Net 获得的结果最接近参考图像。EUVP-US 测试集的视觉增强效果证明了 UWT-Net的有效性。

针对 UWT-Net 的消融实验分为两部分：第一部分是 UWT-Net 训练过程所用损失的对比研究，第二部分是网络结构的对比研究。消融实验使用的是UIEB数据集。

（1）损失的消融实验

实验中采用多损失组合优化的方式对神经网络的参数进行更新，分别为像素损失、结构损失、边缘损失和特征损失。像素损失是主体损失，其余3种损失是辅助损失。为了验证所选择的损失是有效的，需要对损失组合情况下获得的水下图像增强效果进行量化分析。针对损失的消融实验分为以下4组。

· 设置1：L1损失+结构损失+边缘损失+特征损失。

· 设置2：L1损失+结构损失+边缘损失。

· 设置3：L1损失+边缘损失+特征损失。

· 设置4：L1损失+结构损失+特征损失。

设置1包含了所有损失项，设置2、设置3和设置4中分别在设置1的基础上去掉一种损失，目的是验证去掉的损失对水下图像增强任务是否具有促进作用。

L1 损失消融实验结果见表3，从 PSNR、SSIM与UIQM这3种指标值可以看出，使用设置1获得的增强性能最好，此时3种辅助损失都被使用。实验结果表明，获得的最高的PSNR值为21.163，最高的SSIM值为0.870，最高的UIQM值为3.084。

（2）网络结构的消融实验

UWT-Net采用了并行连接的Transformer模块。为了验证PCTB是否能够提升模型的性能，对比了是否使用PCTB获得的表现。根据表1，使用 PCTB 获得的 3 种指标值分别为 21.163、0.870与3.084。将PCTB换成单层下采样模块加Transformer 模块后，测得的 3 种指标值分别为20.595、0.857与3.021，证明了PCTB能够提升增强结果。

为了证明UWT-Net的收敛性，对水下图像增强网络的训练过程中各项损失值进行记录，包括整体损失L、像素损失 L_p、结构损失 L_s、边缘损失 L_e和特征损失 L_f。整体损失 L 的收敛情况如图6 所示，从图6 可以看出，损失值在迭代 300轮时已经趋近于收敛，当迭代轮数达到 400 时基本处于稳定收敛状态。水下图像复原任务将数据分为训练集和测试集，损失的计算是在训练集上进行的，针对训练集的过度优化可能会导致模型在测试集上的效果降低，这种情况即机器学习中的过拟合现象。为了防止模型出现过拟合现象，在400轮时需要停止对UWT-Net的训练。4种损失的收敛情况如图7所示，像素损失收敛最快，边缘损失与结构损失的收敛趋势比较接近，特征损失基本趋近于0。从图6和图7的损失收敛情况可以看出，增强网络的训练过程稳定，损失的收敛趋势明显。

本文针对水下图像的颜色失真与对比度降低问题，将视觉Transformer模型作为基础网络模块，构建了端到端的水下图像增强模型UWT-Net。模型的训练过程将4种不同类型的损失组合作为图像增强网络的整体损失，4 种损失分别从像素、结构、边缘和特征的角度设计，能够有效地对增强网络的参数更新过程进行指导。本文在两个公开的水下图像数据集上进行了定量和定性的实验，将提出的水下图像增强模型与其他水下图像增强模型进行对比与分析，实验结果表明，本文设计的模型能够获得较高的峰值信噪比，并保证增强后的水下图像与参考图像具有较高的结构相似性。