CN115546338A - Image Colorization Method Based on Transformer and Generative Adversarial Network - Google Patents

Abstract

The invention discloses a method for coloring images based on a Transformer and a generation countermeasure network, which solves the problem of coloring images by using the generation countermeasure network GAN and the Transformer instead of simply using a convolutional neural network CNN. The local enhanced forward propagation network and the hopping connection ensure that shallow features can be efficiently transmitted and utilized in the network so that the transform-GAN can efficiently capture the correlation between global and local information. An optimal training process is also explored through data enhancement and objective function selection, and a color image generator and a discriminator are formed to enable the Transformer-GAN to perform well in the aspect of image colorization. The best visual effect is achieved.

Description

Image Colorization Method Based on Transformer and Generative Adversarial Network

technical field

The invention belongs to the technical field of image processing, and relates to an image coloring method based on Transformer and generation confrontation network.

Background technique

In the image colorization task, our goal is to generate a color image from an input grayscale image. According to the category, from the early traditional non-skip connection algorithm based on the CNN structure, to the network in which the image color is specified by the user (these networks require the user to input the color value in a specific layer). and an end-to-end feed-forward architecture for animated image colorization using Generative Adversarial Networks (GAN), as well as domain-specific infrared colorization, radar image colorization, etc., and later multimodal coloring models (text-based shading network). The diverse colorization network compensates for the lack of diversity by generating different color images, the network architecture contains a multi-path network that learns different features in different network paths or levels, and the user gives a reference image as the colorization network input sample. All the above models have one thing in common that they are all based on the convolutional neural network CNN. However, unlike the previous work, I use the transformer and the generation confrontation network GAN to build the image colorization network. As far as I know, this is the first A study of image colorization using Transformer for major networks.

Contents of the invention

The purpose of the present invention is to provide an image coloring method based on Transformer and generating confrontation network, which solves the problems of poor coloring effect and poor coloring diversity of current image coloring network.

The technical solution adopted in the present invention is, based on Transformer and generating an image coloring method against a network, the method is implemented according to the following steps:

Step 1, constructing an image coloring model based on generating an adversarial network, the image coloring model includes a color image generator and a discriminator; the color image generator is the same as generating a color image, and the discriminator is used to judge whether the input image is real Color images or pseudo-color images;

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

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, alternately input the false color image and the real color image corresponding to the gray image into the discriminator, and then calculate the real color image and label corresponding to the gray image according to the loss function The loss between the value 1, and the loss between the pseudo-color image generated by the gray image and the label value 0 according to the loss function calculation, and finally use the backpropagation algorithm to update the parameters of the discriminator; wherein the label value A value of 1 represents the real image, and a label value of 0 represents the 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, and finally update the parameters of the color image generator by using the backpropagation algorithm.

Step 3.3: Repeat step 3.1 and step 3.2 to update the discriminator and color image generator parameters until the loss value converges, and the color image generator generates a pseudo-color image with good effect, that is, the optimized image coloring model is obtained;

Step 4, using the optimized image coloring model to directly color the gray image.

The present invention is also characterized in that,

In step 1, multiple MWin-transformer modules are included in the color image generator, and the function of the Mwin-transformer module is to extract and reconstruct image features, and output 3 channel effective color images:

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 enhanced forward propagation network LeFF.

The flow of the color image generator to generate a pseudo-color image is as follows:

X'＝Embedded Tokens(X _in )

X"=W-MSA(LN(X))+X'

X _out ＝LeFF(LN(X″))+X″

Among them, X _in represents the 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 layer normalized result LN(X') of the vector X' is input into the window-based multi-head self-attention mechanism W-MSA to obtain a vector with extracted feature information, and then added to X' to obtain more The vector X" of feature information; X" represents the output obtained by inputting X' into the window-based multi-head self-attention mechanism and layer normalization operation;

Continue to normalize the vector X″, and input the normalized LN(X″) into the local enhanced forward propagation network to obtain a vector that extracts more local feature information, and then add it to X″ to get a converged The vector X _out of more local feature information, X _out represents the output obtained by inputting X″ into the local enhanced forward propagation network LeFF and the layer normalization operation.

The layer normalization LN operation is to solve the internal covariate offset problem, and the calculation process of the layer normalization operation is:

X代表向量，μ以及δ分别代表每个样本的均值和方差，

和

为仿射学习参数，d_k是隐藏维度，

表示该数是一个k维的向量。Among them, the object of the LN layer is

X represents the vector, μ and δ represent the mean and variance of each sample, respectively,

and

is the affine learning parameter, d _k is the hidden dimension,

Indicates that the number is a k-dimensional vector.

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

Divide the pseudo-color image into multiple windows, and then perform self-attention calculations in these different windows. Since the number of patches in a window is much smaller than the number of all small blocks in a picture, and the number of windows remains the same, so The computational complexity of the window-based multi-head self-attention mechanism and the image size have changed from a square relationship to a linear relationship, which greatly reduces the computational complexity of the model.

Convolutions are added to the forward propagation network in the Mwin-transformer module, resulting in a locally enhanced forward propagation network LeFF.

The loss function is:

in,

表示条件生成对抗网络损失，

表示Charbonnier损失，λ表示Charbonnier损失的权重系数；where G ^* represents the sum of loss functions,

Denotes the conditional generative adversarial network loss,

Represents Charbonnier loss, λ represents the weight coefficient of Charbonnier loss;

x represents the input gray image;

y represents the real color image corresponding to the input gray image;

log represents a logarithmic function with base 2;

表示自变量为x,y；

Indicates that the independent variable is x, y;

表示自变量为x；

Indicates that the independent variable is x;

ε represents a constant coefficient with a value of 10 ^-3 ;

|||| means to find the absolute value.

The beneficial effects of the present invention are: the present invention is a method for image coloring based on Transformer and generation confrontation network, which can color gray images. In the whole invention, Transformer is beneficial to the present invention to capture the global features of the image, and the invention includes The forward propagation network (LeFF) of the present invention is conducive to capturing the local features of the image, and the generative confrontation network is conducive to better training the entire image coloring model. The image coloring method has a good coloring effect on the gray image both in details and as a whole, and the coloring method is suitable for gray images of any size, which has high versatility.

Description of drawings

Fig. 1 is a structural diagram of the image coloring model of the present invention;

Fig. 2 is the structural diagram of color image generator G;

Fig. 3 (a) is the structural diagram of discriminator D of the present invention;

Fig. 3 (b) is the structural diagram of MWin-transformer of the present invention;

Fig. 3(c) is a structural diagram of the local enhanced forward propagation network LeFF of the present invention.

detailed description

The image coloring method based on Transformer and generation confrontation network of the present invention, the method is implemented according to the following steps:

Step 1, constructing an image coloring model based on generating a confrontation network, the image coloring model includes a color image generator and a discriminator; the color image generator is the same as generating a color image, and the discriminator is used to judge whether the input image is real Color images or pseudo-color images;

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

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, alternately input the false color image and the real color image corresponding to the gray image into the discriminator, and then calculate the real color image and label corresponding to the gray image according to the loss function The loss between the value 1, and the loss between the pseudo-color image generated by the gray image and the label value 0 according to the loss function calculation, and finally use the backpropagation algorithm to update the parameters of the discriminator; wherein the label value A value of 1 represents the real image, and a label value of 0 represents the 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, and finally update the parameters of the color image generator by using the backpropagation algorithm.

Step 3.3: Repeat step 3.1 and step 3.2 to update the discriminator and color image generator parameters until the loss value converges, and the color image generator generates a pseudo-color image with good effect, that is, the optimized image coloring model is obtained;

Step 4, using the optimized image coloring model to directly color the gray image.

As shown in Figure 1, an image coloring model based on Transformer and GAN is constructed. G and D represent the color image generator and discriminator respectively. Specifically, the gray image x∈R ^3×H×W is used as input into Color image The color image generator G generates a pseudo-color image G(x), and then alternately inputs the pseudo-color image G(x) and the real color image y to the discriminator.

First, fix the parameters of the color image generator, alternately input the false color image and the real color image corresponding to the gray image into the discriminator, and then calculate the true color image corresponding to the gray image and the label value as 1 according to the loss function The loss between, and according to the loss function to calculate the loss between the pseudo-color image generated by the gray image and the label value 0, and finally use the backpropagation algorithm to update the parameters of the discriminator; where the label value 1 represents is the real image, and the label value of 0 represents the generated pseudo-color image.

The parameters of the discriminator are fixed, and the loss between the generated pseudo-color image and the label value of 1 is calculated according to the loss function, and finally the parameters of the color image generator are updated using the back propagation algorithm. Continuously repeat steps 3.1 and 3.2 to update the discriminator and color image generator parameters until the loss value converges, and the color image generator generates a pseudo-color image with good effect, that is, the optimized image coloring model is obtained. The gray image is directly colored by using the optimized image coloring model.

The design of this method mainly includes two main points, namely the design of the color image generator and discriminator and the design of their components, and their detailed constructions are introduced one by one below.

①. Design of color image generator G:

The input and output of image colorization is a mapping relationship, and the depth conversion process of these components should be a symmetrical relationship. Based on this idea, as shown in Figure 2, the entire color image generator is designed as a U-shape; in general Say, in the encoder stage, the gray image x first undergoes input dimension adjustment, which is a convolutional layer with a convolution kernel of 3×3, and the activation function LeakyReLU is used to adjust the input dimension and extract low-level features. Then, after the designed window-based transformer module MWin-transformer, it reaches the downsampling layer consisting of a 4×4 convolutional layer with a stride of 2, and this step is repeated twice. Then, the image passes through the MWin-transformer module as the bottleneck level. The decoder and the encoder correspond to each other, and their design is completely symmetrical and consistent: first, it passes through the window MWin-transformer module and performs upsampling, where the upsampling operation is a convolution kernel with a stride of 2 and a 2×2 Transposed convolution. To maintain the symmetry of the network, this step is repeated 2 times as in the encoder. Finally, the output dimensionality is adjusted using an output dimensionality adjustment consisting of 3×3 convolutions to ensure that the output is a 3-channel effective color image.

②. MWin-transformer module design

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

The layer normalization operation LN is as follows:

以及

分别代表每个样本的均值和方差，

和

为仿射学习参数，d_k是隐藏维度。MWin-transformer模块的计算流程如下所示：The LN layer is an important guarantee for the rapid training and stable convergence of the image coloring model. The role of the LN layer is

as well as

represent the mean and variance of each sample, respectively,

and

is the affine learning parameter, and d _k is the hidden dimension. The calculation process of the MWin-transformer module is as follows:

X'＝Embedded Tokens(X _in )

X"=W-MSA(LN(X))+X'

X _out =LeFF(LN(X″))+X ^″

Among them, X _in represents the input, which is a gray image or a pseudo-color image;

Embedding Tokens Embedding Tokens means converting X _in into a vector;

X' represents the vector output obtained by inputting X _in into Embedding Tokens;

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

Continue to normalize the vector X″, and input the normalized LN(X″) into the local enhanced forward propagation network to obtain a vector that extracts more local feature information, and then add it to X″ to get a converged The vector X _out of more local feature information, X _out represents the output obtained by inputting X″ into the local enhancement forward propagation network and layer normalization operation.

③Local Enhanced Forward Propagation Network (LeFF)

In order to enhance the ability of the image coloring model to capture local features, we add convolution to the forward propagation network to form a local enhanced forward propagation network (LeFF). The specific design is shown in Figure 3(c): first, the input sequence , and then through the sequence into an image module to turn the sequence into an image, we use a convolution with a convolution kernel of 1*1 for the image, then activate it through an activation function, and then pass through a convolution with a convolution kernel of 3*3 and convolution The kernel is a 1*1 convolution, and then activated by the activation function, and finally the image is converted into a sequence to complete the local enhanced forward propagation network.

④.Design of discriminator D

The essence of the discriminator is to judge whether the given data is "real", that is, to judge whether it is the real training data or the fake data generated by the color image generator G. As shown in Figure 3(a), we first input the real color image or Pseudo-color images, and expand them into patches by linear flattening, linear flattening consists of convolutional layers, and then stack 4 MWin Transformer blocks with the same structure as G, and finally output true or False to achieve the function of discrimination.

Example 1

In order to prove the effectiveness of the image coloring model of the present invention, experiments were carried out on animal face images and landscape images, and we compared them with other currently popular models: Research by Yoo et al. (Yoo, S., Bahng, H .,Chung,S.,Lee,J.,Chang,J.,Choo,J.:Coloring with limited data:Few-shotcolorization via memory augmented networks.In:Proceedings of the IEEE/CVFConference on Computer Vision and Pattern Recognition, pp.11283–11292(2019)), research by Su et al. (Su, J.-W., Chu, H.-K., Huang, J.-B.:Instance-aware imagecolorization.In:Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.7968–7977(2020)), we compared three indicators: Frechet Inception distance score, peak signal-to-noise ratio, and structural similarity. In these two data In the test set, these three indicators have all improved, and the lowest increases are 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 is better than a CNN to learn the distribution of large amounts of data. In addition, we observe that the overall colorization effect of our model is smoother and more uniform than other methods, without excessive color mutations, which proves that the transformer can better capture the global information of the image.

Claims (7)

1. Image rendering method based on Transformer and generative confrontation network, characterized in that it is implemented according to the following steps,

step 1, constructing an image coloring model based on a generation countermeasure network, wherein the image coloring model comprises a color image generator and a discriminator; the color image generator is used for generating a color image, and the discriminator is used for judging whether the input image is a real color image or a false color image;

step 2, inputting the gray image into a color image generator of the image coloring model to generate a pseudo color image;

and 3, respectively updating parameters of the discriminator and the color image generator:

step 3.1: firstly, fixing parameters of a color image generator, inputting the false color image and the real color image corresponding to the gray image into an identifier in sequence, then calculating the loss between the real color image corresponding to the gray image and a label value of 1 according to a loss function, calculating the loss between the false color image generated by the gray image and a label value of 0 according to the loss function, and finally updating the parameters of the identifier by using a back propagation algorithm; wherein a label value of 1 represents a real image and a label value of 0 represents a generated pseudo-color image;

step 3.2: fixing parameters of the discriminator, calculating the loss between the generated pseudo color image and the label value of 1 according to a loss function, and finally updating the parameters of the color image generator by utilizing a back propagation algorithm;

step 3.3: continuously circulating the process of updating parameters of the discriminator and the color image generator in the steps 3.1 and 3.2 until the loss value is converged, and generating a false color image with good effect by the color image generator, namely obtaining an optimized image coloring model;

and 4, directly coloring the gray image by using the optimized image coloring model.

2. The method for rendering images based on Transformer and generation countermeasure network as claimed in claim 1, wherein in step 1, the color image generator comprises multiple MWin-Transformer modules, and the function of the MWin-Transformer module is to extract and reconstruct the features of the images, and output a 3-channel effective 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 enhanced forward propagation network LeFF.

3. The Transformer-based and confrontational network-generating image rendering method of claim 2, wherein the color image generator generates a pseudo-color image as follows:

X′＝Embedded Tokens(X _in )

X″＝W-MSA(LN(X))+X′

X _out ＝LeFF(LN(X″))+X″

wherein X _in Representing an input as a gray image or a false color image;

embedding Tokens denotes to assign X to _in Converting into a vector;

x' represents X _in Inputting the vector obtained by Embedding Tokens and outputting the vector;

then inputting a result LN (X ') obtained by carrying out layer normalization on the vector X' into a multi-head self-attention mechanism W-MSA based on a window to obtain a vector with extracted characteristic information, and adding the vector with X 'to obtain a vector X' with more collected characteristic information; x 'represents the input of X' into the output of the window-based multi-headed autofocusing mechanism and layer normalization operation;

continuously carrying out layer normalization on the vector X ', inputting the normalized LN (X ') into a local enhanced forward propagation network to obtain a vector with more extracted local characteristic information, and adding the vector with X ' to obtain a vector X with more gathered local characteristic information _out ，X _out Represents the input of X "into the locally enhanced forward propagation network LeFF and the resulting output of the layer normalization operation.

4. The transform-and-generative-confrontation-network-based image rendering method according to claim 3, wherein the layer normalization operation is calculated by:

wherein, the action object of LN layer is

X represents a vector, μ and δ represent the mean sum of each sampleThe variance of the measured values is calculated,

and

as affine learning parameters, d _k Is the dimension of the concealment that is to be hidden,

indicating that the number is a k-dimensional vector.

5. The transform-based and generative countermeasure network-based image rendering method of claim 3, wherein the window-based multi-headed self-attention mechanism is as follows:

the method comprises the steps of dividing a pseudo-color image into a plurality of windows, and then performing self-attention calculation in the different windows, wherein the patch number in one window is far less than all small blocks in one picture, and the number of the windows is kept unchanged, so that the calculation complexity of the multi-head self-attention mechanism based on the windows and the size of the image are changed into a linear relation from a square relation, and the calculation complexity of a model is greatly reduced.

6. The transform-and-generate countermeasure-network-based image shading method of claim 2, wherein a convolution is added to the forward propagation network in the Mwin-transform module, thereby forming a local enhanced forward propagation network LeFF.

7. The Transformer-based and generative countermeasure network-based image rendering method of claim 1, wherein the loss function is:

wherein,

wherein, G ^* The sum of the loss functions is represented,

the presentation conditions are generated to counter network loss,

a weighting coefficient representing a Chardonnier loss, and λ represents the Chardonnier loss;

x represents an input gray image;

y represents a real color image corresponding to the input gray image;

log represents a base 2 logarithmic function;

representing the independent variables as x, y;

represents an independent variable of x;

ε represents a value of 10 ^-3 Constant coefficient of (d);

the absolute value is obtained.

Images