一种基于Transformer的图像颜色校正方法

一种基于transformer的图像颜色校正方法
技术领域
1.本发明涉及图像处理的技术领域，特别是一种基于transformer的图像颜色校正方法。


背景技术：

2.随着智慧城市建设不断深入，对自动驾驶、智慧交通、智能安防的需求越来越强，对图像的质量要求也越来越高。在复杂光照下，如夜间受到车灯、路灯等因素影响，使获得的图像带有不同光照的投影，会造成物体色度信息的改变，对于同一物体表面，不同的光照可能会产生完全不同的颜色。颜色特征是人类感知外界物体的重要特征，但颜色特征又是一种极不稳定的视觉特征，是影响图像识别与检索性能的重要因素之一。颜色矫正是计算机视觉领域中既基本又重要的任务之一，其处理结果的好坏决定着后续任务的性能，是计算机视觉领域重要预处理阶段。近些年，深度学习理论也被用于解决颜色计算问题，主要思路是首先利用网络模型对场景光源色度进行估计，然后再根据估计后的光源色度信息利用对角模型对原始图像进行矫正。
3.基于传统卷积神经网络的方法中，使用相同的卷积核来处理不同的图像区域，这种只考虑图像局部信息进行建模，影响了模型的性能。transformer作为卷积神经网络的替代方案，它利用一种自注意力机制来捕获上下文之间的全局信息。然而，该类方法通常将输入图像分成固定大小的块(如32
×
32)，并独立处理每个块。这种方式也不可避免地会产生两个缺点：首先，边界像素不能利用块之外的相邻像素进行注意力信息获取。其次，得到的图像会在每个块周围引入伪影现象。虽然这个问题可以通过块重叠来缓解，但它会带来额外的计算开销。
4.申请公开号为cn112581373a的发明专利申请公开了一种深度学习的图像颜色矫正方法，首先，收集不同光照环境下的图像数据构建色偏图像数据集。之后，对数据进行数据增强和预处理，使数据符合神经网络训练要求。然后，构建一种encoder-decoder结构的基础卷积神经网络，并依据该基础神经网络构建级联的卷积神经网络。同时，针对级联神经网络的特点，设计了特定的神经网络损失函数。通过训练之后，该级联卷积神经网络可将一张色偏的原图片转化为无色偏的图片。该方法的缺点是该网络结构采用u型结构，它能够进行除噪和色彩校正，但恢复的结果图像在某些情况下可能存在一些缺陷，例如模糊、光晕现象，主要原因是过多的池化和上采样很容易产生光晕和模糊。


技术实现要素：

5.为了解决上述的技术问题，本发明提出了一种基于transformer的图像颜色校正方法，融合了cnn和transformer的优点，一方面具有cnn处理大尺寸图像的优势；另一方面具有transformer的优势，可以使用移位窗口策略对全局信息进行建模。
6.本发明提供一种基于transformer的图像颜色校正方法，包括输入原始带颜色投影的图像，还包括以下步骤：
7.步骤1：构建图像颜色校正模型，所述图像颜色校正模型包括浅层特征提取模块、深层特征提取模块和图像重构模块；
8.步骤2：对所述图像颜色校正模型进行训练；
9.步骤3：将所述原始带颜色投影的图像输入到训练后的所述图像颜色校正模型中，得到最后校正后彩色图像。
10.优选的是，所述浅层特征提取模块的输入为所述原始带颜色投影的图像，大小为h
×w×
3，利用1
×
1卷积使得输出是大小为h
×w×
120的特征图。
11.在上述任一方案中优选的是，所述深层特征提取模块的stb层的输入为所述大小为h
×w×
120的特征图，输出为提取后的深层语义特征。
12.在上述任一方案中优选的是，所述stb层的结构依次为：输入特征、ln层、w-msa子模块或sw-msa子模块、残差连接层、ln层、多层感知机层、残差连接层和输出特征。
13.在上述任一方案中优选的是，所述stb层包括2部分，左侧部分采用w-msa子模块，右侧部分采用sw-msa子模块。
14.在上述任一方案中优选的是，所述ln层用于对输入数据进行归一化处理。
15.在上述任一方案中优选的是，所述多层感知机层由两层全连接神经网络组成，第一层全连接网络先把输入向量从120维度映射到480维度，激活函数为relu函数，第二层全连接网络再从480维度映射回120维度，不使用激活函数。
16.在上述任一方案中优选的是，所述全连接网络的表达式为
17.ffn(x)＝max(0,xw1+b1)w2+b218.其中，x为输入特征图,w1为第一层权重,b1为第一层偏置,w2为第二层权重,b2为第二层偏置。
19.在上述任一方案中优选的是，所述w-msa子模块用于对输入特征进行窗口划分，在每个划分的小窗口上进行多头自注意力计算。
20.在上述任一方案中优选的是，所述w-msa子模块用于将输入特征映射到不同的子空间中，再分别在所有子空间上做点乘运算计算注意力向量，最后把所有子空间计算得到的注意力向量拼接起来，并利用w0做进一步映射得到最终的注意力向量作为输出。
21.在上述任一方案中优选的是，所述w-msa子模块的表达式为
22.multihead(q,k,v)＝concat(head1,
…
,headh)w0[0023][0024]
其中，q,k,v分别为w-msa子模块的查询向量、键向量、值向量，w0为线性映射权重，h为子空间的个数，attention为头注意力函数，为第j子空间中q的映射矩阵，为第j子空间中k的映射矩阵，为第j子空间中v的映射矩阵。
[0025]
在上述任一方案中优选的是，在一个单独子空间上的所述注意力向量的计算方式依次为：先将查询向量q和键向量k点乘，再除以键向量k的维度的平方根dk得到查询向量q的分数矩阵，然后通过softmax函数将分数矩阵归一化得到权重矩阵，最后乘以值向量v即得到一个子空间的注意力向量，具体形式为：
[0026][0027]
其中，k
t
为键向量k的转置，dk为键向量k的维度的平方根。
[0028]
在上述任一方案中优选的是，所述w-msa子模块的工作方法包括以下子步骤：
[0029]
步骤11：将所述大小为h
×w×
120的特征图循环左移和循环上移半个窗口的大小；
[0030]
步骤12：按照w-msa子模块切分窗口，得到和w-msa不同的窗口内容；
[0031]
步骤13：结合mask矩阵，计算每个窗口的attention；
[0032]
步骤14：得到的图像特征再循环下移和循环右移半个窗口的大小，恢复到原来的位置。
[0033]
在上述任一方案中优选的是，每个窗口的attention，公式为：
[0034][0035]
其中，mask为初始位置值均为0的矩阵。
[0036]
在上述任一方案中优选的是，所述步骤2包括利用l1损失和感知损失函数加权平均作为模型的总损失函数l
total
，公式为
[0037]
l
total
＝0.5l1+0.5l
perc
[0038][0039][0040]
其中，l1表示像素级别的l1范数损失，l
perc
表示感知损失，i
gt
代表真实图像，ic代表校正后图像，cj代表通道，hj和wj分别代表第j特征图的高度和宽度，代表在预训练模型vgg16中最后卷积层输出的特征图。
[0041]
本发明提出了一种基于transformer的图像颜色校正方法，能够恢复场景中真实内容信息，在智能安防、智能驾驶、智慧交通等多个领域有广泛的应用前景。
[0042]
swin transformer是指一种基于窗口的transformer方法。
[0043]
stb层是指基于swin transformer的块结构。
[0044]
ln层是指layer normalization，一种层归一化方法。
[0045]
w-msa子模块是指基于窗口的多头自注意力机制。
[0046]
sw-msa子模块是指基于移动窗口的多头自注意力机制。
附图说明
[0047]
图1为按照本发明的基于transformer的图像颜色校正方法的一优选实施例的流程图。
[0048]
图2为按照本发明的基于transformer的图像颜色校正方法的一优选实施例的整体架构示意图。
[0049]
图3为按照本发明的基于transformer的图像颜色校正方法的浅层特征提取模块
的一实施例的示意图。
[0050]
图4为按照本发明的基于transformer的图像颜色校正方法的stb层机构的一实施例的示意图。
[0051]
图5为按照本发明的基于transformer的图像颜色校正方法的mask矩阵计算的一实施例的示意图。
[0052]
图6为按照本发明的基于transformer的图像颜色校正方法的图像重构的一实施例的示意图。
具体实施方式
[0053]
下面结合附图和具体的实施例对本发明做进一步的阐述。
[0054]
实施例一
[0055]
如图1所示，执行步骤100，输入原始带颜色投影的图像。
[0056]
执行步骤110，构建图像颜色校正模型，所述图像颜色校正模型包括浅层特征提取模块、深层特征提取模块和图像重构模块。
[0057]
所述浅层特征提取模块的输入为所述原始带颜色投影的图像，大小为h
×w×
3，利用1
×
1卷积使得输出是大小为h
×w×
120的特征图。
[0058]
所述深层特征提取模块的stb层的输入为所述大小为h
×w×
120的特征图，输出为提取后的深层语义特征。所述stb层的结构依次为：输入特征、ln层、w-msa子模块或sw-msa子模块、残差连接层、ln层、多层感知机层、残差连接层和输出特征。
[0059]
所述ln层用于对输入数据进行归一化处理。
[0060]
所述多层感知机层由两层全连接神经网络组成，第一层全连接网络先把输入向量从120维度映射到480维度，激活函数为relu函数，第二层全连接网络再从480维度映射回120维度，不使用激活函数。所述全连接网络的表达式为
[0061]
ffn(x)＝max(0,xw1+b1)w2+b2[0062]
其中，x为输入特征图,w1为第一层权重,b1为第一层偏置,w2为第二层权重,b2为第二层偏置。
[0063]
所述stb层包括2部分，左侧部分采用w-msa子模块，右侧部分采用sw-msa子模块。
[0064]
所述w-msa子模块用于对输入特征进行窗口划分，在每个划分的小窗口上进行多头自注意力计算。所述w-msa子模块用于将输入特征映射到不同的子空间中，再分别在所有子空间上做点乘运算计算注意力向量，最后把所有子空间计算得到的注意力向量拼接起来，并利用w0做进一步映射得到最终的注意力向量作为输出。所述w-msa子模块的表达式为
[0065]
multihead(q,k,v)＝concat(head1,
…
,headh)w0[0066][0067]
其中，q,k,v分别为w-msa子模块的查询向量、键向量、值向量，w0为线性映射权重，h为子空间的个数，attention为头注意力函数，为第j子空间中q的映射矩阵，为第j子空间中k的映射矩阵，为第j子空间中v的映射矩阵。
[0068]
在一个单独子空间上的所述注意力向量的计算方式依次为：先将查询向量q和键
向量k点乘，再除以键向量k的维度的平方根dk得到查询向量q的分数矩阵，然后通过softmax函数将分数矩阵归一化得到权重矩阵，最后乘以值向量v即得到一个子空间的注意力向量，具体形式为：
[0069][0070]
其中，k
t
为键向量k的转置，dk为键向量k的维度的平方根。
[0071]
所述w-msa子模块的工作方法包括以下子步骤：
[0072]
执行步骤111，将所述大小为h
×w×
120的特征图循环左移和循环上移半个窗口的大小；
[0073]
执行步骤112，按照w-msa子模块切分窗口，得到和w-msa不同的窗口内容；
[0074]
执行步骤113，结合mask矩阵，计算每个窗口的attention，每个窗口的attention的公式为：
[0075][0076]
其中，mask为初始位置值均为0的矩阵。
[0077]
执行步骤114，得到的图像特征再循环下移和循环右移半个窗口的大小，恢复到原来的位置。
[0078]
执行步骤120，对所述图像颜色校正模型进行训练，利用l1损失和感知损失函数加权平均作为模型的总损失函数l
total
，公式为
[0079]
l
total
＝0.5l1+0.5l
perc
[0080][0081][0082]
其中，l1表示像素级别的l1范数损失，l
perc
表示感知损失，i
gt
代表真实图像，ic代表校正后图像，cj代表通道，hj和wj分别代表第j特征图的高度和宽度，代表在预训练模型vgg16中最后卷积层输出的特征图。
[0083]
执行步骤130，将所述原始带颜色投影的图像输入到训练后的所述图像颜色校正模型中，得到最后校正后彩色图像。
[0084]
实施例二
[0085]
本发明主要基于swin transformer模型而提出的一种图像颜色校正方法，它由三个模块组成：浅层特征提取模块、深层特征提取模块和图像重建模块。浅层特征提取模块利用卷积层实现，作为残差连接通道，直接传输到重构模块，以保留图像底层纹理信息。深度特征提取模块主要由多个residual swin transformer(rst)层组成，每个层利用1个swin transformer块进行局部注意力信息和全局信息建模。重建模块中首先融合浅层和深层特征，再利用两个卷积层实现图像复原，得到最后校正后彩色图像。整个框架图如图2所示：
[0086]
(1)浅层特征提取模块利用patch embedding技术，直接以原始图像作为输入，patch embedding将原始图像分通过1
×
1卷积将图像维度应用到线性嵌入层，本发明实施中嵌入层的维度设置为120。
[0087]
(2)深层特征提取模块利用移动窗口(shift window)策略，对输入的图像特征图进行窗口划分。先在划分后的窗口里单独使用多头自注意力机制(multi-head self-attention,msa)对局部依赖关系进行建模，从而提取布局图像特征。然后通过平移窗口位置，解决不同窗口之间信息共享，利用全局依赖关系进行建模。
[0088]
(3)重建模块包含两层卷积，用于图像维度映射，使得输出特征图的维度为3，特征图通道的变化为120
→
64
→
3。
[0089]
本发明的方法，利用图2所示的网络框架，按照以下步骤具体实施：
[0090]
步骤1：构建浅层特征提取模块：如图3所示，该处理模块的输入为原始带颜色投影的图像，大小为h
×w×
3，利用1
×
1卷积使得输出是大小为h
×w×
120的特征图。
[0091]
步骤2：构建高层特征提取模块：基于swin transformer提取图像高层语义信息，如图4所示。stb层(swin transformer block)的输入是步骤1的输出特征，大小为h
×w×
120。输出是提取后的高层语义特征，大小同样为h
×w×
120。stb利用其自注意力机制建模不同空间位置特征之间的依赖关系，有效地捕获上下文信息。另外，stb在网络中融入局部特征和全局信息，其结构避免了cnn堆叠卷积层来获取全局信息的模式，使模型具有很好的特征提取能力。
[0092]
如图4所示，stb块的结构依次为：步骤1的输出特征作为输入特征
→
ln(layer norm)层
→
w-msa子模块(窗口多头自注意力层)或者sw-msa子模块(移位窗口多头自注意力层)
→
残差连接层
→
ln层
→
mlp(多层感知机层)
→
残差连接层
→
输出特征。stb包括2部分，左侧采用的是wmsa子模块，右侧采用sw-msa子模块。ln层主要作用是对输入数据进行归一化处理，将其归到0到1之间，从而保证输入层数据分布是一样的；mlp由两层全连接神经网络组成，第一层全连接网络先把输入向量从120维度映射到480维度，激活函数为relu函数，第二层全连接网络再从480维度映射回120维度，不使用激活函数，全连接网络的表达式如下式(1)：
[0093]
ffn(x)＝max(0,xw1+b1)w2+b2ꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀ
(1)
[0094]
残差连接层主要是进行残差连接，解决梯度消失问题。
[0095]
w-msa子模块首先对输入特征进行窗口划分，本发明实施例中设置的划分窗口大小为8
×
8，在每个划分的小窗口上进行多头自注意力计算。w-msa子模块将输入特征映射到不同的子空间中，再分别在所有子空间上做点乘运算计算注意力向量，最后把所有子空间计算得到的注意力向量拼接起来，并利用w0做进一步映射得到最终的注意力向量作为输出，w-msa子模块的表达式如公式(2)所示：
[0096]
multihead(q,k,v)＝concat(head1,
…
,headh)w0[0097][0098]
其中，q,k,v分别为w-msa子模块的查询向量、键向量、值向量，为第j子空间中q的映射矩阵，为第j子空间中k的映射矩阵，为第j子空间中v的映射矩阵，本步骤
实施例中设置的子空间的个数h设为6，在一个单独子空间上的注意力向量(20维)的计算方式依次为：先将查询向量q和键向量k点乘，再除以键向量k的维度的平方根得到查询向量q的分数矩阵(定义为a)，然后通过softmax函数将分数矩阵归一化得到权重矩阵，最后乘以值向量v即得到一个子空间的注意力向量，具体形式如公式(3)所示：
[0099][0100]
w-msa子模块通过将输入特征映射到不同的子空间，利用获得的注意力分数捕获特征在不同子空间上的依赖关系。单独使用w-msa子模块得到的网络的建模能力有限，因为它将每个窗口当作一个独立区域计算。它忽略了窗口之间的联系。为了解决这个问题，本发明利用sw-msa子模块获取全局依赖关系。sw-msa子模块首先将图像特征进行尺寸为半个窗口大小的像素位移操作，然后基于mask机制进行w-msa子模块操作。即形成和w-msa子模块所划分同一位置的窗口里面包含的图像特征信息不同，即解决了不同窗口之间信息间交互的问题。具体操作流程如下：
[0101]
将步骤1的输出特征循环左移和循环上移半个窗口的大小，在移位的基础上再按照w-msa子模块切分窗口，就会得到和w-msa不同的窗口内容；然后计算每个窗口的mask,如图5所示，假设图像大小为48
×
48，窗口大小为16
×
16，得到9个窗口，标号1-9，其中标号1，2，4，5，mask矩阵所有位置值为0，按照w-msa模块计算方式进行attention计算。标号3，6共享一组mask矩阵,其中黑色位置mask值为0，白色位置mask值为-100，同理，标号7，8共享一组mask矩阵，标号9窗口单独设置一个mask矩阵。结合mask矩阵，每个窗口计算attention，如公式(4)所示：
[0102][0103]
得到的图像特征再循环下移和循环右移半个窗口的大小，恢复到原来的位置。
[0104]
步骤3：图像重构模块，该重构模块的输入数据是步骤2的输出特征，大小为h
×w×
120；该重构模块的输出是校正后彩色图像，大小h
×w×
3，利用2次卷积操作实现，如图6所示。经过conv1层，卷积核大小为1
×
1，步长为1，输出层大小h
×w×
64，再经过conv2层，卷积核大小为3
×
3，步长为1，输出通道为3，构成校正后图像。
[0105]
模型训练实施阶段，构建4650组图像匹配对作为训练样本。对于损失函数，由于l1损失函数在图像的整体对比度、均匀区域的平滑效果方面表现较好，感知损失函数能够约束真实图像和预测图像之间的差异，保持更多的原图像的细节信息。本专利中，利用l1损失和感知损失函数加权平均作为模型的总损失函数，如公式(5)所示：
[0106]
l
total
＝0.5l1+0.5l
perc
[0107][0108][0109]
式中，l1表示像素级别的l1范数损失，l
perc
表示感知损失，i
gt
代表真实图像，ic代表校正后图像，cj代表通道，hj和wj分别代表第j特征图的高度和宽度，代表在预训练模型
vgg16中最后卷积层输出的特征图。
[0110]
实施例3
[0111]
本专利提出的颜色校正方法能够恢复场景中真实内容信息，在智能安防、智能驾驶、智慧交通等多个领域有广泛的应用前景。同时，生活娱乐方面，手机、相机等已经成为人们日常生活必不可少的一部分，人们对图像质量的要求越来越高，各大手机商(如华为、苹果)内部就配有白平衡(颜色恒常性)功能，但这类算法不是端到端的学习方式，是建立图像内容和场景光照之间的映射关系，利用得到的光照颜色之后，还要进行后续处理。本专利提出的方法是一种端到端方式，模型的输出结果就是校正后彩色图像，去除光照的影响，恢复物体本来颜色，使之有利于进行进一步的图像分析和特征提取，具有较高的应用价值。
[0112]
为了更好地理解本发明，以上结合本发明的具体实施例做了详细描述，但并非是对本发明的限制。凡是依据本发明的技术实质对以上实施例所做的任何简单修改，均仍属于本发明技术方案的范围。本说明书中每个实施例重点说明的都是与其它实施例的不同之处，各个实施例之间相同或相似的部分相互参见即可。对于系统实施例而言，由于其与方法实施例基本对应，所以描述的比较简单，相关之处参见方法实施例的部分说明即可。