Hierarchical attention aggregation with multi-resolution feature learning for GAN-based underwater image enhancement
引言
近年来,为了优化海上作业效率并推动海事行业自动机器学习的发展,水下图像增强与复原技术变得日益重要。针对水下图像中存在的色彩偏差、曝光不足和模糊等问题,本文提出了一种基于生成对抗网络(GAN)的水下图像增强方法,该方法采用了新的分层注意力聚合和多分辨率特征学习技术。所提方法包括一个生成器和一个判别器。具体而言,生成器由编码器、瓶颈层和解码器组成。生成器引入了块间串行连接,以更好地适应复杂的图像场景和任务需求,同时采用并行连接来提取多级特征并增强网络的表达能力。
为了提取语义和上下文信息,在编码器中设计了分层注意力密集聚合,其中包括多尺度特征层次和密集特征层次。此外,在瓶颈层中设计了多尺度空间注意力机制,以应对水下图像场景的变化。在解码器中,重点考虑了特征通道层,并提出了一种多通道注意力机制,以将多分辨率通道特征恢复为三通道增强图像。
此外,还引入了角度损失函数作为额外的监督手段,以提高生成图像与原始图像之间的相似性,增加图像清晰度,并减少色彩偏差。同时,本文采用了补丁判别器来增强机器视觉。大量实验表明,所提方法优于目前最先进的方法。
介绍
随着人类对海洋的探索和资源开发的不断深入,海事行业已成为一个热门领域。然而,由于海洋的广阔和环境的复杂性,海上交通流预测(Cohen and Klein, 2022)和自主作业(Gao et al., 2023)仍然是海事行业学者关注的焦点。在海洋活动中,光学成像系统常被用于记录水下元素。海事工作的效率主要取决于所捕获图像的质量;然而,退化的水下场景可能会阻碍高级视觉任务的性能和信心。
使用光学成像技术拍摄的水下照片常常因为光的散射和衰减(Zhou et al., 2022c)而具有低对比度、能见度差、特征模糊和颜色失真等问题。根据光学成像原理,光散射是由于光线在到达相机之前先照射到水中的颗粒上,从而降低了图像的能见度和对比度(Zhou et al., 2022a)。此外,不同波长的光在水中传播时会经历不同程度的衰减,其中波长较长的红光首先衰减,而波长较短的绿光和蓝光衰减较慢,这意味着水下图像以蓝绿色调为主(Zhuang and Ding, 2020; Jiang et al., 2022a)。
因此,水下图像增强作为计算机视觉领域的一项基本视觉任务,越来越受到人们的关注。
计算机视觉技术不断发展进步,吸引了越来越多的关注。为了解决水下环境中的图像退化问题,计算机视觉技术已被应用于水下图像增强领域,逐渐成为提高水下图像质量的主要技术。先前的工作提出,可以通过直方图变换、灰度边缘检测和白平衡等技术直接调整图像的像素值来改善视觉效果。然而,简单地调整像素值可能会导致噪声放大和颜色失真(Zhou et al., 2022b)。许多学者使用小波变换去除噪声,取得了显著效果。随着水下图像增强和复原方法成为热门研究话题,研究人员开始考虑水下图像的独特特性。虽然已提出传统方法,但这些方法往往忽略了水下成像的光学特性,导致颜色偏差和伪影。许多学者开始模拟水下图像的退化过程,估计模型参数并反转退化过程以获得清晰图像。这种基于物理模型的方法依赖于物理先验,具有局限性。最近,计算机视觉任务中提出了深度学习技术,并显示出强大的学习能力。
为了减轻水下图像中的颜色失真和低能见度问题,已经提出了许多用于水下图像增强和恢复的深度学习方法。现有的基于深度学习的图像增强方法可以分为三类:基于卷积神经网络(CNN)的方法、基于生成对抗网络(GAN)的方法和基于Transformer的方法。其中,基于CNN的方法中最常用的结构是基于Unet的编码器-解码器架构(Zhou et al., 2023b,c; Liu et al., 2022),它可以融合来自多个分辨率的特征信息来进行图像增强。同时,用于水下图像增强的基于GAN的方法也展示了有效的结构,可以促使恢复后的图像特征更接近真实图像。
在计算机视觉领域,Transformer在各种任务中取得了巨大成功(Jiang et al., 2022b; Xiao et al., 2023; Wang et al., 2023)。在水下图像增强中,U-Shape是首个引入Transformer的方法,它结合了通道Transformer和空间全局特征Transformer。然而,与基于GAN的方法相比,基于Transformer的方法存在一些缺点。首先,基于GAN的方法可以直接从真实样本中学习更真实的图像,而基于Transformer的方法通常通过自回归模型生成图像,这可能导致图像模糊或失真。此外,基于Transformer的方法的生成过程受到输入序列的限制,不完整或错误的输入信息可能导致图像质量降低。其次,基于Transformer的方法需要大量的参数进行训练,并且每个时间步都需要计算复杂的自注意力机制,这导致训练时间长和计算资源成本高。另一方面,基于GAN的方法可以通过对抗训练快速学习一个优秀的生成模型,更适合处理水下图像增强中的复杂场景和各种类型的图像退化。
因此,考虑到水下图像的特点以及训练时间和计算资源的限制,我们选择使用基于GAN的架构进行水下图像增强。我们提出了基于GAN的水下图像增强方法——具有多分辨率特征学习的分层注意力聚合(HAAM-GAN),该方法结合了Unet架构和对抗学习。与传统的Unet结构不同,我们区分了编码器和解码器的主要功能,并在它们之间引入了瓶颈层。我们的方法利用具有不同焦点的注意力机制来引导编码器、瓶颈层和解码器关注不同的图像特征,从而促进了我们图像增强网络的不同功能。我们提出的编码器主要关注图像的语义和上下文信息。鉴于水下图像增强的最终目标是促进后续的高级视觉任务,如海事行业中的自动机器学习,我们引入了补丁级判别器来鼓励网络关注像素级特征信息。与传统损失函数通常优先考虑与人类更相关的信息相比,补丁级判别器使得HAAM-GAN网络在机器视觉方面更接近真实图像。
与以往的工作不同,HAAM-GAN考虑了模型的实时性能和表达能力,并选择了一种并行和串行连接相结合的方式,包括块间串行连接和块内并行连接。在水下图像增强任务中,使用块间串行连接可以更好地适应复杂的图像场景和任务需求,从而提高模型的性能和有效性。通过模块间的串行连接,可以逐层提取高级特征,使水下图像特征的学习和处理更加有效,并增强模型的鲁棒性和泛化能力。此外,为了提取多级特征并实现它们的交互,我们在编码器(HADA)和瓶颈层(MSA)中使用了并行连接,其中不同的特征提取器独立地学习特征,从而增强了网络的表达能力。
具体来说,为了生成能够有效恢复水下图像的潜在向量,HAAM-GAN在编码器中更加关注输入水下图像的上下文语义信息。因此,我们提出了分层注意力密集聚合(HADA)来聚合多尺度特征层次和密集特征层次。考虑到多个下采样特征具有更大的感受野,我们在最深层设计了基于多尺度空间注意力机制(MSA)的瓶颈层,以进一步改进潜在向量的几何结构和语义信息。
在解码器中,主要目标是从获得的潜在向量中恢复出清晰的图像。由于解码器和瓶颈层为潜在向量生成了丰富的结构信息,HAAM-GAN在解码器中设计了一个多通道注意力机制(MCA)来研究通道间的权重信息(将多分辨率通道特征恢复到三通道增强图像)。
在损失函数中,考虑到机器视觉在后续高级视觉任务中的重要性,我们基于GAN的对抗理论引入了机器视觉判断。我们还引入了感知损失来保留图像特征,并结合角度损失来限制RGB颜色空间中的像素,避免颜色不匹配,有效克服过曝光问题,并减轻颜色失真。
我们的主要贡献总结如下:
(1)我们在编码器中提出了分层注意力密集聚合(HADA),通过两尺度卷积(多尺度特征层次)获得不同感受野下的语义信息和上下文特征信息,然后设计了一个密集连接层来获得密集特征层次,这增强了编码器的表达能力,并获得了更多可行的潜在向量。
(2)为了进一步获取图像的全局几何结构信息,我们在最深层(具有更大的感受野)设计了瓶颈层,并提出了多尺度空间注意力机制(MSA),以进一步改进潜在向量的几何结构和语义信息。
(3)为了有效地将具有丰富结构信息和上下文信息的潜在向量恢复为三通道清晰图像,我们在解码器中设计了多通道注意力机制(MCA)来进行通道加权,并串联两种不同的注意力机制来生成更接近真实情况的图像。
(4)我们结合了GAN的对抗理论来增强机器视觉并促进后续任务,并考虑到水下图像中严重的色彩偏差,我们提出了角度损失函数作为更有效的监督,这有效地使所提出的网络学习到的映射更加清晰。
本文的其余部分组织如下:第2节讨论了传统方法和基于深度学习的方法的优缺点。第3节介绍了所提出的方法。第4节进行了定性和定量评估,以将所提出的方法与最先进的方法进行比较。最后,我们在第5节对所提出的方法进行了总结。
相关工作
由于水下探测的重要性,已经提出了许多水下图像增强与复原方法。目前的方法可以分为两类:传统水下图像增强与复原方法和深度学习驱动的水下图像增强与复原方法。在本节中,我们将详细讨论上述两种方法。
2.1 传统方法
传统方法可以分为图像增强方法和图像复原方法。图像增强方法通过直接修改图像像素值来改善视觉图像质量。
基于物理模型的图像复原方法使用诸如透射率和背景光等关键参数,利用水下成像模型估计清晰图像。它通过对水下成像模型进行反演来获得增强后的图像。图像复原方法主要关注透射图估计的准确性。
2.1.1 图像增强方法
由于颜色衰减、对比度低和曝光不足等因素,提高水下图像的视觉质量面临着重大挑战。为了应对水下图像中的这些问题,许多学者提出了各种增强方法。一个常见的范式是先利用白平衡解决颜色偏差问题,然后使用诸如直方图均衡化或基于Retinex理论的技术来增强颜色校正后的图像的对比度和细节信息。
Ancuti等人(2012)提出了一种基于融合的水下图像增强方法,该方法结合了颜色校正和对比度增强的图像。
Zhou等人(2022c)提出了一种基于多特征先验融合(MFPF)的图像增强方法,该方法通过提取和融合水下图像的多个特征先验来改善其视觉质量。基于视网膜理论,Zhuang等人(2021)提出了BRE方法,以解决由于水下环境光照变化引起的曝光不足和模糊问题。Bai等人(2020)提出了四个阶段来增强水下图像:像素强度中心区域化、直方图全局均衡化、直方图局部均衡化和多尺度融合。由于水下图像固有的颜色衰减,颜色校正和增强方法已被同时或顺序地应用于各种水下图像复原算法中。Fu等人(2014)首先提出了一种简单的水下颜色偏差校正算法,然后提出了一个基于视网膜皮层理论的直接光和反射光分离的可变框架,最后通过各种增强策略改进了分离出的光分量。
无模型方法未能考虑水下环境的独特成像特性,特别是物体与摄像机之间深度上的强烈依赖性。因此,这些方法对水下图像应用统一的增强操作,这阻碍了其解决不同深度下降质问题的能力。因此,无物理模型的方法会导致不恰当的(过度/不足)增强结果(Li等人,2021)。
2.1.2 图像复原方法
水下图像因传输和大气光的影响而发生退化,且退化的程度受图像场景深度的影响。为了恢复不同程度的退化,研究人员利用水下成像模型取得了进展。通过反转成像模型,他们成功地在改进后的照片中实现了增强,使照片看起来更加自然。为了进一步解决退化问题,已经提出了多种技术。
Zhou等人(2021a)设计了一个稳健的成像模型来估计背景光,该模型通过补偿红色通道来修正颜色校正图像的传输图。Sea-thru(Akkaynak和Treibitz,2019)提出了两个新观点:一是信号衰减系数并非恒定不变,二是与后向散射相关的系数与信号衰减系数不同。基于这两个观点,Sea-thru开发了一个新的物理成像模型,以提高水下图像颜色失真的鲁棒性。Lee等人(2020)提出将连续颜色校正算法与暗通道先验方法(SDCP)相结合,以最小化增强图像中的红色伪影。Chang等人(2018)改进了SDCP,并开发了一种基于深度估计和传输补偿的水下图像恢复方法。
传统方法在水下图像增强领域取得了良好效果。其中,不依赖于物理模型的方法不考虑水下图像的成像特性(水下图像增强任务强烈依赖于物体与相机之间的距离),而是对水下图像进行统一的增强操作,导致无法解决不同深度下的图像退化问题。物理模型方法基于一个或几个固定的先验来估计模型中的参数,然后反转得到真实的水下图像,但水下环境多样,受到悬浮颗粒和水质的影响。现有的先验无法满足所有水下图像,导致模型中的参数估计错误,进而使水下图像的恢复出错。
基于深度学习的方法
由于深度学习网络具有强大的学习能力,近年来越来越多的人将其应用于水下图像改进,并取得了令人瞩目的成果。以下介绍几种水下图像增强和复原的深度学习方法。
基于卷积神经网络(CNN)的方法
在确保网络结构不会过于庞大以充分利用每层卷积提取的特征的前提下,Li等人(2020a)提出了基于先验水下场景的轻量级网络UWCNN。UWCNN采用模块化结构构建卷积网络,并在每个模块之间使用密集连接。值得注意的是,UWCNN简化了卷积网络,该网络仅包含10个卷积层,兼顾了网络学习的时效性。网络模型的大小会对其本身的学习能力产生限制。Liu等人(2019)提出了一种基于残差学习的方法,并针对数据集不足的问题,使用循环生成对抗网络(CycleGAN)生成水下合成图像。这显著改善了水下图像的视觉效果,并有助于基于视觉的水下任务分割和跟踪。随后,Li等人(2021)提出了由介质传输引导的多色彩空间编码的UColor网络。Li等人(2019b)发现,如果输入图像在进入卷积网络之前进行预处理,那么生成的图像质量将会提高,因此他们提出了WaterNet网络。
在解决低级视觉任务方面,深度学习技术取得了显著进展。然而,为CNN训练创建逼真的水下图像生成模型具有挑战性,因为它依赖于各种因素,如场景和照明条件、温度和浊度。
基于生成对抗网络(GAN)的方法
生成对抗网络(GAN)利用损失函数指导生成器和判别器之间的对抗,从而使生成器获得最佳的生成效果。生成器的目的是生成尽可能接近实际图像的图像,以使判别器无法区分生成的图像。
判别器的目的是尽可能准确地区分生成的图像和真实图像。如果判别器无法正确区分生成的图像,则需要进一步学习。
不难从上文得出结论,GAN的训练过程需要大量的训练数据集(Goodfellow等人,2020)。
Fabbri等人(2018)提出了UGAN网络。该模型使用循环一致性对抗网络(CycleGAN)(Zhu等人,2017)生成数据集样本。此外,还使用生成对抗网络(GAN)替代了单一的卷积神经网络(CNN),以提高水下视觉场景的质量。随后,Islam等人(2020b)基于条件生成对抗网络(cGAN)(Vishwakarma等人,2020)提出了FUnIE-GAN网络。前者采用了一套成对和不成对的水下图像进行训练,这些图像是在不同能见度条件下使用七种不同的成像工具拍摄的。该网络是一个轻量级的实时水下图像增强模型,能够迅速响应。
论文方法
我们在图1中展示了HAAM-GAN的总体架构,它主要由生成器和判别器组成。生成器网络以原始水下图像为输入,通过学习多分辨率特征来生成清晰的水下图像。在编码器部分,我们提出了层次注意力密集聚合(HADA)方法来提取有利于水下图像增强的潜在向量,重点关注提取语义和上下文信息,以便在解码器中进行图像重建。为了进一步确定潜在向量的特征信息以进行细节恢复,我们在瓶颈层设计了多尺度空间注意力机制,以帮助模型更好地理解图像的几何结构和语义信息。在提取出图像的结构语义信息后,我们在解码器中设计了多通道注意力机制,以关注通道之间的权重信息(将多分辨率通道特征恢复为三通道增强图像)。考虑到传统损失函数是基于人工设置的(更多关注人类视觉感知),并且不能自适应地区分增强图像和真实图像,我们从机器视觉的角度引入了一个判别器来关注图像的特征信息,这使得我们的方法更有利于后续的海洋工作。
简而言之,HAAM-GAN通过结合层次注意力密集聚合、多尺度空间注意力机制和多通道注意力机制,以及一个专注于图像特征信息的判别器,实现了对水下图像的清晰化增强,为后续的海洋工作提供了更好的图像基础。
层次注意力密集聚合
在编码器部分,我们在每次下采样步骤之间采用了层次注意力密集聚合模块。如图1(a)所示,由于水下图像的特性,不同水下图像的场景信息具有不同的尺寸,因此需要不同的感受野来提取场景信息。基于这一特性,HADA首先使用3×3和5×5的卷积来获取不同尺度的上下文特征信息和场景结构信息,这被称为多尺度特征层次。在获得不同尺度的场景信息后,为了增强编码器的鲁棒性,HADA引入了密集连接,生成包含丰富上下文信息的密集特征层次,用于图像恢复,实现了语义和结构特征信息的重用。与ResNet相比,这种方法提高了效率。
多尺度空间注意力机制
多次下采样后的特征具有更大的感受野,这使得更容易获取图像的全局信息。为了进一步增强潜在向量中整体图像场景的几何和语义信息,我们提出了一个作为瓶颈层的多尺度空间注意力机制,如图1(b)所示。在MSA中,我们使用多尺度(3×3、5×5、7×7)卷积来获取水下图像的整体场景结构信息和图像上下文信息,这可以覆盖最低分辨率特征层中不同比例的场景信息(如鱼、潜水员、珊瑚礁等)。
在获取图像的整体场景信息后,我们引入了一个通道注意力机制(Sandler等人,2018)来关注通道之间的不同特征信息,这增强了瓶颈层的表达能力,并提高了HAAM-GAN的泛化能力。考虑到深层特征的通道维度相对较高,通道注意力机制可以过滤掉一些不重要的特征,有效降低了模型的计算复杂度,提高了HAAM-GAN的效率。
多通道注意力机制
在获得包含丰富结构信息的多分辨率潜在向量后,HAAM-GAN的解码器旨在利用这些潜在向量,并将多分辨率通道特征恢复为三通道增强图像。因此,HAAM-GAN的解码器更关注通道间的信息(在解码器中,每次上采样都伴随着特征通道数减半)。为此,我们提出了一个多通道注意力机制,它通过级联两个不同的通道注意力机制来过滤掉多通道特征信息中的一些不重要特征信息,同时保留有利于图像恢复的通道特征信息。
为了在确保网络正确表达的同时降低HAAM-GAN的时间复杂度,MCA级联了一个简单的通道注意力模块(SCAM)和一个通道注意力模块(CAM)。在SCAM中,我们仅使用平均池化操作来保留特征图的详细信息,并避免过度简化导致的信息丢失。此外,平均池化还可以平滑特征图,减少噪声,并提高模型的稳定性。平均池化具有一定的不变性,即当特征图中的对象位置发生变化时,池化后的结果不会变化太大,从而提高了模型的鲁棒性。在SCAM之后,我们提出了通道注意力模块(CAM),通过引入最大池化操作来提取同一通道内最显著的特征,同时保留平均池化的平滑效果和特征不变性。通过结合这两种池化操作的结果,模型可以保持鲁棒性和准确性。最后,解码器输出与输入图像大小相同的增强图像。
Patch-GAN 网络
相较于水下图像中的低频信息,高级视觉任务更关注图像中的高频信息。因此,HAAM-GAN的判别器被设计为引入Patch-GAN来建模高频结构。Patch-GAN(Isola等人,2017)与其他常用的判别器不同之处在于,它将输入转换为n×n的矩阵X,并在块级别上区分生成的图像和真实图像,通过使用局部像素块来专注于建模高频输入。
此外,与在图像级别进行全局识别的判别器相比,Patch-GAN判别器的结构在计算上更为高效,因为它需要的参数更少。该判别器使用四个卷积层将256×256×6的真实图像输入转换为16×16×1的图像输出,这反映了判别器的平均有效性响应。每个卷积层都采用了4×4的滤波器,步长为2,填充为1。输出随后经过Leaky-ReLU非线性激活和批量归一化(BN)处理。
损失函数
为了监督对抗训练和条件生成对抗网络的损失函数,我们在目标函数中加入了其他必要的损失函数,包括全局相似度函数、内容感知损失函数和角度损失函数。
条件生成对抗模型的损失函数
我们从博弈论的角度(Vishwakarma等人,2020)在GAN网络的基础上加入了生成条件,损失函数可以表示为:
其中,𝑍代表输入的水下图像,𝑋表示真实图像(即ground truth)。𝐸表示期望。在传统的条件对抗生成网络中,生成器𝐺试图减小损失函数𝐿𝑐𝐺𝐴𝑁(𝐺, 𝐷),而判别器𝐷则试图最大化它。此外,我们还在所提出的模型中包含了全局相似度函数、内容损失函数和角度损失函数,以衡量生成图像的质量。
全局相似度损失函数
众所周知,L1和L2范数损失函数在图像生成方面会存在模糊问题。L1损失函数,也称为最小绝对偏差函数,比L2损失函数具有更高的鲁棒性和更少的模糊性。因此,我们引入L1损失函数,通过测量预测结果与ground truth之间的像素级值差异,鼓励生成器𝐺生成与实际图像相似度更高的图像。损失函数如下:
其中,𝑌是与𝑋对应的ground truth。
内容感知损失函数
我们在目标函数中增加了一个内容损失项,以鼓励𝐺生成具有与目标(实际)图像相似内容(即特征表示)的改进图像。受Li等人(2019a)的启发,图像内容损失函数𝛷(∙)定义为从预训练的VGG19网络的block5层检索到的高级特征。内容损失函数主要通过比较从图像中提取的特征与已经训练过的VGG19网络中的特征来计算。我们提出的内容损失函数如下:
请注意,这里的𝛷(∙)表示从VGG19网络中提取的特征映射函数,它接受一个图像作为输入并输出该图像的高级特征表示。
角度损失函数
与陆地上低质量图像(如雾天或雨天)的增强不同,水下图像由于光在水下传播过程中的选择性吸收而遭受严重的蓝绿色失真。此外,由于水介质吸收光线,水下拍摄时通常使用带有照明效果的成像仪器来增强水下图像的可见性。然而,这些措施可能导致环境光和相机补光灯的干扰,从而在水下图像中产生亮点。因此,我们引入了角度损失函数来约束生成器,其目的是衡量生成器的结果图像与真实图像之间的差异。与在空间域中使用基于欧几里得距离的损失相比,角度损失函数(Sidorov,2020)对微小的角度差异更为敏感,能够更好地在空间域中区分相似的样本数据,并指导神经网络的后续优化方法。
在水下图像增强网络中,使用角度损失函数可以使网络更好地学习复杂特征,从而提高增强效果。此外,由于水下图像通常受到颜色偏移和散射等多种因素的影响,使用角度损失函数可以帮助网络更好地区分不同类别的样本,并更准确地恢复图像的细节和结构。
从公式(1)我们可以推断出,损失可以衡量真实图像和生成图像之间的像素级值差异。由于光在水下传播过程中会被吸收和散射,我们在水下拍摄时会使用带有照明效果的成像仪器,以使水下图像具有更高的可见性。然而,上述措施会导致水下图像受到场景光和相机补光灯之间相互作用的影响,从而产生颜色偏差。
为了解决上述问题,我们没有考虑同时拟合方向和长度的损失函数。相反,我们将角度损失函数添加到目标函数中,并且仅使用一个向量(方向)和一个增加灵活性的参数,以改善图像的颜色恒常性。角度损失函数的公式如下:
其中,𝑎𝑛𝑔𝑢𝑙𝑎𝑟是角度误差的度量,用于比较预测值与准确真实图像之间的差异。
总损失函数
综上所述,我们将这些损失函数与模型训练相结合,形成了一个目标函数。条件生成对抗网络的损失函数指导模型进行训练。𝐿1损失函数采用全局相似性,并添加了内容感知函数,以鼓励生成器生成与实际图像内容相似的图像。角度损失函数考虑了预测值与真实值之间角度的预测值,并引入了尺度不变性,以提高目标对特征多样性的鲁棒性。
模型训练的总损失函数如下:
总的来说,将角度损失函数添加到目标函数中,提高了生成图像的清晰度,以及生成图像与原始图像之间的一致性,并消除了颜色偏差。