AL-GAN:一种融合注意力机制的轻量级GAN水下图像增强模型

冯建新,韩亚军,潘成胜,孙传林,蔡远航

(1.辽宁省通信网络与信息处理重点实验室, 辽宁 大连 116622; 2.大连大学 信息工程学院, 辽宁 大连 116622)

0 引言

海洋有着很多重要的信息资源,这些信息资源对研究海洋生物以及勘探资源有着重要的意义,因此安全可靠地获取海洋信息资源非常重要。然而海洋环境不同于陆地,水下环境情况复杂,且可见光在水下会遇到吸收和散射等问题,这些会造成采集到的水下信息(如视频或图像)会出现偏色、细节模糊等问题,这类问题会对后续水下目标跟踪和检测造成严重的影响。因此需要提高水下图像质量,消除偏色、细节模糊等问题。

用于图像增强的方法根据成像模型不同可以分为物理模型和非物理模型2种。水下图像增强的物理模型是根据光学原理建模,主要进行去噪和颜色校正,如暗通道先验(DCP)[1],水下DCP等。水下图像的精确恢复取决于模型的先验知识,然而先验知识在水下场景不同时往往鲁棒性较弱,从而会导致严重的估计偏差。同时,在实际应用中水下深度和光传播系数基本参数也很难获得。水下图像增强的非物理模型不考虑水下图像的物理退化机制,它是通过修改退化图像的具体像素值增强图像效果,主要增强方法如直方图滑动拉伸[2]、多尺度融合[3]等。此类方法不依赖于物理成像模型,因此这些方法通常不足以恢复原始场景特征,尤其是颜色特征。

由于深度学习可以实现复杂非线性系统的端到端建模,能够提高感知图像的质量,目前已经在去噪、去模糊、感知增强和对比度调整等低级视觉任务中取得了令人信服的成功[4-5]。尤其是近些年生成对抗网络(generative adversarial networks,GAN)[6]发展迅速,因此越来越多的研究人员将目光转移到了应用深度学习进行水下图像增强的研究中。JieLi等[7]提出了一种无监督生成对抗网络Water-GAN实现对单目图像水下偏色问题的矫正,它具有很好的实时性。Fabbri等[8]用扩充数据集方法提出适用于水下场景的UGAN(underwater GAN),解决了水下图像数据集缺乏问题,并且能够有效地提升水下图像的对比度和清晰度。Islam等[9]提出了一种基于GAN实时性很高的水下图像增强模型FUnIE-GAN,该模型根据图像的整体内容、细节样式、风格和局部纹理信息建立损失函数,可以实时增强水下图像,网络模型推理速度较快,但是针对不同的水体情况处理效果不好,例如过度增强或增强不足,同时模型鲁棒性会降低。Zhang等[10]改进原有生成对抗网络,为了避免图像细节模糊,通过引入梯度损失锐化图像边缘提高图像清晰度,但是在处理偏色问题上会有颜色失真问题。Hambarde等[11]提出一种端到端的水下生成对抗网络UW-GAN,能够完成单个水下图像深度预测,进一步提出了一种用于大规模数据库合成水下图像生成方法。UW-GAN网络模型在U-Net框架生成网络中使用VGG16模型作为编码器,虽然图像增强效果较好,但模型参数量较大,网络模型的推理速度很慢。

上述基于GAN的模型直接学习退化图像到清晰图像的映射,针对不同水体情况无法做到增强水下图像质量的同时提高网络模型的推理速度。为了做到增强不同水体的水下图像质量,兼顾网络模型推理速度的提升,本文中提出了一种融合注意力机制的轻量级GAN模型,以下是本文中的主要研究内容:

1) 提出了一种轻量级GAN模型,生成网络在FUnIE-GAN模型基础上使用MobileNet替换原U-Net生成网络参数量极大的VGG16模型,减少了模型参数量,从而解决生成网络推理速率慢的问题。与其他模型比较推理速度提高了1.2～3倍。

2) 在特征提取网络中引入了通道和空间双重注意力机制,实现了去除水下图像颜色干扰和增强图像细节的目的。

3) 所提出的模型在数据集上表现出更好、更稳定的性能。在EUVP数据集上,与其他算法相比,水下图像评价指标(underwater image quality measure,UIQM)[12]提高约0.21,自然图像评价指标(natural image quality evaluator,NIQE)[13]降低约0.65。

1 网络结构

网络模型包括生成网络和判别网络2部分,如图1所示。本文中生成网络的输入为水下退化图像,判别网络的输入为生成图像和地面真实图像。在模型训练阶段,生成网络和判别网络互相监督,循环交替训练更新网络模型的参数。最终训练出高性能并且稳定的生成网络,用生成网络来增强水下图像。

U-Net[14]是一个包含上采样和下采样的全卷积网络,它通过在上采样过程中组合不同下采样的特征重新构建清晰图像。我们的生成网络受U-Net结构启发,设计为5层编解码结构,编码器模块由深度可分离卷积模块和注意力模块组成,解码器模块由多个上采样卷积模块组合而成。为了减少网络在下采样期间造成的信息丢失,U型网络连接编码器输出作为解码器输入。判别网络模型通过定义PatchGAN[15]实现。

1.1 生成网络模型

1.1.1编码器模块

根据MobileNet[16]模型体积小、精度高特性,深度可分离卷积作为MobileNet的核心组件,由深度卷积和逐点卷积组成。在编码器模块中使用深度可分离卷积替换普通卷积,深度可分离卷积用于特征提取的卷积大小为3×3,用来调整通道数的卷积块大小为1×1。

CBAM[17]注意力模块由通道注意模块(CAM)和空间注意模块(SAM)组成。输入图像的特征首先通过卷积模块提取出来,然后再输入到CAM得到输入特征的权重分布图,权重分布图可以显示输入图像信息的重要度。SAM侧重于输入特征信息在哪里更重要。2种注意机制结合增强了网络的学习功能,使网络更加注重水下图像退化严重部分的信息,从而能够更好地增强图像,本文中CBAM结构如图2所示。

编码器模块经过前3层卷积模块后接入CBAM注意力模块,卷积模块的每一层先是使用大小3×3,步长为2的深度卷积提取特征,然后用1×1卷积调整通道数,每个卷积层后面连接的是归一化BN及ReLU激活函数。为探寻注意力机制和深度可分离卷积模块的最优组合,在3.2节进行了生成网络结构设计实验,实验结果表明,前3层卷积模块和CBAM组合结构是最优组合。

1.1.2解码器模块

解码器由5个上采样卷积模块组成,解码器中前4个卷积模块的每一层采用大小为3×3的滤波器,卷积步长为2。在每个卷积层之后添加BN和激活函数ReLU,最后一个卷积模块将特征图转换为256×256×3输出图像。

图1 网络模型结构

图2 通道空间注意力模块

1.2 判别网络模型

本文中判别网络使用马尔可夫鉴别器(PatchGAN),与常规鉴别器不同,PatchGAN输入是当前输入图像Ii的地面真值图像IGT和其经过生成网络处理得到的输出图像IG,输出大小为16×16特征矩阵,这在判断之前就将图像进行了小块分割,从而提高了鉴别器精度,同时也促使网络在训练过程中可以获得更多的细节信息。本文中判别网络前4层使用3×3的卷积层并进行2倍下采样,在每一卷积层后添加BN层和Leaky ReLU激活层,最后一层在卷积层后添加Tanh激活层,5个卷积层将256×256×6的输入图像Ii(生成图像和真实图像)转换为16×16×1的输出,输出矩阵的每一个元素都表示输入图像的一个感受野,从而可以更好地捕获图像局部特征,如局部纹理和细节形式[18],判别网络结构如图3所示。

图3 判别网络结构

1.3 损失函数

为了更好恢复图像视觉效果,同时尽可能保留细节特征,损失函数包括对抗损失、全局相似性损失、内容损失3个方面,损失函数定义如下:

(1)

式(1)中:LWGAN[19]为对抗损失函数;LL1[20]为全局相似性损失函数;Lcon[21]为内容损失函数;λ1、λ2为平衡各项损失函数的权重因子,其中λ1=0.7,λ2=0.3。G和D分别表示生成网络和判决网络。

1.3.1对抗损失LWGAN

传统GAN损失函数是基于散度做优化,但是在这种优化下网络损失训练会很不稳定,会带来梯度消失等问题,这种问题会导致模型训练崩溃[22]。因此,使用带有梯度惩罚的WGAN(WGAN-GP)[23],对抗损失函数表示为

LWGAN(G,D)=E[D(y)]-E[D(G(x))]+

(2)

1.3.2全局相似性损失LL1

由于L1损失不容易引入模糊,所以使用输出图像Io与生成网络输出图像IG之间L1距离优化鉴别器,从而保证图像信息的一致性,其定义如下:

(3)

1.3.3内容损失Lcon

为使生成图像更加真实,提升增强图像的视觉效果,在目标函数中添加内容损失,本文中使用感知距离作为内容损失形式[24]。引入VGG19预训练网络,通过提取第4次最大池化前第3个卷积层输出的高级特征构造内容损失,定义如下,其中φ代表网络提取到的高级特征。

(4)

2 实验设计

2.1 实验参数

本文中实验是基于Pytorch深度学习框架实现的,实验CPU为12th Gen Intel(R) Core(TM) i5-12500H 3.10 GHz,GPU为NVIDIA GeForce RTX3050Ti。训练模型的优化器采用Adam[25],批处理大小Batch-size为8,初始学习率设为0.01,一共150个epoch,每30个epoch学习率乘以0.5衰减1次。全局相似性损失权重系数λ1为0.7,内容损失权重系数λ2为0.3。

2.2 生成网络结构

从训练过程和训练结果出发,对生成网络U-Net框架编码器部分CBAM模块和卷积模块进行组合训练测试。选择UIQM和网络模型推理速度作为衡量网络性能指标。针对不同的CBAM模块与下采样卷积模块组合,表1是采用EUVP数据集中相同测试集进行测试得到的性能,其中F1表示第1层卷积模块与CBAM模块组合,F2表示第2层卷积模块与CBAM模块组合,F3和F4同理。

从表1可以看出,第2层卷积模块与CBAM模块组合方案F4在模型推理速度上与F1相当,略高于F2,远低于F3,但其UIQM为2.327,说明F4在处理图像效果上明显优于其他组合。

表1 单层卷积模块和CBAM组合结果对比

表2为2层上采样卷积模块和CBAM组合的测试结果。从表2可以看出,2层卷积模块与CBAM组合方案明显比单层卷积模块与CBAM结合方案效果好,其中F2+F3的UIQM可达2.465,但是模型推理速度相对变慢,F2+F4为1.73 s。

表2 两层卷积模块和CBAM组合结果对比

表3为3/4层上采样卷积模块和CBAM组合方案的测试结果。从表3可以看出,3/4层卷积模块和CBAM组合方案的UIQM比2层组合更优,但同时模型推理速度也变慢。3层卷积模块与CBAM组合方案和每一层卷积模块都与CBAM组合方案相比,UIQM值几乎一致,其中F1+F2+F3方案推理速度为1.75 s最快,该值与2层采样卷积模块和CBAM组合方案相近。

综上所述,从推理速度和图像效果出发,本文中生成网络确定采用F1+F2+F3方案,即前3层卷积模块与CBAM组合方案构造生成网络的编码器。

2.3 损失函数性能

图4和图5是在对抗损失、全局相似性损失、内容损失3个方面构造损失函数的性能分析。其中,Lall为3种损失函数组合,LWGAN+L1为对抗损失函数与全局相似性损失函数组合,LWGAN+C为对抗损失函数与内容损失函数的组合;n为迭代次数,loss为损失函数值。

从图4可以看出,只使用对抗性损失函数训练时震动幅度大,收敛过程不稳定,所需要循环次数多。采用3种损失函数组合形式的训练过程更加平稳,离散程度更小。

从图5可以看出,Lall曲线、LWGAN+L1曲线表现相近且损失函数值明显低于LWGAN+C曲线。但随着迭代次数增加,LWGAN+L1曲线收敛后又出现了小幅度的波动,Lall曲线相比于LWGAN+L1曲线更加稳定。综上所述,使用3种损失函数组合Lall具有更好、更稳定的性能。此外,从峰值信噪比(PSNR)和结构相似性(SSIM)[26]可知,3种损失函数组合Lall的PSNR是23.749,SSIM是0.765,明显高于其他损失函数组合,因此,本文中损失函数考虑对抗损失、全局相似性损失、内容损失3个方面,损失函数为3种损失函数组合Lall,如式(1)所示。

图5 3种损失函数训练过程对比

表4为不同损失函数组合训练出的模型在评价指标上的实验结果对比。

表4 不同损失函数组合评价指标实验结果对比

3 结果与讨论

选用水下公开数据集EUVP验证本文中算法的有效性,EUVP数据集包含12 000多对水下图像和其对应的地面清晰图像,选取其中7 000幅水下图像和其对应的清晰图像作为训练集。将本文中算法和传统算法、深度学习算法从主客观两方面做对比,由实验结果可知本文中算法在复杂情况下的水下图像增强效果更好。

3.1 主观评价

为了验证本文中算法对水下真实图像增强效果,测试数据随机选取150张真实水下图像。图6为测试数据部分图像在不同算法上的效果图,对比的传统算法包括水下暗通道先验(UDCP),基于图像模糊的水下图像恢复(IBLA)[27],基于快速场景深度估计模型的水下图像恢复(ULAP)[28];对比的深度学习算法包括基于对抗的水下图像增强网络(UGAN),基于残差网络的生成模型(DeepSESR)[29],基于GAN网络的快速水下图像增强(FUnIE-GAN)。图6(a)为输入图像,图6(I)为真实水下图像,图6(b)—图6(h)分别是UDCP、IBLA、ULAP、UGAN、Deep-SESR、FUnIE-GAN以及本文中方法的处理结果图。可见本文中方法相比于其他方法结果更加贴近真实图像,解决了水下图像由于水体原因造成的偏色问题,处理得到的图像有更高的对比度和清晰度,色彩更加鲜艳。

图6 不同方法处理结果对比

3.2 客观评价

客观地分析和评价算法性能,使用水下图像质量指标UIQM和自然图像评价指标NIQE作为评价指标。UIQM是由水下图像色彩测量UICM,水下图像清晰度测量UISM,水下图像对比度测量UIconM,3个评价指标组成的。UIQM的值越大,表明图像的质量越高,其计算公式为

UIQM=c1×UICM+c2×UISM+c3×UIconM

(5)

根据文献[30]设置c1=0.028 2,c2=0.295,c3=3.575。

NIQE无须利用人眼评分的失真图像进行训练,NIQE相比于传统评价指标PNSR和SSIM能更有效地反映图像重建的质量。NIQE指标的值越小,表明图像的质量越高,其计算公式为

D(v1,v2,∑1,∑2)=

(6)

式(6)中:v1和v2分别表示自然多元高斯模型和失真图像多元高斯模型的均值向量;∑1和∑2分别表示自然MVG模型和失真图像MVG模型的协方差矩阵。

对150张水下图像进行评估,表5为各种方法处理图像的评价指标UIQM和NIQE。可以看出,经本文中算法处理的图像评价指标最优。与6种对比算法相比,本文中算法处理的图像UIQM明显得到提升,同时NIQE明显降低。与FUnIE-GAN算法相比,UIQM提升了约0.21 ,NIQE降低了约0.65。

表5 UIQM和NIQE的定量比较

综上所述,根据各种算法在无参考图像质量评价指标UIQM和NIQE上的对比,证明本文中算法能够更好的恢复水下图像,提高了图像的清晰度和对比度。

3.3 消融实验

表6为采用VGG16模型做编码器和本文中编码器在处理单幅图像所需时间。随机选取5张图片,分别在VGG16和MobileNet为编码器的网络中测试,结果表明,本文中编码器比VGG16编码器运算速度提高很大,表明MobileNet替换VGG16的有效性。

表6 不同编码器运行时间对比(s)

图7为有无CBAM模块处的图像对比,从图7前2列对比图可以看出,无CBAM模块处理的图像较有CBAM模块处理的图像虽然解决了水下图像偏绿色问题,但是增强后的图像明显出现偏黄。从后面3列对比图可以看出,无CBAM模块处理的图像在细节上虽然得到增强,但是相比于有CBAM模块处理仍然有模糊问题。有CBAM模块处理的图像,不仅解决了色偏问题,而且清晰度也得到提升,因此本文中模型引入了CBAM模块。

图7 有无通道空间注意力模块结果对比

表7用客观指标UIQM和NIQE进行了分析,实验结果如表7所示。从表7可见,有CBAM模块处理的图像UIQM明显比无CBAM模块处理的高,同时NIQE明显更低,定量说明有CBAM模块处理的图像质量更高。

表7 有无CBAM模块UIQM/NIQE评价指标结果对比 Table 7 Comparison of UIQM/NIQE evaluation index results with or without CBAM module

3.4 运行时间

进一步,为验证本文中提出的GAN水下图像增强模型轻量化特征,对方法运行时间进行了估计处理,其中Deep-SESR、UGAN和FUnIE-GAN是基于深度学习的方法,结果如表8所示。该结果为CPU处理1张大小256×256图像所需要的时间,可以看出本文中方法虽然不是最快的,但是处理时间与FUnIE-GAN相当。

4 结论

本文中提出一种AL-GAN水下图像增强模型。实验结果表明,相比与对比算法,本文中提出的水下图像增强模型UIQM提高了约0.21,NIQE降低了约0.65,说明本文中模型能够更好地增强图像细节,提升图像的对比度和清晰度,对于水下图像增强效果表现最优,同时网络模型的推理速度同比于其他模型提高了1.2～3倍。

但在不同的水下环境中,水下图像退化的程度也会有所不同。严重退化的图像单独使用深度学习网络处理后,仍然会有严重的颜色干扰,并可能产生冗余噪声点。对深度学习网络添加物理模型可以更好地估计水下图像成像环境,进一步,通过环境参数约束网络可以更好地恢复出清晰的水下场景图像。在未来的工作中,我们将尝试结合物理模型,进一步关注退化严重的水下图像处理。