Image Coloring Method Based on Transformer and Generative Adversarial Network

Abstract

The present invention discloses an image coloring method based on Transformer and generative adversarial network, which uses the generative adversarial network GAN and Transformer instead of simply using the convolutional neural network CNN to solve the image coloring problem. The proposed Transformer-GAN reduces excessive computing resources through a window-based multi-head self-attention mechanism and a discriminator friendly to computing resources. The local enhancement forward propagation network and skip connections ensure that shallow features can be effectively transmitted and utilized in the network, enabling Transformer-GAN to effectively capture the correlation between global and local information. The best training process is also explored through data augmentation and objective function selection. The formed color image generator and discriminator enable Transformer-GAN to perform well in image colorization, achieving the best visual effect.

Description

Technical Field

[0001] The present invention belongs to the technical field of image processing, and relates to an image coloring method based on Transformer and generative adversarial network. Background Art

[0002] In the image coloring task, our goal is to generate a color image from the input grayscale image. By category, from the early traditional CNN structure-based algorithms without skip connections to the later networks that allow users to specify image colors (these networks require users to input color values in specific layers). And the end-to-end feedforward architecture for colorizing animated images using generative adversarial networks (GANs), as well as infrared colorization, radar image colorization, etc. for specific fields, and the later multi-modal coloring models (text-based coloring networks). Diverse coloring networks compensate for the lack of diversity by generating different color images. The network architecture includes multi-path networks that learn different features in different network paths or levels, and the user provides a reference image as an input sample for the coloring network. All of the above models have one thing in common, that is, they are all networks based on convolutional neural network CNN. However, different from previous work, I use Transformer and generative adversarial network GAN to construct an image colorization network. As far as I know, this is the first study using Transformer as the main network for image colorization. Summary of the Invention

[0003] The purpose of the present invention is to provide an image coloring method based on Transformer and generative adversarial network, which solves the problems of poor coloring effect and poor coloring diversity of the current image coloring network.

[0004] The technical solution adopted by the present invention is an image coloring method based on Transformer and generative adversarial network, and the method is implemented according to the following steps:

[0005] Step 1, construct an image coloring model based on generative adversarial network, the image coloring model includes a color image generator and a discriminator; the color image generator is used to generate a color image, and the discriminator is used to judge whether the input image is a real color image or a pseudo-color image;

[0006] Step 2, input the gray image into the color image generator of the image coloring model to generate a pseudo-color image;

[0007] Step 3, update the parameters of the discriminator and the color image generator respectively:

[0008] Step 3.1: First, fix the parameters of the color image generator. Input the pseudo-color image and the real color image corresponding to the gray image into the discriminator alternately. Then, calculate the loss between the real color image corresponding to the gray image and the label value of 1 according to the loss function, and calculate the loss between the pseudo-color image generated from the gray image and the label value of 0 according to the loss function. Finally, use the backpropagation algorithm to update the parameters of the discriminator; where the label value of 1 represents a real image, and the label value of 0 represents a generated pseudo-color image.

[0009] Step 3.2: Fix the parameters of the discriminator. Calculate the loss between the generated pseudo-color image and the label value of 1 according to the loss function. Finally, use the backpropagation algorithm to update the parameters of the color image generator.

[0010] Step 3.3: Continuously loop the process of updating the parameters of the discriminator and the color image generator in Step 3.1 and Step 3.2 until the loss value converges and the color image generator generates a good pseudo-color image, that is, an optimized image coloring model is obtained.

[0011] Step 4: Use the optimized image coloring model to directly color the gray image.

[0012] The features of the present invention also lie in that

[0013] In Step 1, the color image generator includes multiple MWin-transformer modules. The function of the Mwin-transformer module is to extract and reconstruct the features of the image and output a 3-channel effective color image:

[0014] The Mwin-transformer module consists of three core parts: window-based multi-head self-attention mechanism, layer normalization operation LN, and local enhancement forward propagation network LeFF.

[0015] The process of the color image generator generating a pseudo-color image is as follows:

[0016] X′ = Embedded Tokens(X in )

[0017] X″ = W-MSA(LN(X)) + X′

[0018] X out = LeFF(LN(X″)) + X″

[0019] Where, X in represents the input, which is a gray image or a pseudo-color image;

[0020] Embedding Tokens means to embed Xin Convert to a vector;

[0021] X' represents the vector output obtained by inputting X in into Embedding Tokens;

[0022] Then, the result LN(X') after layer normalization of the vector X' is input into the window-based multi-head self-attention mechanism W-MSA to obtain a vector with feature information extracted, and then added to X' to obtain a vector X'' that aggregates more feature information; X'' represents the output obtained by inputting X' into the window-based multi-head self-attention mechanism and the layer normalization operation;

[0023] Continue to perform layer normalization on the vector X'', input the normalized LN(X'') into the local enhanced forward propagation network to obtain a vector with more local feature information extracted, and then add it to X'' to obtain a vector X that aggregates more local feature information out , X out represents the output obtained by inputting X'' into the local enhanced forward propagation network LeFF and the layer normalization operation.

[0024] The layer normalization LN operation is to solve the problem of internal covariate shift, and the calculation process of the layer normalization operation is as follows:

[0025]

[0026] Among them, the object of action of the LN layer is X represents a vector, μ and δ respectively represent the mean and variance of each sample, and are affine learning parameters, d k is the hidden dimension, indicating that this number is a k-dimensional vector.

[0027] The window-based multi-head self-attention mechanism is as follows:

[0028] The pseudo-color image is divided into multiple windows, and then self-attention calculations are performed in these different windows. Since the number of patches in a window is much smaller than the total number of patches in an image, and the number of windows remains unchanged, the computational complexity of the window-based multi-head self-attention mechanism changes from a square relationship with the image size to a linear relationship, greatly reducing the computational complexity of the model.

[0029] Convolution is added to the forward propagation network in the Mwin-transformer module to form the local enhanced forward propagation network LeFF.

[0030] The loss function is:

[0031]

[0032] Among them,

[0033]

[0034]

[0035] Among them, G * represents the sum of loss functions, represents the loss of the conditional generative adversarial network, represents the Charbonnier loss, and λ represents the weight coefficient of the Charbonnier loss;

[0036] x represents the input grayscale image;

[0037] y represents the real color image corresponding to the input grayscale image;

[0038] log represents the logarithmic function with base 2;

[0039] represents that the independent variables are x and y;

[0040] represents that the independent variable is x;

[0041] ε represents a constant coefficient with a value of 10 -3 ;

[0042] |||| represents taking the absolute value.

[0043] The beneficial effects of the present invention are as follows: The present invention is a method for image coloring based on Transformer and generative adversarial network, which can color grayscale images. Throughout the invention, the Transformer is beneficial for the present invention to capture the global features of the image, and the forward propagation network (LeFF) included in the invention is beneficial for the present invention to capture the local features of the image. The generative adversarial network is beneficial for better training the entire image coloring model. This image coloring method has a good coloring effect on grayscale images both in terms of details and overall, and the coloring method is applicable to grayscale images of any size, having high versatility. BRIEF DESCRIPTION OF THE DRAWINGS

[0044] Figure 1 is the structural diagram of the image coloring model of the present invention;

[0045] Figure 2 is the structural diagram of the color image generator G;

[0046] FIG. 3(a) is the structural diagram of the discriminator D of the present invention;

[0047] Figure 3(b) is the structural diagram of the MWin-transformer of the present invention;

[0048] Figure 3(c) is the structural diagram of the local enhanced forward propagation network LeFF of the present invention. Specific implementation manner

[0049] The image coloring method based on Transformer and generative adversarial network of the present invention is implemented according to the following steps:

[0050] Step 1, construct an image coloring model based on a generative adversarial network, where the image coloring model includes a color image generator and a discriminator; the color image generator is used to generate a color image, and the discriminator is used to judge whether the input image is a real color image or a pseudo-color image;

[0051] Step 2, input a gray image into the color image generator of the image coloring model to generate a pseudo-color image;

[0052] Step 3, update the parameters of the discriminator and the color image generator respectively:

[0053] Step 3.1: First, fix the parameters of the color image generator, alternately input the pseudo-color image and the real color image corresponding to the gray image into the discriminator in sequence, then calculate the loss between the real color image corresponding to the gray image and the label value of 1 according to the loss function, and calculate the loss between the pseudo-color image generated from the gray image and the label value of 0 according to the loss function. Finally, use the backpropagation algorithm to update the parameters of the discriminator; where the label value of 1 represents a real image, and the label value of 0 represents a generated pseudo-color image;

[0054] Step 3.2: Fix the parameters of the discriminator, calculate the loss between the generated pseudo-color image and the label value of 1 according to the loss function, and finally use the backpropagation algorithm to update the parameters of the color image generator.

[0055] Step 3.3: Continuously loop the process of updating the parameters of the discriminator and the color image generator in steps 3.1 and 3.2 until the loss value converges and the color image generator generates a good pseudo-color image, that is, an optimized image coloring model is obtained;

[0056] Step 4, directly color the gray image using the optimized image coloring model.

[0057] As Figure 1 shown, construct an image coloring model based on Transformer and generative adversarial network GAN, where G and D represent the color image generator and the discriminator respectively. Specifically, let the gray image x ∈ R 3×H×WAs the input to the color image generator G, a pseudo-color image G(x) is generated, and then the pseudo-color image G(x) and the real color image y are alternately input into the discriminator.

[0058] First, fix the parameters of the color image generator. Input the pseudo-color image and the real color image corresponding to the grayscale image into the discriminator alternately. Then, calculate the loss between the real color image corresponding to the grayscale image and the label value of 1 according to the loss function, and calculate the loss between the pseudo-color image generated from the grayscale image and the label value of 0 according to the loss function. Finally, use the backpropagation algorithm to update the parameters of the discriminator; where the label value of 1 represents the real image, and the label value of 0 represents the generated pseudo-color image.

[0059] Fix the parameters of the discriminator. Calculate the loss between the generated pseudo-color image and the label value of 1 according to the loss function. Finally, use the backpropagation algorithm to update the parameters of the color image generator. Continuously loop the process of updating the parameters of the discriminator and the color image generator in steps 3.1 and 3.2 until the loss value converges, and the color image generator generates a good pseudo-color image, that is, an optimized image coloring model is obtained. Use the optimized image coloring model to directly color the grayscale image.

[0060] The design of this method mainly includes two key points, namely the design of the color image generator and the discriminator and the design of their components. Their detailed structures are introduced one by one below.

[0061] ①. Design of the color image generator G:

[0062] The input and output of image colorization is a mapping relationship, and the depth conversion process of these components should be a symmetric relationship. Based on this idea, as Figure 2As shown in the figure, the entire color image generator is designed in a U shape. Generally speaking, in the encoder stage, the gray image x first undergoes input dimension adjustment, which is a convolutional layer with a 3×3 convolutional kernel, and the activation function LeakyReLU is used to adjust the input dimension and extract low-level features. Then, after the designed window transformer module MWin-transformer, it reaches the downsampling layer composed of a convolutional layer with a 4×4 convolutional kernel and a stride of 2, and this step is repeated 2 times. Then, the image passes through the MWin-transformer module as the bottleneck stage. The decoder corresponds to the encoder, and their designs are completely symmetric and consistent: First, it passes through the window MWin-transformer module and undergoes upsampling, where the upsampling operation is a transposed convolution with a 2×2 convolutional kernel and a stride of 2. To maintain the symmetry of the network, this step is repeated 2 times as in the encoder. Finally, the output dimension adjustment composed of 3×3 convolution is used to adjust the output dimension to ensure that the output is a 3-channel valid color image.

[0063] ②. Design of the MWin-transformer module

[0064] In Figure 3(b), the constructed MWin-transformer module consists of three core parts: the W-MHSA mechanism, the layer normalization operation LN, and the local enhancement forward propagation network LeFF network;

[0065] Among them, the layer normalization operation LN is as follows:

[0066]

[0067] The LN layer is an important guarantee for the fast training and stable convergence of the image coloring model. The object of action of the LN layer is and represent the mean and variance of each sample respectively, and are affine learning parameters, and d k is the hidden dimension. The calculation process of the MWin-transformer module is as follows:

[0068] X′ = Embedded Tokens(X in )

[0069] X″ = W-MSA(LN(X)) + X′ [[ID=३६]]

[0070] X out = LeFF(LN(X″)) + X ″

[0071] Among them, X inRepresents the input, which is a grayscale image or a pseudo-color image;

[0072] Embedding Tokens means converting X in into a vector;

[0073] X′ represents the vector output obtained by inputting X in into the Embedding Tokens;

[0074] Then, the result LN(X′) after layer normalization of the vector X′ is input into the window-based multi-head self-attention mechanism W-MSA to obtain a vector with feature information extracted, and then added to X′ to obtain a vector X″ that aggregates more feature information; X″ represents the output obtained by inputting X′ into the window-based multi-head self-attention mechanism and layer normalization operation;

[0075] Continue to perform layer normalization on the vector X″, input the normalized LN(X″) into the local enhanced forward propagation network to obtain a vector with more local feature information extracted, and then add it to X″ to obtain a vector X out , X out represents the output obtained by inputting X″ into the local enhanced forward propagation network and layer normalization operation.

[0076] ③ Local Enhanced Forward Propagation Network (LeFF)

[0077] To enhance the ability of the image coloring model to capture local features, we add a convolution to the forward propagation network to form the local enhanced forward propagation network (LeFF). The specific design is shown in Figure 3(c): First, the input sequence, and then through the sequence-to-image module to convert the sequence into an image. We use a convolution with a kernel size of 1*1 on the image, then activate it through an activation function, then perform a convolution with a kernel size of 3*3 and a convolution with a kernel size of 1*1, and then activate it through an activation function. Finally, convert the image back into a sequence to complete the local enhanced forward propagation network.

[0078] ④ Design of the Discriminator D

[0079] The essence of the discriminator is to judge whether the given data is "real", that is, to judge whether it is real training data or fake data generated by the color image generator G. As shown in Figure 3(a), we first input real color images or pseudo-color images, and flatten them into small patches through linear flattening, which is composed of convolutional layers. Then, stack 4 MWin Transformer blocks with the same structure as G. Finally, output true or false through a linear layer to achieve the discrimination function.

[0080] Example 1

[0081] To prove the effectiveness of the image coloring model of the present invention, experiments were conducted on animal face images and landscape images respectively. We compared it with other currently popular models: the research of Yoo et al. (Yoo, S., Bahng, H., Chung, S., Lee, J., Chang, J., Choo, J.: Coloring with limited data: Few-shot colorization via memory augmented networks. In: Proceedings of the IEEE / CVF Conference on Computer Vision and Pattern Recognition, pp. 11283–11292 (2019)), the research of Su et al. (Su, J.-W., Chu, H.-K., Huang, J.-B.: Instance-aware image colorization. In: Proceedings of the IEEE / CVF Conference on Computer Vision and Pattern Recognition, pp. 7968–7977 (2020)). We compared on three metrics respectively: Frechet Inception distance score, peak signal-to-noise ratio, and structural similarity. On the tests of these two datasets, all three metrics were improved, with the lowest improvements being 0.003, 0.263, and 0.014 respectively. We found that the more complex the scene, the more realistic the details of the model coloring. We hypothesize that this is because the transformer can learn the distribution law of a large amount of data better than a CNN. In addition, we observed that the overall coloring effect of our model is smoother and more uniform than other methods, and there are no excessive color mutations, which proves that the transformer can better capture the global information of the image.

Claims

1. An image coloring method based on Transformer and generative adversarial network, characterized in that, The method is implemented according to the following steps: Step 1: Construct an image coloring model based on a generative adversarial network. The image coloring model includes a color image generator and a discriminator. The color image generator is used to generate color images, and the discriminator is used to determine whether the input image is a real color image or a pseudo-color image. In Step 1, the color image generator contains multiple MWin-transformer modules. The function of the MWin-transformer module is to extract and reconstruct the features of the image and output a 3-channel valid color image. The MWin-transformer module consists of three core parts: a window-based multi-head self-attention mechanism, a layer normalization operation LN, and a local enhancement forward propagation network LeFF. Step 2: Input the grayscale image into the color image generator of the image coloring model to generate a pseudo-color image. The process of the color image generator generating a pseudo-color image is as follows: X′ = Embedded Tokens(X in ) X″ = W-MSA(LN(X′)) + X′ X out = LeFF(LN(X″)) + X″ wherein, X in represents an input, which is a gray image or a pseudo-color image; Embedding Tokens means converting X in into a vector; X' represents the vector output obtained by inputting X in into Embedding Tokens; Then, the result LN(X′) after layer-normalizing the vector X′ is input into the window-based multi-head self-attention mechanism W-MSA to obtain a vector with feature information extracted. Then, it is added to X′ to obtain a vector X″ that aggregates more feature information. X″ represents the output obtained by inputting X′ into the window-based multi-head self-attention mechanism and the layer normalization operation. Continue to perform layer normalization on the vector X″, input the normalized LN(X″) into the local enhancement forward propagation network to obtain a vector with more local feature information extracted, and then add it to X″ to obtain a vector X that aggregates more local feature information. out , X out denotes the output obtained by inputting X″ into the local enhancement forward propagation network LeFF and performing layer normalization operations; Step 3: Update the parameters of the discriminator and the color image generator respectively: Step 3.1: First, fix the parameters of the color image generator. Input the pseudo-color image and the real color image corresponding to the grayscale image into the discriminator alternately. Then, calculate the loss between the real color image corresponding to the grayscale image and the label value of 1 according to the loss function, and calculate the loss between the pseudo-color image generated from the grayscale image and the label value of 0 according to the loss function. Finally, use the backpropagation algorithm to update the parameters of the discriminator. Among them, the label value of 1 represents a real image, and the label value of 0 represents a generated pseudo-color image. Step 3.2: Fix the parameters of the discriminator. Calculate the loss between the generated pseudo-color image and the label value of 1 according to the loss function. Finally, use the backpropagation algorithm to update the parameters of the color image generator. Step 3.3: Continuously loop the process of updating the parameters of the discriminator and the color image generator in Step 3.1 and Step 3.2 until the loss value converges and the color image generator generates a good pseudo-color image, that is, an optimized image coloring model is obtained. Step 4: Directly color the grayscale image using the optimized image coloring model.

2. The image coloring method based on Transformer and generative adversarial network according to claim 1, characterized in that, The calculation process of the layer normalization operation is as follows: Among them, the object of action of the LN layer is X represents a vector, and μ and δ represent the mean and variance of each sample respectively. and are affine learning parameters, and d k is the hidden dimension. indicates that this number is a k-dimensional vector.

3. The image coloring method based on Transformer and generative adversarial network according to claim 1, characterized in that The window-based multi-head self-attention mechanism is as follows: The pseudo-color image is divided into multiple windows, and then self-attention calculations are performed in these different windows. Since the number of patches in one window is much smaller than the total number of patches in an image, and the number of windows remains unchanged, the computational complexity of the window-based multi-head self-attention mechanism changes from a quadratic relationship with the image size to a linear relationship, greatly reducing the computational complexity of the model.

4. The image coloring method based on Transformer and generative adversarial network according to claim 1, wherein Convolution is added to the forward propagation network in the Mwin-transformer module to form a locally enhanced forward propagation network, LeFF.

5. The image coloring method based on Transformer and generative adversarial network according to claim 1, wherein The loss function is as follows: Where, Among them, G * represents the sum of loss functions, represents the loss of the conditional generative adversarial network, represents the Charbonnier loss, and λ represents the weight coefficient of the Charbonnier loss; x represents the input grayscale image; y represents the true color image corresponding to the input grayscale image; log represents the logarithmic function with base 2; Indicating that the independent variables are x and y; It is indicated that the independent variable is x; ε represents a constant coefficient with a value of 10 -3 ; || || represents taking the absolute value.

Images