一种基于上下文感知离散码本的图像着色方法

技术领域

本申请属于图像着色技术领域，尤其涉及一种基于上下文感知离散码本的图像着色方法。

背景技术

图像着色的目的是从给定的灰度或亮度通道中预测缺失的色度通道，在旧资料恢复、图像或视频压缩以及艺术创作中有广泛的应用。着色任务将单通道的灰度图I

然而，着色是一个本质上具有不确定性的任务，这是由于一张灰度图像经过着色后的每一个部分可以有不同的颜色，且可以都是合理的(例如，花朵可以是红色、黄色或白色的)，因此，这造成了着色任务中的不确定性和多样性本质。

图像着色方法大致可以分为传统的基于参考信息的方法和全自动的方法。对于传统的基于参考信息的方法，部分研究人员倾向于寻求用户的协助，如指定个别像素点的颜色，然后使用模型将对应的区块着色；部分使用示例图像或颜色提示来引导上色，如给定模板的颜色图，让模型根据该模板图的总体颜色或者语义信息，来减少不确定性并指导着色的过程。这些方法通过匹配来自亮度通道的空间内容，然后将示例图像的色度信息注入到灰度图像中。然而，获取可行的示例图像或用户协助并不容易，并且需要很多费力的人工指引以保证着色结果的准确性。在全自动的着色方法当中，部分研究者使用端到端的模式训练着色模型，但是如此通常会导致模型学习到较为平均的结果，导致低对比度的着色图像。因此研究者引入了额外的信息来辅助训练并着色，如使用了显著图，直方图模块等。但是在现实的应用当中，想要获取高质量的辅助信息仍然存在一定的困难。

发明内容

本申请的目的是提供一种基于上下文感知离散码本的图像着色方法，以解决现有图像着色技术所存在的上述技术问题，通过在特定方向遍历潜在空间改变编码序列，解决了着色当中存在的复杂的多输出问题。

为了实现上述目的，本申请技术方案如下：

一种基于上下文感知离散码本的图像着色方法，包括：

构建预训练模型和图像着色模型，所述预训练模型包括色度编码器、亮度编码器、第一注意力模块、向量量化模块和解码器，所述图像着色模型包括亮度编码器、第二注意力模块、向量量化模块和解码器，训练所述预训练模型；

将预训练模型中训练好的亮度编码器、向量量化模块和解码器的网络参数，作为图像着色模型中对应模块的网络参数，固定向量量化模块和解码器的网络参数，训练所述图像着色模型；

采用训练好的图像着色模型，对待着色图像进行着色。

进一步的，所述训练所述预训练模型，包括：

将图像的色度通道图像和亮度通道图像分别输入到亮度编码器和色度编码器，得到亮度编码特征和色度编码特征；

在第一注意力模块中，首先对亮度编码特征和色度编码特征分别执行自注意力操作，得到亮度自注意力特征和色度自注意力特征，然后将亮度自注意力特征和色度自注意力特征执行交叉注意力操作，得到交叉注意力编码特征；

在向量量化模块中，对交叉注意力编码特征进行矢量量化操作，得到量化特征；

在解码器中，将量化特征重建为彩色图像。

进一步的，所述将亮度自注意力特征和色度自注意力特征执行交叉注意力操作，包括：

依次执行第一预设次数的交叉注意力操作；

在第一次交叉注意力操作时，通过亮度自注意力特征映射得到查询向量，通过亮度自注意力特征和色度自注意力特征映射再拼接得到键向量和值向量，执行交叉注意力操作；

之后的交叉注意力操作时，通过前一次交叉注意力操作的输出特征映射得到查询向量，通过前一次交叉注意力操作的输出特征和色度自注意力特征映射再拼接得到键向量和值向量，执行交叉注意力操作。

进一步的，所述训练所述预训练模型，采用的联合损失函数如下：

其中λ

其中，sg[·]表示梯度计算停止算子，

进一步的，所述固定向量量化模块和解码器的网络参数，训练所述图像着色模型，包括：

固定向量量化模块和解码器的网络参数，将图像的亮度通道图像输入到亮度编码器，得到亮度编码特征；

在第二注意力模块中，对亮度编码特征进行锚定注意力操作，得到锚定注意力特征；

将锚定注意力特征归一化后，输入向量量化模块进行矢量量化操作，得到量化特征；

在解码器中，将量化特征重建为彩色图像。

进一步的，所述对亮度编码特征进行锚定注意力操作，包括：

依次执行第二预设次数的锚定注意力操作；

在第一次锚定注意力操作时，以亮度编码特征作为输入特征和锚定特征，执行锚定注意力操作；

之后的锚定注意力操作时，以前一个锚定注意力操作的输出特征作为输入特征，以亮度编码特征作为锚定特征，执行锚定注意力操作。

进一步的，所述训练所述图像着色模型，采用的联合损失函数如下：

其中，λ

其中，sg[·]表示梯度计算停止算子，

进一步的，所述图像着色模型还包括一个线性层，所述基于上下文感知离散码本的图像着色方法，还包括：

固定图像着色模型中亮度编码器、第二注意力模块、向量量化模块和解码器模块，训练线性层；

在采用训练好的图像着色模型对待着色图像进行着色时，将预设的位移和方向输入训练好的线性层，对初始化得到的正交矩阵M进行线性变换，得到线性变换后的矩阵；

将第二注意力模块输出的锚定注意力特征与线性变换后的矩阵相加，然后经过向量量化模块和解码器模块，输出多样化的着色图像。

本申请提出的一种基于上下文感知离散码本的图像着色方法，无需任何额外的辅助信息引导着色过程，确保了在实际应用中更高的实用性；通过学习到的可感知图像内容的离散码本和提出的查询模式，模型具备记忆图像空间特征和颜色的映射关系的能力，确保生成的图像的高质量和高语义匹配度。本申请还解决了着色任务当中的多输出的难题，提供了用户可选的方式来产生多样性输出内容的同时，保证了产出图像的语义合理性并降低了颜色溢出的效益。

附图说明

图1为本申请基于上下文感知离散码本的图像着色方法流程图。

图2为本申请预训练模型示意图。

图3为本申请图像着色模型示意图。

图4为本申请图像着色模型另一个实施例示意图。

具体实施方式

为了使本申请的目的、技术方案及优点更加清楚明白，以下结合附图及实施例，对本申请进行进一步详细说明。应当理解，此处描述的具体实施例仅用以解释本申请，并不用于限定本申请。

在一个实施例中，如图1所示，提出了一种基于上下文感知离散码本的图像着色方法，包括：

步骤S1、构建预训练模型和图像着色模型，所述预训练模型包括色度编码器、亮度编码器、第一注意力模块、向量量化模块和解码器，所述图像着色模型包括亮度编码器、第二注意力模块、向量量化模块和解码器，训练所述预训练模型。

本实施例通过构建预训练模型，然后进行预训练，以得到图像着色模型的向量量化模块和解码器。预训练模型如图2所示，包括色度编码器、亮度编码器、第一注意力模块、向量量化模块和解码器。

在训练预训练模型时，执行如下操作：

步骤1.1、将图像的色度通道图像和亮度通道图像分别输入到亮度编码器和色度编码器，得到亮度编码特征和色度编码特征。

本步骤利用亮度编码器和色度编码器，对输入的色度通道图像I

具体的，先将给定图像样本转换到CIELAB颜色空间，得到

本步骤通过两个单独的编码器E

步骤1.2、在第一注意力模块中，首先对亮度编码特征和色度编码特征分别执行自注意力操作，得到亮度自注意力特征和色度自注意力特征，然后将亮度自注意力特征和色度自注意力特征执行交叉注意力操作，得到交叉注意力编码特征。

本步骤在第一注意力模块中完成，第一注意力模块中包括自注意力操作和交叉注意力操作。首先分别对亮度编码特征和色度编码特征执行自注意力操作，得到亮度自注意力特征x

之后，对亮度自注意力特征x

在一个具体的实施例中，本申请将亮度自注意力特征和色度自注意力特征执行交叉注意力操作，包括：

依次执行第一预设次数的交叉注意力操作；

在第一次交叉注意力操作时，通过亮度自注意力特征映射得到查询向量，通过亮度自注意力特征和色度自注意力特征映射再拼接得到键向量和值向量，执行交叉注意力操作；

之后的交叉注意力操作时，通过前一次交叉注意力操作的输出特征映射得到查询向量，通过前一次交叉注意力操作的输出特征和色度自注意力特征映射再拼接得到键向量和值向量，执行交叉注意力操作。

具体的，本实施例中需依次执行交叉注意力N

在第一次交叉注意力操作时，亮度自注意力特征x

q＝q

k

k＝Concat(k

其中，w

在之后的其他次的交叉注意力操作时，与第一次交叉注意力操作不同的是，以前一次交叉注意力操作的输出特征代替亮度自注意力特征进行交叉注意力操作，具体的交叉操作相同，这里不再赘述。

在执行交叉注意力之后，获得的交叉注意力编码特征

步骤1.3、在向量量化模块中，对交叉注意力编码特征进行矢量量化操作，得到量化特征。

具体的，通过最近邻方法(Nearest Neighbor)对交叉注意力编码特征

其中，

步骤1.4、在解码器中，将量化特征重建为彩色图像。

具体的，本实施例使用解码器D，在给定特征编码z

重建色度通道图像

在一个具体的实施例中，解码器的结构与VQGAN相同，由12个残差模块以及5个上采样模块组成。

步骤S2、将预训练模型中训练好的亮度编码器、向量量化模块和解码器的网络参数，作为图像着色模型中对应模块的网络参数，固定向量量化模块和解码器的网络参数，训练所述图像着色模型。

本实施例通过对预训练模型的训练，可以得到预训练模型中亮度编码器、向量量化模块和解码器的网络参数，将这些网络参数传递到图像着色模型中对应模块，然后固定图像着色模型中向量量化模块和解码器，训练图像着色模型。

其中，所述训练图像着色模型，包括：

步骤2.1、固定向量量化模块和解码器的网络参数，将图像的亮度通道图像输入到亮度编码器，得到亮度编码特征。

图像着色模型如图3所示，包括亮度编码器、第二注意力模块、向量量化模块和解码器，其中第二注意力模块包括多个锚定锚块，执行锚定注意力操作。

训练图像着色模型时，固定向量量化模块和解码器的网络参数。以图像样本的亮度通道图像I

步骤2.2、在第二注意力模块中，对亮度编码特征进行锚定注意力操作，得到锚定注意力特征。

具体的，将步骤2.1得到的亮度编码特征

在一个优选的实施例中，所述对亮度编码特征进行锚定注意力操作，包括：

依次执行第二预设次数的锚定注意力操作；

在第一次锚定注意力操作时，以亮度编码特征作为输入特征和锚定特征，执行锚定注意力操作；

之后的锚定注意力操作时，以前一个锚定注意力操作的输出特征作为输入特征，以亮度编码特征作为锚定特征，执行锚定注意力操作。

具体的，本实施例中需依次执行N2锚定注意力操作，目的是预测更准确的码序列，在所有N2个锚定注意力模块中集成相同的亮度编码特征

在第一次锚定注意力操作时，以亮度编码特征作为输入特征和锚定特征，采用公式表示为：

其中q

在之后的其他次的锚定注意力操作时，与第一次锚定注意力操作不同的是，以前一次锚定注意力操作的输出特征作为输入特征，锚定特征不变，具体的锚定注意力操作相同，这里不再赘述。

本实施例执行锚定注意力N

步骤2.3、将锚定注意力特征归一化后，输入向量量化模块进行矢量量化操作，得到量化特征。

本步骤将锚定注意力特征

其中，q(·)表示向量量化操作，LayerNorm是一种层归一化方法。

步骤2.4、在解码器中，将量化特征重建为彩色图像。

然后，冻结的解码器D

本实施例为了提高码序列预测的效果，引入了锚定自注意力，并且在保留有效性的同时简化了结构，使其适应图像着色和码预测任务。训练图像着色模型时，冻结了离散码本

步骤S3、采用训练好的图像着色模型，对待着色图像进行着色。

在训练好图像着色模型之后，将待着色图像输入到亮度编码器，然后经过第二注意力模块、向量量化模块和解码器，得到着色后的图像，完成图像的着色。

在一个具体的实施例中，训练预训练模型时，联合损失函数包括重建损失、感知损失、对抗损失和双重承诺损失，联合损失函数

其中λ

其中：重建损失

其中，

感知损失

其中，MSE代表均方误差，

对抗损失

其中，Discriminator(*)表示判别器。

本实施例引入判别器D以及对抗损失

双重承诺损失

其中，sg[·]表示梯度计算停止算子。

本实施例提出需要约束码序列层级的特征，使用双重承诺损失

在另一个具体的实施例中，训练图像着色模型时，联合损失函数包括散度损失和码本损失，联合损失函数

其中，λ

其中，散度损失

其中，

码本损失

本实施例码本损失

在另一个实施例中，所述图像着色模型还包括一个线性层，如图4所示。本实施例一种基于上下文感知离散码本的图像着色方法，还包括：

固定亮度编码器、第二注意力模块、向量量化模块和解码器模块，训练线性层；

在采用训练好的图像着色模型对待着色图像进行着色时，将预设的位移和方向输入训练好的线性层，对初始化得到的正交矩阵M进行线性变换，得到线性变换后的矩阵；

将第二注意力模块输出的锚定注意力特征与线性变换后的矩阵相加，然后经过向量量化模块和解码器模块，输出多样化的着色图像。

具体的，前面的实施例已经通过查找预训练的离散码本和通过解码器生成图像，产生了最可能的着色结果。然而，图像着色是一个对于单个灰度图片有多种可行和合理结果的问题。本实施例希望通过在特定方向遍历潜在空间并改变码序列产出多样但合理的着色结果。

如图4所示，在第二注意力模块之后，引入了一个线性层，该线性层用于对正交矩阵

本实施例可以根据不同的位移δ

其中，z

在训练时，固定亮度编码器、第二注意力模块、向量量化模块和解码器模块，将训练样本图像输入到模型中，第二注意力模块输出锚定注意力特征z

I

I

第一幅图像是给定z

将这两幅图像输入到一个位移和方向预测器

在训练时根据不同的幅度δ

训练时，使用两个预测级别的损失，即预测的

其中，CrossEntropy表示交叉熵损失，MAE表示均方误差损失。

此外，引入另外两个图像级别的损失：因为需要保持变换后的图像与原始图像具有相同的类别，因此添加了分类损失：

具体的，其中本发明使用预训练的ResNet-50模型作为分类器class(·)。变换后的图像虽然具有多样化的颜色，但是它们可能会遭受更严重的颜色渗透溢出效应，因此在训练过程当中添加了一个基于梯度的损失函数：

其中

上述四种损失函数的联合损失函数如下：

其中λ

在计算上述联合损失函数之后，进行反向传播，从而完成对线性层的训练。

本申请一种基于上下文感知离散码本的图像着色方法，首先学习一个离散码本，通过将着色任务转化为潜在空间的编码序列来简化I

本申请技术方案还通过在特定方向遍历潜在空间并改变码序列产出多样但合理的着色结果，目标是通过在特定方向k遍历潜在空间并改变码序列产出多样但合理的着色结果。由于步骤S2已经通过查找预训练的离散码本和通过解码器生成图像，产生了最可能的着色结果。引入了变换矩阵M以及移动方向k和距离δ的预测器P，配合针对此情况设计的损失函数，可以以无监督的方式探索出真实且高质量的多着色结果。

以上所述实施例仅表达了本申请的几种实施方式，其描述较为具体和详细，但并不能因此而理解为对发明专利范围的限制。应当指出的是，对于本领域的普通技术人员来说，在不脱离本申请构思的前提下，还可以做出若干变形和改进，这些都属于本申请的保护范围。因此，本申请专利的保护范围应以所附权利要求为准。