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

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

背景技术：

1、图像着色的目的是从给定的灰度或亮度通道中预测缺失的色度通道，在旧资料恢复、图像或视频压缩以及艺术创作中有广泛的应用。着色任务将单通道的灰度图il作为输入，目标是去预测定义在cielab颜色空间的2个色度通道iab。

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

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

技术实现思路

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

17、

18、其中λ1，λ2和λ3均代表对应损失函数的权重，表示训练所述预训练模型的联合损失函数；表示重建损失，表示感知损失，表示对抗损失，表示双重承诺损失；

19、

20、其中，sg[·]表示梯度计算停止算子，表示交叉注意力编码特征，zq表示量化特征。

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

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

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

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

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

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

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

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

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

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

31、

32、其中，λ4表示对应损失函数的权重，表示训练所述图像着色模型的联合损失函数，表示散度损失，表示码本损失：

33、

34、

35、其中，sg[·]表示梯度计算停止算子，表示锚定注意力特征经过归一化后的特征，表示预训练模型的交叉注意力编码特征，表示图像着色模型中量化模块输出的量化特征。

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

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

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

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

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

技术特征：

1.一种基于上下文感知离散码本的图像着色方法，其特征在于，所述基于上下文感知离散码本的图像着色方法，包括：

2.根据权利要求1所述的基于上下文感知离散码本的图像着色方法，其特征在于，所述训练所述预训练模型，包括：

3.根据权利要求2所述的基于上下文感知离散码本的图像着色方法，其特征在于，所述将亮度自注意力特征和色度自注意力特征执行交叉注意力操作，包括：

4.根据权利要求2所述的基于上下文感知离散码本的图像着色方法，其特征在于，所述训练所述预训练模型，采用的联合损失函数如下：

5.根据权利要求1所述的基于上下文感知离散码本的图像着色方法，其特征在于，所述固定向量量化模块和解码器的网络参数，训练所述图像着色模型，包括：

6.根据权利要求5所述的基于上下文感知离散码本的图像着色方法，其特征在于，所述对亮度编码特征进行锚定注意力操作，包括：

7.根据权利要求5所述的基于上下文感知离散码本的图像着色方法，其特征在于，所述训练所述图像着色模型，采用的联合损失函数如下：

8.根据权利要求1所述的基于上下文感知离散码本的图像着色方法，其特征在于，所述图像着色模型还包括一个线性层，所述基于上下文感知离散码本的图像着色方法，还包括：

技术总结
本发明公开了一种基于上下文感知离散码本的图像着色方法，属于图像着色人工智能应用领域，首先训练预训练模型，学习一个可感知上下文的离散码本。然后将预训练模型中训练好的亮度编码器、向量量化模块和解码器的网络参数，作为图像着色模型中对应模块的网络参数，固定向量量化模块和解码器的网络参数，训练所述图像着色模型。最后采用训练好的图像着色模型，对待着色图像进行着色。本发明无需任何额外的辅助信息引导着色过程，确保了在实际应用中更高的实用性，可生成生动且逼真的着色图像。

技术研发人员：穆攀,刘哲源,许函宁,白琮
受保护的技术使用者：浙江工业大学
技术研发日：
技术公布日：2024/1/15