基于改进Swin Transformer的交互式灰度图像着色方法的研究

Research on Interactive Grayscale Image Coloring Method Based on Improved Swin Transformer

摘要: 交互图像着色的目的是在用户提供特定位置的颜色时对灰度图像进行着色。我们使用了改进的Swin Transformer，局部稳定层和像素重组的高效方案。该模型利用Swin Transformer的层次化结构和偏移窗口机制，相比Transformer能够实现在保持计算效率的同时，增强模型对多尺度特征的捕获能力，从而更好地传播和融合用户提供的颜色提示。此外我们改进Swin Transformer从而提高计算效率和模型的可训练性。后面通过局部稳定层去解决像素重组中的局部伪影和边界模糊的问题，最后使用了像素重组取缔解码器来实现高效的图像着色。经实验结果表明，我们的方法优于现有的交互着色方法，可以在用户的提示下产生准确而真实的着色图像。

Abstract: The purpose of interactive image coloring is to color a grayscale image based on the colors specified by the user at specific locations. We utilize an improved Swin Transformer, combined with a local stability layer and an efficient pixel shuffling scheme. This model leverages the hierarchical structure and shifted window mechanism of the Swin Transformer, which enhances the model’s ability to capture multi-scale features while maintaining computational efficiency, compared to the traditional Transformer. This leads to better propagation and integration of user-provided color hints. Additionally, we have optimized the Swin Transformer to improve computational efficiency and model trainability. The local stability layer addresses local artifacts and boundary blurring issues during pixel shuffling, and the pixel shuffling technique replaces the decoder to achieve efficient image colorization. Experimental results show that our method outperforms existing interactive colorization methods, producing accurate and realistic colorized images based on user hints.

1. 引言

从本质上来说，图像着色是一个病态问题，因为自然界物体包含多种多样的颜色，为同一个物体着色可以有多种结果，受这种不确定性影响，灰度图像着色是一个充满挑战的任务。传统着色方法需要投入大量人力，着色结果也很大程度上依赖于个人理解和专业水平。图像上色是一种将灰度图像添加生动颜色的技术，它可以用于给老旧照片上色、重制经典黑白电影以及在绘画作品之间传递颜色等广泛应用。近年来，基于深度学习的方法被引入来解决这一问题，并主要分为三种解决方案：全自动上色[1]-[3]、基于范例的上色[4] [5]和用户提示引导[6]-[8]的上色。依靠深度神经网络及其设计和大规模数据集学习，可以实现完全自动的着色方法，但是全自动着色有个问题是上色结果容易存在颜色饱和度不足、颜色晕染等问题。而基于范例的图像着色是一种使用参考彩色图像来对灰度图像进行着色的方法。但这种方法在很大程度上取决于参考彩色图像的选择，可能会产生不期望的物体轮廓，而且需要找到与被着色图像高度一致的参考图像，有点浪费时间。基于用户交互的图像着色，在最少的用户交互提示下(可以是涂鸦，笔画，像素点等等)对灰度图像进行着色，这样着色效果就可以围绕提示而着色，也不需要寻找参考图像，可以大大增加自主性和艺术性，但是该方法也有痛点，用户提示难以到达其他区域导致着色不理想或者着色容易受到用户提示而导致着色不符合预期，本文是基于用户交互提示的灰度图像着色而展开的。

2. 相关工作

用户交互的图像着色允许用户通过在输入灰度图像的不同点位置指定颜色(即用户提示)来逐个对图像进行着色。由于用于指定空间位置的常用大小范围是从2 × 2到10 × 10像素，因此用户提示仅覆盖整个图像的小部分。故需要一个点交互着色模型来将用户提示传播到整个图像，以便在最小的用户提示下产生合理的结果，然后使用先进优秀的方法在每个区域内传播用户提示的颜色。Yin等人[6]使用手工制作的滤镜来确定用户提示应该填充的区域是一种早期的方法，用于对图像内的简单图案进行着色。Zhang等人[7]提出了一种基于深度学习的方法，通过扩展现有的无条件着色模型来生成给定灰度图像和用户提示的彩色图像。尽管这些方法使用用户提示作为生成彩色图像的条件，将提示传播到较远的相关区域只能在较深层中完成，这使得对较大的语义区域进行着色比对附近的区域更具挑战性。Yun等人[8]提出视觉变换器应用于点交互式着色，视觉变换器的自注意力机制使得模型能够在少量局部提示的情况下，有选择性地为相关区域着色，Transformer的自注意力机制能够在图像的全局范围内捕捉长距离的依赖关系，使得图像中远离颜色提示的区域也能被准确地着色。但是该模型计算开销大和高分辨率着色不稳定。为此我们使用Swin Transformer [9]，Swin Transformer引入了层级式结构和滑动窗口(Shifted Window)机制，将自注意力的计算限制在局部窗口内，大大降低了计算复杂度，同时保留了自注意力的全局建模能力，这样可以使模型在着色灰度图像时更高效。图1是本文提出的灰度图像着色的网络架构。

Figure 1. The overall architecture of the interactive grayscale image colorization based on the improved Swin Transformer

图1. 基于改进SwinTransformer的交互式灰度图像网络架构

3. 方法

3.1. Swin Transformer传递用户提示

首先，我们准备灰度图像$I_g\in R^{H\times W\times 1}$ 和模拟的用户提示$I_{h\mathrm{int}}\in R^{H\times W\times 3}$ 作为训练样本。这里的H和W分别代表图像的高度和宽度，且图像$I_g$ 只具有一个颜色通道(灰度)。灰度图像$I_g$ 通过将大规模数据集中的RGB图像转换到CIELab颜色空间并提取L (明度)通道来获得。用户提示$I_{h\mathrm{int}}$ 包含了非提示区域的a和b通道值，这些值被设置为0，同时添加第三个通道来将提示区域记为1和非提示区域记为0。

在训练过程中，我们通过确定提示的位置和颜色来模拟用户提示。提示位置从图像上的均匀分布中采样，以反映用户可能在图像的任何位置提供提示。一旦确定了提示位置，就通过计算提示区域内每个通道的平均颜色值来获取用户提示的颜色。最终，给定灰度图像$I_g$ 和模拟的用户提示$I_{h\mathrm{int}}$ ，我们通过通道级联⊕获得输入X：

$X=I_g\oplus I_{h\mathrm{int}}$ (1)

其中⊕表示通道级联，得到的输入$X\in R^{H\times W\times 4}$ 然后被重塑为Swin Transformer的图像块序列。

在Swin Transformer中，我们使用层次化的处理方式来逐步细化特征表示。首先，在较低层次，我们使用较大的移位窗口来捕获图像的全局上下文信息。随着层次的提升，我们逐渐减小移位窗口的大小，以获得更细粒度的特征表示。这种层次化的处理方式有助于模型在不同尺度上理解和传播用户提示。

在每个层次上，Swin Transformer通过self-attention机制来计算每个token与其他所有tokens的关系。我们使用如下公式来表示Swin Transformer中的self-attention机制：

$Attention\left(Q,K,V\right)=soft\max \left(\frac{Q{K}^T}{\sqrt{\text{d}}}+B\right)V$ (2)

Swin Transformer的关键优势在于其能够通过移位操作来捕获长距离依赖关系，这使得模型能够更有效地将用户提示传播到图像的远距离相关区域。此外，通过层次化的处理，Swin Transformer能够在保持计算效率的同时，捕获不同尺度的特征，这对于点交互式图像着色来说是至关重要的。

在Swin Transformer的每个层次上，我们都添加了位置编码，以确保模型能够利用空间位置信息。我们使用如下公式来计算位置编码：

${\text{Z}}_0={\text{X}}_{\text{p}}+{\text{E}}_{\text{pos}}$ (3)

其中${\text{E}}_{\text{pos}}$ 是正弦余弦位置编码。

最终，通过Swin Transformer的层次化特征学习和移位自注意力机制，我们能够有效地将用户提示传播到整个图像的相关区域，实现高质量的点交互式图像着色。

3.2. Swin Transformer模型的改进

在本节中，我们改进的模型架构有三个组成部分，分别是Convolution Token Embedding (CTE)、Swin Transformer Block以及Patch Merging三个部分，其中功能如下：

(1) CTE主要负责对输入图像进行分割，切成一小块一小块，然后将其转化成向量矩阵，代表着小块特征。

(2) Swin Transformer Block是利用自注意力机制对每个窗口内的局部特征进行处理。这些窗口内的特征经过层次化处理后，能够捕捉图像中的局部细节和全局上下文信息。随后通过特征的整合和平均池化，产生一个包含丰富信息的全局特征向量。

(3) PatchMerging是一种用于下采样的方法，通过合并相邻的小图像块来形成更大的featuremap，从而增加感野和上下文信息。

我们使用了Convolution Token Embedding (CTE)替代了初始模型的PatchPartition和LinearEmbedding。这样做主要是为了减少参数，提高模型的效率和训练性。Convolution Token Embedding (CTE)是一种新型的TokenEmbedding方法，它利用一维卷积神经网络(CNN)取代了传统的线性嵌入转换方法。它能够将输入数据序列有效地转换为特征丰富的嵌入向量。

Convolution Token Embedding网络结构图如图2，其中CTE模块的原理如下：

(1) 卷积：输入图像$X\in R^{H\times W\times 4}$ 通过两个特殊的Conv 1 × 1，不改变输入的图像的尺寸，但是改变通道数，生成三个不同通道数的特征向量，分别代表“键”(K)、“查询”(Q)和“值”(V)。

(2) 重塑：将K、Q、V维度变成(N, H × W, C)的二维向量，其中N为批量大小，H × W为输入图像的像素数，C为输出通道数。

(3) 生成偏置项：输入图像通过线性层LinearLayer和多层感知机MLP产生一个与K、Q、V维度相同的偏置项β，其目的为了加权K和V输出。

(4) 多头注意力机制处理：首先通过线性投影层处理输入的K、Q、V和β，生成K'、Q'、V'和β'。然后将Q'分解为多个头部，每个头的维度为D model/Num heads，其中D model为向量的维度，Num_heads为头的数量。K'、V'和β'也是相同的处理方法。对于每个头部计算注意力权重，并将权重与对应的V'相乘，得到Num heads个输出张量。将所有头的输出拼接起来，加上偏置量β，得到最终的输出张量。

(5) 重塑：将输出结果重塑与输入图像相同的空间维度，但具有新的通道数，最后生成一个新的特征向量。

Figure 2. Network structure diagram of Convolution Token Embedding

图2. Convolution Token Embedding网络结构图

3.3. 像素重组和局部稳定层

为了将Swin Transformer编码器的输出特征y上采样为全分辨率的彩色图像，我们没有使用传统意义上的解码器[7]，取而代之地采用了Pixel Shuffling (像素重组) [10]，这是一种上采样技术。具体来说，像素洗牌将尺寸为(H/P, W/P, C × P²)的特征图转换为(H, W, C)的尺寸，其中H和W分别是输入图像的高度和宽度，C是通道数，而P是上采样的比例。

尽管像素重组在上采样过程中非常高效，但它在处理较大的上采样比时可能会在图像块的边界处产生可见的伪影。为了解决这个问题，我们引入了局部稳定层[8]，该层限制了模型利用邻近特征生成颜色，从而确保了邻近图像块之间的颜色一致性。通过这种方式，我们显著减少了像素洗牌可能引起的不真实感和细节丢失问题。总结来说，我们的上采样过程可以表示为：

$I_{pred}^{ab}=PS\left(LS\left(y\right)\right)$ (4)

其中$\text{PS}\left(\text{*}\right)$ 代表像素洗牌操作，$LS(\cdot )$ 代表局部稳定层，而$I_{pred}^{ab}\in {\text{R}}^{\text{H}\times \text{W}\times 2}$ 是预测的CIELab颜色空间中的ab颜色通道输出。最终预测的彩色图像$I_{pred}$ 通过将灰度输入$I_{\text{g}}$ (L通道)与$I_{pred}^{ab}$ (ab通道)进行通道级联得到：

$I_{pred}=I_g\oplus I_{pred}^{ab}$ (5)

通过结合像素洗牌和局部稳定层，我们能够在不使用额外解码器的情况下，有效地获得全分辨率的彩色图像，可以实现在用户的提示下有效的着色。

3.4. 目标函数

我们使用CIELab色彩空间中预测图像与原始彩色图像之间的Huber损失[11]来训练我们的模型，来衡量预测图像与原始彩色图像之间的差异：

$L_{recon}=\frac{1}{2}{\left(I_{pred}-I_{GT}\right)}^2\cdot 1_{\left|I_{pred}-I_{GT}\right|<1}+\left|I_{pred}-I_{GT}\right|-\frac{1}{2}\cdot 1_{\left|I_{pred}-I_{GT}\right|\ge 1}$ (6)

4. 实验

4.1. 实验参数

本文的算法主要是依托深度学习主流开源框架PyTorch，Python版本Python 3.7，实验设置的迭代次数为500，初始学习率设置为0.01，学习率衰减类型[12]选择的是CosineAnnealingLrUpdater，它根据余弦函数的形状递减降低学习率[13]。这种学习率下降策略可以帮助模型在训练后期更稳定地收集，并且有助于跳出局部的最佳点。在模型训练过程中使用的优化算法是AdamW [14]，这是一种常用的优化器，结合了Adam优化计算和权衡技术，可以有效地更新模型参数。

我们使用COCO2017训练图像着色，它由163,957张图像组成。我们在训练过程中不使用分类标签，因为我们的模型是以自监督的方式训练的，然后和现有的交互式灰度图像着色模型进行了比较。

4.2. 实验结果

本节展示了我们方法的着色效果，并与现有方法进行了比较。由于是自然图像的灰度图像进行着色，而每个物体都具有复杂颜色分布，但是大部分物体都有一个统一的基调颜色，比如天空，大海是蓝色。所以引入峰值信噪比(Peak Signal to Noise Ratio, PSNR)和学习感知图像块相似度(Learned Perceptual Image Patch Similarity, LPIPS)能一定程度反映着色算法的性能。我们的着色图像数据实验数据如表1和表2，其中@后面的数字是用户提示的数量个数。

Table 1. Comparison table of PSNR data for different algorithm models in image coloring

表1. 不同算法模型着色PSNR数据对比表

	PSNR @1	PSNR @5	PSNR @10	PSNR @50	PSNR @100
Yin et al.	21.963	22.579	23.020	24.634	25.744
Zhang et al.	25.590	26.826	27.558	29.309	30.001
Yun et al.	26.100	28.111	29.094	31.265	32.097
Ours	26.324	28.585	29.664	31.830	33.138

Table 2. Comparison table of LPIPS data for different algorithm models in image coloring

表2. 不同算法模型着色LPIPS数据对比表

	LPIPS @1	LPIPS @5	LPIPS @10	LPIPS @50	LPIPS @100
Yin et al.	0.164	0.151	0.141	0.119	0.104
Zhang et al.	0.100	0.087	0.081	0.068	0.064
Yun et al.	0.088	0.071	0.066	0.052	0.047
Ours	0.085	0.068	0.061	0.048	0.043

4.4. 实验分析

4.4.1. 定量分析

采用上文所提出的评价指标从客观的角度对以上模型进行比较，结果如表1和表2所示，表中计算的结果是对COCO的测试集前4000张图片进行处理，然后计算各个指标，最后得到的平均值。从表中结果可以看出，本文算法具有比其他五种算法更高的PSNR值和更低的LPIPS值，例如在用户提示数量100的时候，我们的PSNR比Yun等提升了3.24%，同时LPIPS降低8.51%等，证明本文提出的着色算法模型优于其他模型，可以生成更为真实，质量更高的彩色图片。

4.4.2. 定性分析

图3展示了四组对比结果。从图中可以看出，本文提出的着色算法在整体上主观效果更好，以下将从着色准确度和细节方面进行分析。在第一行和第三行，Yin等模型虽然大致预测到主体的颜色，但是主体周围的图像而会错误的着色，而且整体饱和度偏低和对比度偏高；在第二行和第四行，可以发现Zhang等模型饱和度有部分提升，但是云朵和车下面的路也会受到其他的影响，存在部分区域着色不理想；而Yun等模型在预测正确颜色的基础上还兼顾了整体性，第三行和第四行对比发现对比度可能略低；最后我们的模型相比Yun等模型饱和度和对比度偏向正常，比较接近真实的自然彩色照片。

输入 用户提示 Yin等 Zhang等 Yun等 Ours 原始图像

Figure 3. Comparison of coloring results for different algorithm models

图3. 不同算法模型图像着色对比图

5. 结论

本文提出基于改进的Swin Transformer的交互式灰度图像着色方法，该方法使用了改进的Swin Transformer编码器、局部稳定层和像素重组，通过定性和定量结果分析，当给定用户的特定提示时，可以生成真实且多彩的图像，其着色结果优于现有的方法。