基于改进扩散模型的图像去雨方法

钱 枫,胡桂铭,祝 能,邓明星,王 洁,许小伟

(武汉科技大学 汽车与交通工程学院, 武汉 430081)

0 引言

雨天环境下,车载摄像头拍摄到的图像有雨痕干扰,严重影响高级驾驶辅助系统(advanced driving assistance systems,ADAS)识别前方车辆信息的准确性,因此图像去雨对智能驾驶具有极其重要的意义。

目前,关于雨水去除的研究主要有传统方法和深度学习方法。传统去雨方法主要有高斯混合模型[1]、字典学习和稀疏编码[2-3]等,这些方法将雨图像分为雨痕层和背景层,并通过剥离雨痕来得到背景层。然而,这些传统去雨方法在处理复杂图像或大雨等情况时容易存在雨痕残留、纹理信息丢失等问题[4]。

相比之下,基于深度学习的图像处理方法具有更好的泛化性,且可以实现端到端训练。在卷积网络中嵌入残差模块、注意力机制等模块可使网络更好地处理细节信息。Han等[5],Yuan等[6],柳博等[7]利用扩展的深度卷积神经网络学习雨的残差图像,避免了参数估计不准确对结果的影响。随着深度学习框架的不断优化,Wang等[8]通过引入循环网络并利用长短期记忆来连接各个阶段,使实验结果更优;陈峥等[9],张焱等[10]运用门控循环单元设计循环神经网络进行去雨。此外,也有其他组合方式和网络结构应用于图像去雨,如生成对抗网络(generative adversarial network,GAN)在图像去雨领域表现出显著优势,朱德利等[11],Yu等[12],李然等[13]使用GAN对抗训练来提高图像保真度,达到优异的去雨效果。但GAN会存在训练不稳定、模式崩溃等问题,主要原因在于生成数据和真实数据的分布重合度不高[14]。上述基于深度学习的去雨方法,使用注意力机制、改进的残差块等对雨痕进行去除,然而,这些方法大多未考虑雨痕类型(如雨滴、雨线)对图像的影响,有较多的雨痕残留,去雨质量仍有较大的提升空间。

为了解决上述图像去雨质量差的问题,本文中提出一种基于扩散模型的图像去雨方法,通过设计残差模块双输入信息通道和添加ECA通道注意力机制对网络进行改进,提高了模型捕获跨通道交互信息的能力,进而有效去除雨痕,提高了去雨质量,泛化性更好。

1 扩散模型

扩散模型作为生成模型之一,是一种依赖先验的条件模型,能够实现从噪声(采样自简单的分布)生成目标数据样本。模型[15]包括前向和逆向2个过程,通过在数据中加入随机噪声(前向过程)并从中采样出所需的数据样本(逆向过程)来实现,扩散过程如图1所示。

图1 扩散过程示意图

1.1 前向过程

前向过程是指逐渐向数据添加高斯噪声直至数据趋于高斯分布。在数据中包含着T个时间(T在训练过程中为可变参数)的扩散过程,当前时间是通过上一时间的数据按照式(1)所示添加高斯噪声。

(1)

式中:q(xt|xt-1)为先验分布;N为高斯分布;xt为时间t对应的加噪数据;I为高斯噪声的方差;βt在扩散过程中从0.0 001到0.02线性增大,并且加入高斯噪声的均值和方差由βt决定。

前向过程采用预先设定好的方差调整方案,使得整个扩散是一个固定的过程。当时间T足够大时,得到的xt则变为完全的随机噪声。

1.2 逆向过程

扩散模型的逆向过程为去噪过程,通过噪声估计网络推测真实分布q(xt|xt-1),从随机噪声开始逐渐去噪,最终生成一个真实的样本。公式如式(2)所示。

p(xt-1|xt)=N[xt-1;μθ(xt,t),σ2I]

(2)

式中:p(xt-1|xt)为后验分布;σ2为方差;μθ(xt,t)为参数化估计均值,通过噪声估计网络对噪声εθ(xt,t)预测进而间接得到。

网络训练公式如式(3)所示,从时间t-1到0逐步重复xt还原成xt-1的步骤(即计算xt-1,将得到的xt-1作为新的xt)。

(3)

2 基于扩散模型的去雨结构

2.1 去雨模型结构

本文中运用扩散模型对图像进行去雨处理,并改进模型的噪声估计网络。设计残差模块的双输入信息通道,生成更多的训练数据,提高模型的泛化性;添加ECA通道注意力机制模块,该模块明确给出每个通道的权重,从而帮助理解模型的决策过程,提高模型的可解释性。模型去雨结构如图2所示,有雨图至加噪数据图(xt)过程为经过T时间加噪直至变为纯噪声的前向过程;加噪数据图至无雨图过程为将噪声输入噪声估计网络学习去噪从而得到无雨图的逆向过程。

图2 扩散模型去雨结构图

模型的训练过程:从有雨图数据集内随机选择一个批次的样本进行训练,每个批次包含8个样本,记为x0,并在1到T的时间中随机选取一个时间t,将其对应的高斯噪声与x0结合形成加噪后的数据结果xt。时间t能够生成的正弦位置编码,与xt一起送入噪声估计网络,准确识别当前所处在的去噪阶段,预测出高斯噪声εθ(xt,t)。通过高斯扩散计算该噪声与随机噪声的损失,更新噪声估计网络权重,并继续训练直至训练完成。

2.2 噪声估计网络

为了最小化实际的噪声分布与预测的噪声分布之间的差异,本文中采用噪声估计网络结构来预测噪声并不断更新权重。

噪声估计网络结构包含了下采样模块、中间层模块、上采样模块、组归一化(group normalization,GN)、激活函数(Swish)等。网络模型结构如图3所示,为了提高模型的泛化性,分别在下采样模块、中间层模块、上采样模块中添加双信息输入的残差模块(residual module,block)中;在中间层模块中添加ECA注意力机制。网络模型是一种编码-解码(encoder-decoder)架构。编码器的每个阶段都包含下采样模块,用来降低特征的空间大小,然后通过解码器将压缩的特征逐渐恢复。为了更好地优化网络,解码器中使用跳跃连接,将编码器中得到的同维度特征连接起来。

图3 改进的噪声估计网络结构示意图

网络预测噪声过程,将有雨图像训练集随机裁剪成64×64的三通道图片,并将其与时间t对应的正弦位置编码一并送入网络中进行训练,计算t时刻后模型的噪声图片,可以预测出噪声,并用预测噪声与实际噪声进行拟合,不断更新权重。有雨图像经过卷积、上采样、中间层和下采样等一系列特征处理,最后通过卷积操作输出得到无雨图。

2.2.1 双信息输入残差模块

针对网络加深而恒等映射难以拟合的问题,引入残差模块作为网络结构的基础组成部分。本文中将残差模块设计为双信息输入,其中一个信息从训练集中抽取雨图作为样本,经过组归一化、激活函数和卷积层处理;另一个信息为从1到T时间随机抽取的t时刻正弦位置编码,经过激活函数和全连接层处理;而后将2个信息得到的输出进行相加操作,从而实现网络信息融合输入,其网络结构如图4所示。该模块中的组归一化对各个组通道方向上的均值和方差进行归一化,从而脱离batchsize大小的约束。在残差模块中加入1×1卷积(convolution,conv)调整因卷积层而改变的通道数,这种方法保留了编码过程中不同层丢失的信息。

图4 残差网络结构示意图

模型采用残差结构,使得特征映射对输出的变化更为敏感,模型的梯度更容易训练。为了保持网络输入输出相同,需使h(x)=x,故令F(x)=0,公式如式(4)所示。

h(x)=F(x)+x

(4)

式中:h(x)为输出;F(x)为残差项;x为输入。网络首先从输入中提取出特征,再与残差结构组成新的特征。在残差模块中,采用跳跃映射的方式来构建残差单元,即通过将输入与输出直接相加的方法,来补充在卷积过程中可能丢失的特征信息。残差模块可以减少数据维度,同时也降低梯度消失与梯度爆炸的风险。

2.2.2 ECA通道注意力机制模块

注意力模块可以定位图像中的目标区域,有效提取图像特征,使网络能更好地理解雨痕区域的特征,并且在处理时保留了图像的细节信息。为了提升模型性能,在网络的中间层模块加入ECA注意力机制。该机制具有较强的信息获取能力,且不会受数据量的影响,能够在不降维的情况下实现全局平均池化,增强整体的通道特征,提升模型性能。ECA通道注意力模块如图5所示,利用1×1卷积来获取通道间的信息,避免了通道的维数下降和参数量减少的问题。

图5 ECA模块结构图

ECA模块首先对输入的特征图进行空间特征的压缩,将其维度从H×W×C变为1×1×C,这一过程采用全局平均池化操作,以获得大小为1×1×C的特征图。然后,模块使用1×1卷积对该特征图进行通道特征学习,了解各通道间的重要性。整个过程,输出的维度不变,仍然为1×1×C。最后,通过通道注意力机制,将通道注意力的特征图1×1×C与原始输入特征图H×W×C逐通道相乘,从而输出具有通道注意力的特征图。

3 实验结果及分析

3.1 数据集设置

本文中对2类雨天场景图像进行实验,一类是相机镜头上有雨滴干扰,另一类是图像背景有雨线遮挡。

对于第一类的雨滴情况,采用文献[16]提供的模拟雨滴数据集,共包含1 119对图片,其中训练集包含861对,测试集包含258对。对于第二类的雨线情况,使用Rain100数据集进行实验,共包含2 000对图片,其中训练集包含1 800对,测试集包含200对。

3.2 训练

3.2.1 训练条件

本实验的训练使用了一台具有24 GB内存的RTX3 090计算机,Python版本为3.9.13,使用的深度学习框架为pytorch.1.13.1。

3.2.2 训练参数设置

本工作对GMZ07膨润土及其不同掺砂率混合物的饱和试样进行直剪试验及扫描电镜测试,研究试样用不同浓度盐溶液饱和时的强度特性,分析和总结试验结果可得到以下结论.

在本文中,网络参数设置初始学习率为 0.000 02,batchsize为8,总迭代次数设置为200 000,并且使用Adam优化器进行梯度优化。为了缓解过拟合现象,模型中引入了Dropout层,其作用是在每个训练批次中随机忽略50%的特征检测器,即将一半的隐层节点值设为0,以此来降低模型的复杂度,并且可以脱离对一部分局部特征的依赖。

3.2.3 训练过程

首先,运用本文提出的网络分别对上述2种数据集进行迭代训练,得到了包含雨痕信息的预训练模型,然后,运用该模型对测试集进行测试。在训练过程中,使用损失函数来估计雨图中雨痕噪声,经过约152 000次迭代,Loss曲线的斜率趋于0,并且数值接近于0。此时,可以认为模型训练已经收敛,整个网络达到了很好的训练效果。

3.3 算法对比及实验结果

本文中使用峰值信噪比(peak signal to noise ratio,PSNR)和结构相似性(structural similarity,SSIM)2项指标对去雨效果进行评价。

为了验证本文算法的可信度,运用改进扩散模型与其他几种去雨算法(GMM、DID-MDN、GAN、文献[16]方法)对测试集图片进行测试,测试得到的图片平均PSNR和SSIM值如表1所示。使用不同方法去雨得到的部分效果图如图6所示。

表1 去雨方法在测试集上的平均PSNR与SSIM值

图6 5种方法去雨结果图

通过比较去雨操作后输出图像和无雨图像之间的关系,得到PSNR和SSIM值。PSNR的计算公式如式(5)所示,其数值越高代表去雨图的失真越小,则效果越好。

(5)

式中:MAXI为无雨图中最大的像素值;MSE为去雨图和无雨图之间的均方差。

(6)

式中:x为雨图;y为去雨图;μx、μy分别为x、y的像素均值;σx、σy分别为x、y的像素方差;σxy为x和y的协方差;c1和c2为常数。

通过表1与图6可知,在模拟雨滴数据集上,GMM和DID-MDN算法的去雨效果并不明显;而GAN网络去除了部分雨滴,但仍有雨滴残留;另外,文献[16]所提算法在去除雨滴效果较好,但损失了部分车辆的边缘信息。在Rain100雨线数据集上,GMM方法去雨有部分雨痕残留;DID-MDN仅能去除少量雨线;GAN网络与文献[16]提出的方法具有良好的去雨效果,但图像的关键细节丢失严重。而改进扩散模型在去除雨滴和雨线方面均取得最高的PSNR和SSIM,其中雨滴和雨线的PSNR值分别为30.328 5和34.896 5,SSIM值分别为0.927 1和0.962 0,表现出良好的性能。通过图6可知,本文算法能够保留图像的细节信息,在一定程度上减少了图像的去雨过度,泛化性更好。

3.4 真实雨天情况下车辆检测验证

本实验选择了2个测试集与真实雨图数据集中包含车辆的图片,使用目标检测算法进行车辆检测。

为了验证本文扩散模型去雨算法的可信度,本文中使用目标检测算法对去雨后的真实雨滴、雨线图像与上述雨滴、雨线测试集图像进行车辆检测,以目标检测置信度作为评价指标。在当前的目标检测器中,YOLOv7在速度和精度上都超过了以往的检测器,故本文中使用YOLOv7进行车辆检测。直接对4个数据集中的原始有雨图进行车辆检测,得到以下置信度结果:1)模拟雨滴[16]为0.49;2)Rain100为0.77;3)真实雨滴图为0.83(白车),0.56(红车);4)真实雨线图为0.75(白车),0.75(黄车)。将上述5种去雨算法生成的无雨图输入目标检测算法中进行检测,不同去雨方法的检测效果如图7所示,计算得出的平均置信度提升值如表2所示。

表2 车辆检测置信度提升平均值 %

图7 5种方法去雨后检测结果图

由表2与图7可知,其他4种算法在真实雨天环境下去雨后图像背景丢失严重,对数据集去雨图片的车辆检测置信度提升有限;而本文算法在雨天情况下对车辆检测置信度的提升均高于其他4种算法。该算法能够通过去除雨痕来有效地提高在雨天环境下的图像质量,从而提高后续计算机视觉任务的准确性,进一步验证了改进扩散模型的实用性。

4 结论

针对图像去雨任务,提出一种基于改进扩散模型图像去雨方法。将扩散模型应用在图像去雨领域,通过雨图和时间来引导前向扩散加噪,从而逆向采样得到无雨图;并且设计残差模块的双输入数据信息和添加ECA注意力机制,提升去雨效果。

与GMM、DID-MDN、GAN和文献[16]算法相比,在模拟雨滴和Rain100数据集去雨上,本文算法取得了最高的PSNR和SSIM值。此外,对去雨后的真实雨滴、雨线图像与上述雨滴、雨线测试集图像通过YOLOv7进行车辆检测,本文所提算法对置信度的提升高于其他几种算法,进一步验证了本文所提算法在图像去雨效果上优于其他方法,同时最大程度保留图像细节部分。