VMA-SWN的整体结构如图1所示, 图中数值表示水下图像与特征图的长宽信息, 网络输入失真的水下图像, 输出增强后的水下图像.记失真水下图像为x, 网络增强后的输出图像为y, 特征图为Ω _i, i表示某层的索引, 卷积操作记为C_s(· ), 池化操作记为P_s(· ), 双三次插值操作记为B_s(· ), s表示步长.VMA-SWN主要包含如下5个模块.

图1

Fig.1

	Figure Option ViewDownloadNew Window
	图1 VMA-SWN网络结构图Fig.1 Network structure of VMA-SWN

1)卷积模块(Convolution Block, CB).第1层CB为

Ω ₁=ReLU(C₁(x)),

最后一层CB为

y=ReLU(C₁(Ω _i)).

2)下采样模块(Down-Sample Block, DB).采用

Ω _i+1=P₂(Ω _i)

实现, 用于VMA-SWN的特征图降维过程.VMA- SWN共包含3个DB, 特征图尺寸从256× 256降至32× 32.

3)跨层白化模块(Skip-Layer Whitening Block, SWB).采用实例白化^[31]实现, 用于空间特征的跨层传递.SWB可以直接嵌入浅层与深层的视觉多头自注意力模块之间, 不需要进行额外的网络修改.以图1左侧第1次跨层路径为例, 该过程为

DB→ SWB→ C→ RDB→ TAB.

DB模块输出特征图尺寸为128× 128× D.首先将该特征图输入SWB进行实例白化.然后将实例白化后的特征与解码特征进行通道拼接操作, 拼接后的特征图尺寸为128× 128× 2D.最后通过降维模块(Reduce Dimension Block, RDB)对尺寸为128× 128× 2D的特征图进行通道降维, 降维后特征图尺寸为128× 128× D, 继续传入下个模块中.

4)基于Transformer的注意力模块(Transformer-Based Attention Block, TAB)与降维的基于Transfor-mer的注意力模块(Reduced Dimensionality TAB, RTAB).TAB采用视觉Transformer中基于窗口的多头自注意力机制实现.RTAB由RDB和TAB构成.RDB通过卷积计算C₁(Ω _i), 降低跨层白化后特征图的维度, 再将特征图传入TAB.RTAB的实现如下:

Ω _i+1=TAB(C₁(Ω _i)).

此时Ω _i的通道数为2D且Ω _i+1的通道数为D.

5)插值模块(Interpolation Block, IB).采用

Ω _i+1=ReLU(C₁(B₂(Ω _i)))

实现, 用于解码过程的特征图升维过程.VMA-SWN共包含3个IB, 特征图尺寸从32× 32升至256× 256.

上述5个模块可构建VMA-SWN, 编码路径与解码路径均含有4个TAB.为了便于下文分析, 将不含跨层白化模块的增强网络记为VMA-SWN-NoSWB.

采用视觉多头自注意力实现的TAB与RTAB是VMA-SWN的基本模块, 计算原理来自视觉注意力模型^{[24, 25]}.首先, 自注意力(Self-Attention, SA)^[24]计算方式如下:

$SA\left( Q, K, V \right)= {softmax}\left( \frac{Q{{K}^{T}}}{\sqrt{d}} \right)V$,

其中, 输入矩阵为Q、K、V, 分别表示查询、键和值, d表示维度缩放值.在计算SA时将相对位置偏差B加入上式中, 即

$SA\left( Q, K, V \right)={softmax}\left( \frac{Q{{K}^{T}}}{\sqrt{d}}+B \right)V$.

构建TAB模块时采用多头自注意力机制, 进行n次SA计算, 即

$MSA\left( Q, K, V \right)=\tau \left[ hea{{d}_{1}}, hea{{d}_{2}}, \ldots , hea{{d}_{c}}, \ldots hea{{d}_{n}} \right]{{W}^{O}}$,

其中, τ [· ]表示拼接操作, W^O表示投射矩阵,

$hea{{d}_{c}}=SA\left( QW_{c}^{Q}, KW_{c}^{K}, VW_{c}^{V} \right)$,

WcQ、 WcK、 WcV均表示投射矩阵.文献[32]指出Swin Transformer^[25]中的循环移位(Cyclic Shift)适用于高级视觉任务, 因为高级视觉任务中图像中心像素的作用大于图像边缘像素, 因此提出在图像复原任务中使用反射填充(Reflection Padding, RP)替代循环移位.本文网络在计算MSA(Multi-head Self Attention)时采用RP方式, 通过对特征图进行反射填充, 保证特征图长宽是窗口尺寸的整数倍数, 即为RPMSA(Reflection Padding MSA).TAB在特征传递过程中采用层归一化(Layer Normalization)函数ψ (· )进行特征归一化, 第i层的特征Ω _i的计算需要使用前一层特征Ω _i-1, 首先计算中间特征:

$\Omega _{i}^{* }=RPMSA\left( \psi \left( {{\text{ }\!\!\Omega\!\!\text{ }}_{i-1}} \right) \right)+{{\text{ }\!\!\Omega\!\!\text{ }}_{i-1}}$. (1)

然后使用多层感知机φ (· )获取第i层的特征:

${{\text{ }\!\!\Omega\!\!\text{ }}_{i}}=\varphi \left( \psi \left( \text{ }\!\!\Omega\!\!\text{ }_{i}^{* } \right) \right)+\text{ }\!\!\Omega\!\!\text{ }_{i}^{* }=\left[ GELU\left( \psi \left( \text{ }\!\!\Omega\!\!\text{ }_{i}^{* } \right){{W}_{1}}+{{b}_{1}} \right){{W}_{2}}+{{b}_{2}} \right]+\text{ }\!\!\Omega\!\!\text{ }_{i}^{* }$, (2)

其中, W₁、W₂表示权重, b₁、b₂表示偏差项, GELU(· )表示激活函数.式(1)与式(2)为TAB的计算过程, RTAB中的TAB计算过程与此相同.

水的类型包括浅海、深海以及浑水等^[20], 如图2所示, 海水深度影响光的吸收程度, VMA-SWN- NoSWB的特征主要通过前向计算中堆叠TAB和上下采样模块获得, 未充分考虑水的类型的影响.

图2

Fig.2

	Figure Option ViewDownloadNew Window
	图2 水下环境光吸收示意图Fig.2 Sketch map of light absorption in underwater environment

为了降低水的类型多样性造成的影响^[20], 并加强特征信息复用, VMA-SWN通过实例白化(IW)^[31]将浅层特征嵌入深层网络.记VMA-SWN第i层的批数量为N且长为M、宽为L的批次特征为Ω _i∈ R^{C× NML}.对于第k个特征信息 Ωik∈ R^{C× ML}, 其中k=1, 2, …, N, 均值向量为:

$\mu =\frac{1}{ML}\text{ }\!\!\Omega\!\!\text{ }_{i}^{k}$.

完成均值向量的计算后, 协方差矩阵为:

$\text{ }\!\!\Lambda\!\!\text{ }=\frac{1}{ML}\left( \text{ }\!\!\Omega\!\!\text{ }_{i}^{k}-\mu \right){{\left( \text{ }\!\!\Omega\!\!\text{ }_{i}^{k}-\mu \right)}^{T}}+\alpha I$.

为了防止Λ 为奇异矩阵, 在Λ 的计算过程中添加α I, I的定义与文献[20]保持一致, α 表示数值较小的正数, 设定为0.000 01.Λ 对角线元素为各个通道的方差, 非对角线元素表示通道间的相关性.实例白化变换的计算如下:

$\text{IW}\left( \text{ }\!\!\Omega\!\!\text{ }_{i}^{k} \right)={{\text{ }\!\!\Lambda\!\!\text{ }}^{-\frac{1}{2}}}\left( \text{ }\!\!\Omega\!\!\text{ }_{i}^{k}-\mu \right)$.

在上述计算中满足的关系是

$\text{IW}\left( \text{ }\!\!\Omega\!\!\text{ }_{i}^{k} \right)\text{IW}{{\left( \text{ }\!\!\Omega\!\!\text{ }_{i}^{k} \right)}^{T}}=I$.

对浅层特征 Ωik进行实例白化处理, 通过跨层连接拼接浅层特征IW( Ωik)与深层特征 Ωjk, 获得强化后的解码特征

${{\left( \text{ }\!\!\Omega\!\!\text{ }_{j}^{k} \right)}^{\tau }}=\tau \left( \text{IW}\left( \text{ }\!\!\Omega\!\!\text{ }_{i}^{k} \right), \text{ }\!\!\Omega\!\!\text{ }_{j}^{k} \right)$,

其中, j表示与i对应的深层特征的索引, τ (· )表示通道拼接操作.SWB即为对批次特征Ω _i执行实例白化.

由于VMA-SWN架构的有效性, 采用常规的损失函数对VMA-SWN进行训练即可获得较优的水下图像增强效果.将VMA-SWN表示的映射函数记为Ψ (· ), 原始失真水下图像为x(长宽分别为H与W), 标签水下图像为y^* .第1部分内容损失为:

${{L}_{mse}}=\frac{1}{H\times W}\underset{a=1}{\overset{H}{\mathop \sum }}\, \underset{b=1}{\overset{W}{\mathop \sum }}\, {{\left\| \text{ }\!\!\Psi\!\!\text{ }\left( x\left( a, b \right) \right)-{{y}^{* }}\left( a, b \right) \right\|}_{2}}$, (3)

其中a、b表示位置索引.

第2部分是结构损失, 采用结构相似性度量(Structural Similarity, SSIM)^[33]实现, 对VMA-SWN获得的增强图像Ψ (x)与参考图像y^* 计算SSIM值, 得到SSIM(Ψ (x), y^* ), 则结构损失定义为

L_ss=1-SSIM(Ψ (x), y^* ).(4)

VMA-SWN的整体损失为式(3)与式(4)损失之和:

$L={{L}_{\text{mse}}}+\lambda {{L}_{\text{ss}}}$,

其中, λ 为权重因子, 用于控制内容损失与结构损失在VMA-SWN的网络参数更新过程中所占的比重.

水下图像增强算法的效果评估可以从定量结果与视觉效果两方面进行, 本文选用水下图像数据集UIEB^[22], 并从水下图像数据库EUVP^[13]上选择水下图像集.

UIEB数据集图像包含不同特性的质量退化, 标签图像的构造采用12种增强方法并结合人工视觉筛选过程, 训练数据为800组, 验证数据为90组.在EUVP数据集上选择水下图像集(Underwater ImageNet, UWIN), 该数据集标签的构造主要通过循环一致对抗网络实现, 训练数据为3 300组, 验证数据数为400组.

实验平台为NVIDIA V100 32 GHz GPU, 并采用PyTorch 1.8深度学习框架.

为了公平对比不同算法, 需要使用量化评估指标对比不同算法增强后的水下图像质量, 水下图像增强任务常用的评估指标包含有参考与无参考两类, 本文选取两种常用的有参考评估指标:峰值信噪比(Peak Signal-to-Noise Ratio, PSNR)和SSIM, 二者的计算采用增强后的水下图像与清晰标签^[22], 并计算无参考评估指标— — 水下图像质量度量(Underwater Image Quality Measure, UIQM)^[34].UIQM主要适用于无清晰标签的情况.无参考评估指标的可靠程度低于有参考评估指标^[34].

为了分析VMA-SWN与已有水下增强算法的性能差异, 选择如下11种对比算法.1)4种非基于深度学习的增强算法.文献[8]算法、ULAP^[9]、文献[10]算法和文献[11]算法.2)7种基于深度学习的增强算法.文献[5]算法、UResnet^[12]、FUnIE-GAN^[13]、UGAN^[14]、ADMNNet^[17]、Shallow-UWnet^[21]、 Water- Net^[22].4种非基于深度学习的增强算法采用OpenCV 4.7图像处理框架实现, 7种基于深度学习的增强算法与VMA-SWN采用PyTorch 1.8深度学习框架实现.

对比实验中11种算法所用的图像尺寸统一缩放为256× 256× 3.对于7种基于深度学习的增强算法和VMA-SWN, 训练阶段的批大小尺寸统一设定为6, 在UIEB数据集上运行200个迭代周期, 在UWIN数据集上运行100个迭代周期.7种基于深度学习的增强算法的学习率初始值与学习率衰减策略均按照对应文献进行设定.在训练过程中VMA-SWN学习率设定为0.000 2, λ 设定为1, 每隔20个迭代周期学习率衰减0.95倍.

在UIEB数据集上对比VMA-SWN与11种算法, 结果如表1所示, 表中黑体数字表示最优值.由表可见, VMA-SWN获得最高的PSNR与SSIM值, 分别为21.862 dB、0.909, 并且获得相对较高的UIQM值, 为2.997.ADMNNet、UGAN与Water-Net获得的PSNR值与VMA-SWN接近, 但是SSIM值明显低于VMA-SWN.UGAN获得的UIQM值最高, VMA-SWN获得的UIQM值位列第3.总之, 表1结果从定量评估角度表明VMA-SWN在UIEB数据集上的有效性.

表1

Table 1

表1(Table 1)

表1 各算法在UIEB数据集上的定量评估结果 Table 1 Quantitative evaluation results of different algorithms on UIEB dataset 算法 PSNR/dB SSIM UIQM ADMNNet 21.138 0.882 2.887 FUnIE-GAN 18.707 0.800 3.081 文献[8]算法 15.470 0.678 1.762 文献[10]算法 13.840 0.605 1.734 文献[5]算法 18.386 0.802 2.785 文献[11]算法 16.647 0.764 2.050 Shallow-UWnet 17.453 0.794 2.773 UGAN 21.472 0.860 3.276 ULAP 14.640 0.694 1.796 UResnet 18.324 0.841 2.892 Water-Net 21.198 0.858 2.978 VMA-SWN 21.862 0.909 2.997

表1 各算法在UIEB数据集上的定量评估结果 Table 1 Quantitative evaluation results of different algorithms on UIEB dataset

各算法在UIEB数据集上视觉效果对比如图3所示.将图3中的3个场景分别记为场景1、场景2与场景3.非基于深度学习的增强算法(文献[8]算法、文献[10]算法与ULAP)在场景1和场景2上存在明显的色差现象; 文献[8]算法与ULAP处理后的场景3亮度偏移较大, 导致场景整体呈现暗色; 文献[10]算法处理后的场景3与标签差异较大.基于深度学习的增强算法(ADMNNet、FUnIE-GAN、文献[5]算法、UResnet)处理后的场景1颜色偏移明显, 未能有效对颜色进行校正.UResnet与Water-Net处理后的场景2存在一定的颜色偏移, 且UResnet处理后的场景2对比度偏低.文献[5]算法与UResnet处理后的场景3亮度相对较低.

图3

Fig.3

	Figure Option ViewDownloadNew Window
	图3 各算法在UIEB数据集上的视觉效果对比Fig.3 Visual effects comparison of different algorithms on UIEB dataset

总之, 场景1的处理难度较大, 尽管VMA-SWN在场景1存在轻微的色偏, 但是处理效果最接近标签.相比11种对比算法, VMA-SWN在3个场景的处理效果都相对较优, 能够有效进行颜色校正并提升3个原始水下场景的视觉质量.

VMA-SWN与11种算法在UWIN数据集上的定量评估结果如表2所示, 表中黑体数字表示最优值.

表2

Table 2

表2(Table 2)

表2 各算法在UWIN数据集上的定量评估结果 Table 2 Quantitative evaluation results of different algorithms on UWIN dataset 算法 PSNR/dB SSIM UIQM ADMNNet 24.668 0.844 3.000 FUnIE-GAN 21.899 0.796 3.032 文献[8]算法 15.688 0.633 2.080 文献[10]算法 15.591 0.626 2.069 文献[5]算法 23.583 0.814 2.983 文献[11]算法 16.083 0.651 2.195 Shallow-UWnet 22.485 0.815 2.960 UGAN 24.586 0.833 3.053 ULAP 17.154 0.685 2.194 UResnet 23.327 0.825 3.016 Water-Net 23.673 0.832 2.991 VMA-SWN 24.994 0.855 2.931

表2 各算法在UWIN数据集上的定量评估结果 Table 2 Quantitative evaluation results of different algorithms on UWIN dataset

由表2可知, ADMNNet与UGAN获得的PSNR值与VMA-SWN相对接近, 其余算法的PSNR值均未超过24 dB.在SSIM指标上只有VMA-SWN超过0.85, 明显高于其余算法.在UIQM指标上, UGAN获得最高值, 为3.053, VMA-SWN的UIQM值为2.931, 但是UGAN获得的PSNR与SSIM值均低于VMA-SWN.

总之, VMA-SWN的定量评估结果相对较优.

各算法在UWIN数据集上视觉效果对比如图4所示.将图4中3个场景由左到右分别记为场景4、场景5与场景6.原始水下图像呈现偏蓝色的状态, 并且与参考图像(标签)的颜色差异明显.非基于深度学习的增强算法(文献[8]算法、文献[10]算法与ULAP)在各场景上增强效果不明显, 图像仍然呈现偏蓝色状态, 文献[11]算法对场景4的处理存在细节丢失与色彩失衡现象, 并且对场景5与场景6的处理效果轻微.基于深度学习的增强算法(ADM- NNet与UGAN)处理后的场景4偏向淡紫色, FUnIE-GAN、文献[5]算法、Shallow-UWnet与URes- net处理后的场景5仍然没有实现较好的颜色复原, Shallow-UWnet与Water-Net处理后的场景6偏向淡蓝色.VMA-SWN处理后的图像最接近相应标签.

图4

Fig.4

	Figure Option ViewDownloadNew Window
	图4 各算法在UWIN数据集上的视觉效果对比Fig.4 Visual effects comparison of different algorithms on UWIN dataset

VMA-SWN含有两个需要进行消融研究的组件, 分别是跨层白化模块SWB与结构损失L_ss, 需要对比SWB与L_ss使用前后VMA-SWN的增强效果变化情况.消融实验采用每次删除一个模块的方式, 共含有下面3种设置.

1)L_mse+L_ss+SWB.使用L_mse与L_ss训练VMA-SWN.

2)L_mse+L_ss.同时使用L_mse+L_ss训练VMA-SWN-NoSWB.

3)L_mse.使用L_mse训练VMA-SWN-NoSWB.

为了保证结果的准确性, 采用PSNR与SSIM作为VMA-SWN量化评估依据, 实验结果如表3所示, 表中黑体数字表示最优值.由表可见, 在UIEB、UWIN数据集上效果最优的是方案1), 方案3)的SSIM值明显低于方案1).该结果表明SWB与L_ss对VMA-SWN的增强效果具有促进作用.

表3

Table 3

表3(Table 3)

表3 VMA-SWN的消融实验结果 Table 3 Results of ablation experiment for VMA-SWN 指标 UIEB UWIN Lmse+Lss+SWB Lmse+Lss Lmse Lmse+Lss+SWB Lmse+Lss Lmse PSNR/dB 21.862 20.437 20.270 24.994 24.288 24.016 SSIM 0.909 0.897 0.879 0.855 0.852 0.839

表3 VMA-SWN的消融实验结果 Table 3 Results of ablation experiment for VMA-SWN

VMA-SWN采用TAB与RTAB作为主要的特征提取模块, TAB与RTAB的特征维度D是一个重要的超参数, 影响VMA-SWN的增强效果与模型大小.D过小可能导致模型的数据学习能力不足, D过大会导致模型的参数量较大.结合已有文献[13]和文献[32]中特征维度的选择方式, 将D的取值设定为16的整数倍数, 分别设定D=32、48、64、96、128, 以此计算相应的评估指标并获取模型参数量.

各算法在UIEB数据集上的结果如图5所示, PSNR与SSIM均取小数点后5位作为有效位.由图可见, D在32、48、64、96与128之间选取(二者维度需要保持一致)时, 定量评估结果存在一定差异, PSNR在21.6 dB左右波动, SSIM在0.909左右波动.根据结果可知, D=64时为较优选择, 此时VMA-SWN获得的PSNR与SSIM值相对较高, 通过计算, 此时模型参数量为8.44 MB.

图5

Fig.5

	Figure Option ViewDownloadNew Window
	图5 D不同时指标值的变化Fig.5 Variation of evaluation metrics with different D

为了验证VMA-SWN训练过程的稳定性, 绘制L_mse与L_ss在训练过程中的损失曲线, 每隔10次迭代记录过去10次迭代的损失均值作为一个数据点, 获取VMA-SWN在UIEB、UWIN数据集上损失函数的收敛情况, 具体如图6与图7所示.

图6

Fig.6

	Figure Option ViewDownloadNew Window
	图6 VMA-SWN在UIEB数据集上的损失收敛曲线Fig.6 Loss convergence curve of VMA-SWN on UIEB dataset

图7

Fig.7

	Figure Option ViewDownloadNew Window
	图7 VMA-SWN在UWIN数据集上的损失收敛曲线Fig.7 Loss convergence curve of VMA-SWN on UWIN dataset

由图6和图7可以看出, L_mse与L_ss在总体上都呈现稳定的收敛趋势, 表明VMA-SWN训练过程的稳定性.