一种基于CNN和ViT的红外图像彩色化方法

本发明涉及图像处理，具体为一种基于cnn和vit的红外图像彩色化方法。

背景技术：

1、将红外图像转换为可见光图像是一项非常有挑战性的任务。在灰度图像彩色化中，只需要将亮度值转换为色度值；而在红外图像彩色化中，需要在给定热特征的情况下同时估计亮度和色度。这导致红外图像彩色化的结果看起来不够真实，存在严重的细节模糊和纹理扭曲等问题。常见的彩色化方法一般采用卷积神经网络(convolutional neuralnetwork，cnn)和生成对抗网络(generative adversarial network，gan)实现彩色化功能。一方面cnn通过共享卷积核提取特征，减少网络参数数量，提高模型效率，另一方面cnn具有平移不变性，即无论特征被移动到图像的哪个位置，网络都能检测到这些特征。尽管cnn存在很多优势，但是其感受野通常很小，不利于捕获全局特征，会产生较为模糊的边缘以及细节信息，导致彩色化质量较差。而vision transformer(vit)通过自注意力模块可以获取长程依赖关系，扩大图像感受野，能够捕捉到更多的上下文信息。但是vit中自注意力模块的计算和开销与输入分辨率呈二次关系，不可避免地会带来gpu内存不足和计算效率低下的问题。

2、中国专利公开号为“cn115170430a”，名称为“基于两阶段条件生成对抗网络的近红外图像着色方法”，首先获取近红外灰度图像；其次构建近红外图像的灰度预处理模块；然后利用近红外图像的灰度预处理模块对近红外灰度图像进行预处理，得到灰度图像；随后构建图像着色模块；接着利用图像着色模块对着色生成图像和训练图像提取到的颜色特征向量进行判别，得到判别结果用于训练图像着色模块；最后对灰度图像进行着色处理，得到着色生成图像；该方法仅使用cnn去实现图像的彩色化，网络常常要设计得很深很宽，这使计算复杂度和存储容量占用率都居高不下；并且cnn对全局语义信息的提取不够完善，降低图像的彩色化质量。因此，如何有效地将cnn和vit结合实现高质量的彩色化结果是本领域技术人员亟需解决的问题。

技术实现思路

1、(一)解决的技术问题

2、针对现有技术的不足，本发明提供了一种基于cnn和vit的红外图像彩色化方法，解决了现有的红外图像彩色化方法得到的彩色化图像质量差和计算复杂度高的问题。

3、(二)技术方案

4、本发明为了实现上述目的具体采用以下技术方案：

5、一种基于cnn和vit的红外图像彩色化方法，该方法具体包括如下步骤：

6、s1，构建网络模型：整个生成对抗网络包括生成器和鉴别器；

7、s2，准备数据集：对整个生成对抗网络先用红外图像数据集一进行训练；

8、s3，训练网络模型：训练红外图像彩色化模型，将s2中准备好的数据集进行预处理，调整数据集中每个图像的尺寸，固定输入图像的大小，将处理好的数据集输入到s1中构建好的网络模型中进行训练；

9、s4，选择最小化损失函数和最优评估指标：通过最小化网络输出图像与标签的损失函数，直到训练次数达到设定阈值或损失函数的值到达设定范围内即可认为模型参数已预训练完成，保存模型参数；同时选择最优评估指标来衡量算法的精度，评估系统的性能；

10、s5，微调模型：用红外图像数据集二对模型进行训练和微调，得到稳定可用的模型参数，进一步提高模型的红外图像彩色化能力；最终使得模型对红外图像彩色化的效果更好；

11、s6，保存模型：将网络训练完成后，需要将网络中所有参数保存，之后用将要彩色化的红外图像输入到网络中就可以得到彩色化好的图像；该网络对输入图像大小没有要求，任意尺寸均可。

12、进一步地，所述生成器包括浅层特征提取模块、卷积耦合模块、transformer模块、卷积层、上采样模块和图像重建模块；

13、所述浅层特征提取模块，用于利用卷积块对图像进行浅层特征提取；

14、所述卷积耦合模块，用于减小特征图的大小，更好地提取局部信息；

15、所述transformer模块，主要由深度多头自注意力模块和局部特征前馈网络；用于对特征图进行特征转换，生成多尺度特征，更容易地捕获长程依赖关系和局部特征；

16、所述卷积层，用于丰富通道信息；

17、所述上采样模块，用于对特征图进行跨通道的融合，逐步恢复特征图大小；

18、所述图像重建模块，用于恢复通道数量，重建高质量的红外彩色化图像。

19、进一步地，所述鉴别器包括多个卷积块、transformer模块和s型函数，增加transformer模块用于加强了鉴别器的鉴别能力，促进生成器的生成能力，有助于快速收敛。

20、进一步地，第一红外图像数据集为kaist数据集，将kaist数据集打乱，构造出未配对的kaist数据集，通过对数据集中图像进行无监督的训练。

21、进一步地，在训练网络模型中预设阈值包括损失函数预设值、训练次数预设值。

22、进一步地，损失函数为复合损失函数，生成器采用的损失函数包括合成损失，对抗损失和patchnce损失；鉴别器采用对抗损失。

23、进一步地，在训练网络模型过程中还包括通过评价指标评估算法彩色化结果的质量和图像失真程度。

24、进一步地，所述第二红外图像数据集为flir数据集。

25、一种基于cnn和vit的红外图像彩色化系统，包括：

26、图像获取模块，用于获取待彩色化图像；

27、图像处理模块，用于对待彩色化图像进行预处理，并分为训练集和测试集；

28、模型训练模块，用于将处理好的待彩色化训练集图像输入到设计好的网络中进行训练，利用训练好的网络对待彩色化测试集图像的每个像素值进行预测，最终生成彩色化图像；

29、质量评估模块，用于评估最终生成的彩色化图像质量是否满足预设质量要求；若满足预设质量要求，则将生成的彩色化图像作为最终的彩色化效果；若不满足预设质量要求，则启动质量提升模块；

30、质量提升模块，用于重新获取数据集，并利用新的数据集继续训练模型，利用重新训练后的网络重新生成待彩色化图像对应的彩色化图像，然后返回质量评估模块。

31、(三)有益效果

32、与现有技术相比，本发明提供了一种基于cnn和vit的红外图像彩色化方法，具备以下有益效果：

33、1、本发明采用的网络结构是基于cnn和vit的生成对抗网络结构，结合了cnn获取局部特征和vit捕获长程依赖关系的优点，获得更大的感受野，降低了计算复杂度，充分利用了全局信息和局部信息，提高了彩色化能力，能够生成高质量的彩色化图像。

34、2、本发明在transformer模块一、transformer模块二和transformer模块三中使用组卷积，利用图像块内部的局部关系和结构信息，更好地提取局部信息。

35、3、本发明在深度多头自注意力模块中使用深度卷积层、线性层一、线性层二和线性层三，可以有效地减少索引值(k)和内容值(v)的空间大小，减轻计算开销。

36、4、本发明在局部特征前馈网络中使用深度卷积层，降低了计算成本，并增加跳跃链接提高梯度跨层的传播能力，以便于获得更好的彩色化效果。

37、5、本发明在鉴别器中加入transformer模块，加强了鉴别器的鉴别能力，促进生成器的生成能力，有助于快速收敛。

38、6、本发明提出了一种由合成损失、对抗损失和patchnce损失组成的复合损失函数，它可以提高彩色化图像的质量，生成精细的局部细节，恢复语义和纹理信息。