基于UNDERWATER-CUT 模型的水下图像增强算法

姚鹏,刘玉会

(中国海洋大学 工程学院,山东 青岛,266100)

0 引言

视觉传感器可以提供丰富的环境色彩信息,利用视觉信息对障碍物进行检测和定位,进而对周围环境进行三维重建,对水下机器人的导航及探索活动具有重要应用价值。然而,水下环境的图像存在对比度低、清晰度不足等问题,对后续任务造成了阻碍。不同于空气中的不清晰图像(如雾化图像),水中图像主要是由水中折射导致的成像不清晰。光线波长影响光线在水中的衰减,波长越短,传输性能越好,因此波长短的光线占成像光线的大部分,导致了水下图像大多呈蓝绿色[1],难以直接对水下图像进行目标识别等后续任务。同时水下存在大量漂浮颗粒物,导致环境光和背景光被吸收,进一步加剧了水下图像对比度低和噪声严重等问题,因此必须对水下图像进行增强。

水下图像增强通常采用基于物理模型和非物理模型2 种方式。基于物理模型的图像增强方法对水下图像的退化过程建立数学模型,并估算出模型的基本参数,然后通过该模型反演图像的退化过程,获得理想状态下的未退化图像[1-3]。然而,由于水下环境多样,难以用具体模型建模,所以该类方法泛化性不强,无法修复多种水域图像。

近年来,随着深度学习的发展,基于深度学习的图像增强算法得到了极大的进步。Li 等[4]提出了WaterGAN 两阶段模型,第1 阶段构建生成式模型,根据成对的陆上图像和深度信息,生成对应的水下图像,第2 阶段构建另一个生成模型,将生成的水下图像和原本的陆上图像作为成对的训练集,训练结束后,第2 阶段的生成模型对水下图像进行脱水增强。Li 等[5]提出了一种基于循环神经网络的算法,该方法在循环生成对抗网络(cycle generative adversarial network,CycleGAN)[6]的基础上添加结构相似性(structural similarity,SSIM)损失约束[7],训练时无需水下数据集和脱水数据集成对训练,而且泛化性强,对于不同水域的水下图像,均有不错的增强效果。该方法通过维护水下图像域X到脱水图像域Y和脱水图像域Y到水下图像域X这2 个方向的映射,避免了多个X域图像映射到同一Y域图像的情形,但同时也引入了不必要的映射。为了修正颜色失真的水下图像,仅需要维护从水下图像到脱水图像的单射,而不是2 个域图像的双射。李庆忠等[8]在CycleGAN 模型的基础上添加了图像强边缘结构相似度约束,同时分两阶段训练网络,第1 阶段进行改进的CycleGAN 模型的训练,第2 阶段,进行少量样本的强监督训练,得到了优于文献的结果,但由于引入了强监督的样本,不利于实际场景的应用。

针对文献[5]存在不必要的双射约束以及文献[8]中存在的需要成对样本参与训练的缺陷,文中基于CUT 模型[9]提出了一种水下图像颜色修正的算法(UNDERWATER-CUT 模型)。CUT 模型是CycleGAN 模型的一个增强模型,仅维护水下图像域到脱水图像域的单射,同时只需要非成对样本参与训练,在图像域转换等方面取得很好的效果,但是CUT 模型未针对水下图像特点进行专门优化,对于水下图像增强效果不佳。文中提出的UNDERWATER-CUT 算法为了保证颜色增强后图像中物体的结构相似性,添加了SSIM 损失进行约束。此外,原始CUT 模型使用了InfoNCE 损失函数进行对比学习的约束[10],但是InfoNCE 函数形式复杂,不利于神经网络收敛,因此文中使用alignment 和uniformity 2 个简单函数的线性组合作为损失函数来进行模型的约束[9],得到了优于CUT模型的结果。

1 UNDERWATER-CUT 模型结构

图1 为UNDERWATER-CUT 网络结构,模型主体是对抗式生成网络[11],由生成器和判别器两部分组成,该网络主要作用为通过非成对的水下数据集和脱水数据集(或陆上数据集)的无监督训练,实现水下图像域到脱水图像域的转换。不同于CycleGAN,文中的方法只需单向网络即可保证高效转换。设水下图像域为X,脱水图像域为Y,为了完成X域到Y域的转换,需要同时训练生成器G和判别器D。生成器需要是一个映射G:X→Y,一方面生成器构造映射y˜=G(x),另一方面判别器衡量生成的质量D(y,y˜),指引生成器梯度向更优的方向移动,当判别器不能分辨生成器产生的图像时,此时达到最优。同时,为了防止模型将多个不同X域图像映射为相同的Y域图片,使用对比学习的方式约束图像的生成,强迫模型学习到的映射G:X→Y为X域到Y域的单射。

图1 UNDERWATER-CUT 模型网络结构Fig.1 Network structure of UNDERWTAER-CUT model

生成器的具体结构如图2 所示,总共分为3 部分子网络:第1 部分由3 层卷积层网络组成,主要作用为从 255×255×3的图像中抽取出高维特征;第2 部分由6 层ResNet[12]组成,主要用来对前一阶段提取出来的高维特征进行特征融合,学习到各层网络之间的隐层关系;最后一部分也由3 层卷积神经网络组成,其主要作用为将上游网络层中的高维特征还原至3 个颜色通道的彩色图像。

图2 生成器模型网络结构Fig.2 Network structure of generator model

为了同时捕捉到图像中的颜色高频分量和低频分量,从而完整且客观地权衡图像的真实度,使用文献[13]中的PatchGAN 判别器结构,将图像分为 70×70的块(patch),对于每一块分别预测真假,最终将所有块的预测值取均值得到最终对整幅图像的预测值。这种判别器具体结构如图3 所示,一共有5 层卷积神经网络组成,和用于图像二分类任务的常规模型不同,文中使用的PatchGAN 最后一层不使用全连接神经网络,而是使用卷积层进行预测值的转换,前4 层的主要作用是进行特征提取和融合,最后一层卷积神经网络将隐层特征转化为N×N×1的预测值,文中选取N=70。

图3 判别器模型网络结构Fig.3 Network structure of discriminator model

2 UNDERWATER-CUT 网络损失函数和训练策略

2.1 UNDERWATER-CUT 网络损失函数

损失函数对于网络模型至关重要,合适的损失函数意味着合适的约束,才能指引网络学习到最优的情况,UNDERWATER-CUT 的损失函数主要由对抗生成损失、对比学习损失和SSIM 损失三部分组成。

1) 对抗生成损失

结合非成对数据集的水下图像数据分布与脱水图像数据分布,对抗生成损失的主要作用是使得生成器和判别器达成纳什均衡,最优时生成器可以生成高质量的脱水水下图像,即

2) 对比学习损失

对比学习损失的目的是最大化保留输入的特征信息,将图像分为多个块,同时将生成器的前半部分Genc作为编码器,便于压缩图像为一个向量表示,为了得到不同特征维度的信息,选取了L层来共同进行信息提取,为了最大化保留向量信息,会在每个l∈L层后面添加2 层的多层感知器(multilayer perceptron,MLP)网络Hl,最终将L个编码向量共同组合成为第l层的输出向量,同样的,对于生成器的输出也存在原始CUT 使用InfoNCE函数作为基本损失函数,即

使用图像块作为基本样本,原图像对应位置作为正样本,其他位置作为负样本,同时对于多层进行计算。损失函数为

式(3)中对于每一层,将图像向量分成Sl份,每一份图像块为第l层位置序号为s的生成图像块,为第l层位置序号为s的原图像块,为第l层位置不为s的其他生成图像块。式(2)由于函数形式复杂使得函数优化比较困难,难以收敛至最优值。使用文献[14]提出的性质,InfoNCE 函数存在2 个性质,即相似样本的相互拉近和其他样本的均匀分布,故使用alignment 和uniformity 2 个简单函数的线性组合来达到对比约束的效果,该损失函数为

将式(4)替换式(3)中的InfoNCE 函数,得出一个优化的PatchNCE 函数,称为Opt-PatchNCE,即

3) SSIM 损失

为了更多的保留水下图像的结构,使用结构一致函数,用于约束原水下图像中的物体结构和增强图像中物体结构一致,其定义为

式中:p为图像块的中心像素;x是X域中 大小为13×13 的图像块;y是G(x)域中大小为13×13 的图像块;ux,uy分别为x,y的均值;σx,σy分别为x,y的方差;σxy为x,y的协变量。SSIM 损失函数主要用于衡量x和G(x)的结构误差,具体为

4) UNDERWATER-CUT 总损失

UNDERWATER-CUT 的总损失函数为对抗生成损失、对比学习损失和SSIM 损失三部分的线性组合,具体为

式中:λ1和λ2分别为约束对比学习损失和SSIM 损失的超参数,选择λ1= 10,λ2=2.5。训练时为了得到最优的生成器,按照下式进行总损失函数的优化

2.2 网络模型训练策略

每一次训练迭代都分为2 个部分,第1 部分,对水下图像域X和脱水图像域Y分别随机采样m个样本,固定生成器参数不变,对判别器的参数进行梯度下降的计算;第2 部分,只对水下图像域X进行随机采样,由于第1 部分已经对判别器参数进行了优化,所以这一部分只优化生成器的模型参数。每次训练判别器和生成器都能根据前一次迭代参数,得以进一步优化,这两部分不断互相指引对方的梯度向更优的方向优化,最终在有限次迭代次数中,二者达到最优,此时生成器增强的图像最为清晰自然。

具体训练配置如下:每次迭代选取8 个样本作为1 批(batch),总的迭代次数(epoch)为200,网络输入为 255×255×3,梯度下降算法为Adam(adaptive moment estimation)算法,其中γ= 0.0002,β1= 0.5,β2=0.999。训练集选取水下图像增强基准(underwater image enhancement benchmark,UIEB)数据集[15],训练时选取800 幅水下图像和脱水图像进行非成对训练。文中算法具体实现使用PyTorch[16]框架,使用NVIDIA GeForce RTX 2080Ti 型号GPU 进行训练,操作系统为Ubuntu 18.04。

3 实验结果和分析

为了验证方法的有效性,将从图像生成效果和客观指标评价两部分来共同评价。图像生成效果上将文中算法和原始CUT、CycleGAN 和SSIMCycleGAN 进行视觉效果上的比较,同时为了证明各个组成部分的有效性,进行了消融实验,比较了SSIM-CUT 和UNDERWATER-CUT 的效果。在客观指标评价上,由于单一指标评测有局限性,例如峰值信噪比(peak signal to noise ratio,PSNR)在某些情况下会出现图像质量很高,但是PSNR 分数不高的情况[17],故采用了FID[18],水下图像质量评价指标(undewater color image quality evaluation,UCIQE)[19]和PSNR 3 种指标来综合评价,使用UIEB 数据集[15]作为训练数据和测试数据,其中800 幅图片为训练集,90 幅图片为测试集,训练集和测试集中的图片分辨率均为256×256。

3.1 图像增强有效性分析

为了客观衡量图像增量后的效果,使用FID,UCIQE 和PSNR 3 种评价指标进行定量比较。FID主要被用来衡量生成模型生成的图像增强结果,FID 的值代表了增强后的图像集和陆上图像集的距离,值越小代表距离越近,同时也代表了脱水后的效果越好。FID 将生成器生成的增强图像和原输入图像看作2 个多维高斯分布,分别用N(µ,σ)和N(µw,σw)表示,然后计算2 个高斯分布的Wasserstein距离,具体为

式中,tr代表矩阵的迹。

UCIQE 指标是色彩浓度、饱和度和对比度的线性组合,用来定量评价水下图像非均匀的色偏、模糊和低对比度的情况,UCIQE 的值和图像质量成正比关系,值越高图像质量越高,具体为

PSNR 反映了两幅图像之间的相似性,值越大表示增强后的图像与陆上图像越接近,即真实性越好,其定义为

式中,MSE为均方误差函数,其具体定义为

式中:I为原始水下图像;K为增强后的水下图像;m和n分别为图像的长度和宽度;MAXI为图像点颜色的最大值,文中图像每个采样点用8 位表示,所以MAXI=255。

各方法增强结果见表1。根据表1 中结果可知,UNDERWATER-CUT 的FID 指标在5 种图像增强算法中最低,说明增强后的图像概率分布和原本水下图像的概率分布最为接近,也就是生成的图像保存信息最为完整;同时,文中方法的UCIQE分数也是最高的,说明在图像的色度、饱和度和清晰度相比其他算法有着更好的表现;此外,文中方法的PSNR 分数接近SSIM-CUT 模型的分数,并都超过了其他算法,综上,文中算法在PSNR 分数接近的情形下,FID 和UCIQE 均领先于其余算法,故可以认为经过UNDERWATER-CUT 增强的图像质量优于其余算法。

表1 各方法增强后图像指标对比Table 1 Comparison of the image indicators enhanced by each method

3.2 图像增强效果对比

为了直观比较各算法的增强效果,将原始图像和各个算法增强后的图像进行对比,结果如图4 所示。原始CycleGAN 算法可以一定程度上增强水下图像,但原始CycleGAN 算法未针对水下环境进行特殊处理,导致水下图像中蓝绿色失衡的问题仍然没有被解决;添加SSIM 的CycleGAN 可以明显提升图像质量,但是由于使用了SSIM 损失函数,只约束了图像中的物体结构,并未对图像的色度和对比度进行处理,图像颜色依然不够真实自然;经过原始CUT 算法处理后,图像中的物体边缘发生了变化,和原始图像相差较大,同时蓝绿色失衡的问题依然存在;对CUT 添加SSIM 约束之后,可以看到增强后的图像和原始图像的物体结构一致,图像较为清晰自然,但是色度和对比度偶有失衡,在结果中可以看到部分增强后的图像背景被错误增强成了红色;最后,UNDERWATER-CUT 效果最为清晰自然,颜色失真得到极大修正,因为使用了SSIM 损失函数又保证了图像中的物体结构不改变,同时使用了相对比较容易优化的损失函数,模型更容易收敛到最优。

图4 不同方法在相同数据集下的增强效果Fig.4 Enhanced results of different methods in the same dataset

4 结束语

文中提出了一种基于UNDERWATER-CUT模型的弱监督水下图像增强算法,不需要水下图像和脱水图像的成对训练,非常适合实际的水下任务。针对水下图像特征,使用SSIM 约束确保修复后的图像不丢失物体结构信息,保证修复后的图像和修复前的结构信息的一致性,同时优化了对比学习损失,使得网络进行梯度下降学习时更容易收敛到最优。实验结果证明,经过UNDERWATERCUT 修复后的图像在校正水下图像颜色效果失真的基础上,最大限度保存了水下物体的结构特征,增强后的图像自然且清晰,有助于其他水下视觉任务如目标识别精度的提升。