Image reconstruction method based on input-output mutual attention and hidden layer self-attention of stylegan

Abstract

The application provides a StyleGAN image reconstruction method based on input-output mutual attention and hidden layer self-attention, which is characterized in that an input image is mapped into features by an encoder, information is extracted from the features by constructing a hidden layer self-attention and input mutual attention module to obtain hidden layer codewords; information is extracted from the hidden layer codewords by constructing an output mutual attention module to obtain style codewords, the output style codewords are sent into a pre-trained StyleGAN generator to obtain a reconstructed image; the encoder, the hidden layer self-attention and input mutual attention module, and the output mutual attention module are combined into an encoding network, the StyleGAN generator is fixed, the encoding network is trained, and the image reconstruction quality is improved. Compared with the prior art, the application has the characteristics that a small amount of parameter quantity and calculation quantity are increased, the reconstructed image quality is significantly improved, and the problem of poor reconstructed image quality caused by insufficient learning of image information of the encoder is better solved.

Description

Technical Field

[0001] This invention belongs to the field of deep learning technology, and in particular to a StyleGAN image reconstruction method based on input-output mutual attention and hidden layer self-attention. Background Technology

[0002] StyleGAN has achieved considerable success in image generation, producing high-resolution and natural-looking images. Its semantically unentangled latent space also provides greater freedom for image editing. Image reconstruction aims to map an input image to a latent space, allowing the style codewords in the latent space to generate a reconstructed image close to the original, which is beneficial for downstream real-world image editing tasks. However, training a StyleGAN model is extremely time-consuming and computationally expensive; therefore, many methods utilize pre-trained StyleGAN models for image reconstruction to improve efficiency.

[0003] Image reconstruction methods using pre-trained StyleGAN models can be divided into two types: optimization-based reconstruction and encoder-based reconstruction. Optimization-based reconstruction methods can reconstruct images that are closer to the original image, but the inference speed is slow and the editability is difficult to control. Encoder-based reconstruction methods have faster inference speed and use fewer computational resources. Encoder-based reconstruction methods require building an encoder to encode the input image into style codewords, and then inputting the style codewords into the generator of the pre-trained StyleGAN model to obtain the reconstructed image. The encoder is trained by calculating the mean squared error loss, perceptual loss, and identity recognition loss between the reconstructed image and the original image.

[0004] Existing encoder-based reconstruction methods typically employ convolutional neural networks (CNNs) as the encoder structure. This structure has a small receptive field, can only establish connections between local information in the image, and cannot perform long-range modeling of image features, risking the loss of global spatial information. The attention module in the Transformer structure focuses on global information, which can compensate for the limitations of the receptive field of convolutional neural networks. Therefore, it is essential to provide a StyleGAN image reconstruction method that rationally utilizes the attention module. Summary of the Invention

[0005] The purpose of this invention is to address the shortcomings of existing technologies by providing a StyleGAN image reconstruction method based on input-output mutual attention and hidden layer self-attention, thereby improving the quality of reconstructed images. The method combines the encoder, hidden layer self-attention module, input mutual attention module, and output mutual attention module into an encoding network. This encoding network is trained, and the style codewords output by the encoding network are fed into a pre-trained StyleGAN generator to obtain the reconstructed image. This method significantly improves the quality of reconstructed images by increasing the number of parameters and computational load by a small amount, effectively solving the problem of poor reconstructed image quality caused by insufficient image information learned by the encoder.

[0006] The objective of this invention is achieved as follows:

[0007] A StyleGAN image reconstruction method based on input-output mutual attention and hidden layer self-attention is proposed. Its key feature is combining the encoder, hidden layer self-attention module, input mutual attention module, and output mutual attention module into a single encoding network. The encoding network is trained, and the style codewords output by the encoding network are fed into a pre-trained StyleGAN generator to improve the quality of the reconstructed image. The method includes the following steps:

[0008] Step 1: Map the input image to features using an encoder, and extract information from the features by constructing a hidden layer self-attention and input mutual attention module to obtain hidden layer codewords;

[0009] Step 2: Extract information from the hidden codewords by constructing an output mutual attention module to obtain style codewords, and feed the output style codewords into the pre-trained StyleGAN generator to obtain the reconstructed image;

[0010] Step 3: Combine the encoder, hidden layer self-attention and input mutual attention modules and output mutual attention modules into a single encoding network, fix the StyleGAN generator, train the encoding network, and improve the image reconstruction quality.

[0011] The hidden layer self-attention and input mutual attention modules are composed of several hidden layer self-attention modules and several input mutual attention modules connected in series. The hidden layer self-attention module is composed of self-attention modules and updates the hidden layer codewords through self-refinement of the hidden layer codewords. The input mutual attention module is composed of cross-attention modules and updates the initial hidden layer codewords with the input image features extracted by the encoder.

[0012] Step 2, which involves constructing an output mutual attention module to extract information from the hidden codewords to obtain style codewords, specifically includes:

[0013] Step 2-1: Output the mutual attention module, which is composed of cross attention modules;

[0014] Step 2-2: The hidden codewords L output by the hidden layer self-attention and input mutual attention modules are used as two types of inputs: key and value. The output mutual attention module has a dimension of B×M×C, where B is the batch size of the input data, M is the number of hidden codewords, and C is the feature dimension.

[0015] Step 2-3: Randomly initialize style codewords with dimensions B×N×C, input them into the StyleGAN style mapping network as learnable parameters to obtain parameterized initial style codewords; input the parameterized initial style codewords as query values ​​with dimensions B×N×C, where B and C are the same as in Step 2-2, and N is the number of style codewords;

[0016] Steps 2-4: Calculate the cross-similarity matrix S using the transpose of Query and Key according to the following formula (1), with a dimension of B×N×M;

[0017] S = softmax(Query × Key) T (1)

[0018] Steps 2-5: According to the following formula (2), multiply S and Value, and then pass through residual connection, normalization and feedforward layer to obtain the updated style codeword W, whose dimension is B×N×C;

[0019] W=FFN(Norm(Residual(S×Value))) (2).

[0020] The number of hidden codewords is not the same as the number of style codewords specified by the StyleGAN network, but a multiple of the number of style codewords, which can further extract image information.

[0021] Compared with the prior art, the present invention has the characteristic that it can significantly improve the quality of reconstructed images by increasing the number of parameters and computation by a small amount, and effectively solves the problem of poor quality of reconstructed images caused by insufficient image information learned by the encoder. Attached Figure Description

[0022] Figure 1 This is a schematic diagram of the process of the present invention;

[0023] Figure 2 This is a schematic diagram of the output of the encoding network. Detailed Implementation

[0024] To more clearly illustrate the technical means, technical improvements, and beneficial effects of this invention, the invention will be described in detail below with reference to the accompanying drawings.

[0025] This invention relates to a StyleGAN image reconstruction method based on input-output mutual attention and hidden layer self-attention, comprising the following specific steps:

[0026] Step 1: Map the input image to features using an encoder, and extract information from the features by constructing a hidden layer self-attention and input mutual attention module to obtain the hidden layer codeword L;

[0027] Step 2: Extract information from the hidden codeword L by constructing an output mutual attention module to obtain the style codeword W. Then, feed the output style codeword W into the pre-trained StyleGAN generator to obtain the reconstructed image.

[0028] Step 3: Combine the encoder, hidden layer self-attention and input mutual attention modules, and output mutual attention modules into a single encoding network. Fix the StyleGAN generator and train the encoding network to improve the image reconstruction quality.

[0029] Step 1 specifically includes: the hidden layer self-attention and input mutual attention module is constructed by cross-connecting several hidden layer self-attention modules and several input mutual attention modules, the hidden layer self-attention module is composed of self-attention modules, and the input mutual attention module is composed of cross-attention modules.

[0030] Step 1-1: Feed the input image into the encoder to obtain features, which serve as the key and value inputs for all input mutual attention modules;

[0031] Step 1-2: Randomly initialize a specified number of hidden codewords with a dimension of B×M×C, where B is the input data batch size, M is the number of hidden codewords, and C is the feature dimension. The number of hidden codewords M is not the same as the number of style codewords N specified by the StyleGAN network, but is k times N, k = 2, 4, 8, ... Input the randomly initialized hidden codewords into the StyleGAN style mapping network and use them as learnable parameters to obtain parameterized initial hidden codewords. The parameterized initial hidden codewords serve as the three inputs of the first-level hidden self-attention module: key, value, and query.

[0032] Steps 1-3: Except for the first-level hidden layer self-attention module, the three inputs of the hidden layer self-attention module, namely the key, value, and query, are taken from the hidden layer codeword L updated by the previous input mutual attention module. The module updates the hidden layer codeword L through self-refinement of the hidden layer codeword.

[0033] Steps 1-4: The query value of the input mutual attention module is taken from the hidden codeword L updated by the previous hidden layer self-attention module. The two inputs, key and value, are taken from the feature F output by the encoder. The module updates the hidden codeword L with the input image information extracted by the encoder.

[0034] Step 2 specifically includes: the output mutual attention module is composed of a cross attention module, which updates the initial style codeword W with the hidden layer codeword.

[0035] Step 2-1: Use the hidden codeword L output by the hidden layer self-attention and input mutual attention module as two types of input: key and value. Its dimension is B×M×C.

[0036] Step 2-2: Randomly initialize style codewords with dimensions B×N×C, input them into the style mapping network of StyleGAN and use them as learnable parameters to obtain parameterized initial style codewords. Here, B and C are the same as in Step 2-1, N is the number of style codewords, and the parameterized initial style codewords are used as the query value Query input.

[0037] Steps 2-3: Calculate the cross-similarity matrix S using the transpose of Query and Key according to Formula 1 below. The cross-similarity matrix S has dimensions of B×N×M.

[0038] S = softmax(Query × Key) T (1)

[0039] Steps 2-4: Following Formula 2 below, multiply S and Value, then pass through residual connection, normalization, and feedforward layer to obtain the updated style codeword W, which has a dimension of B×N×C.

[0040] W=FFN(Norm(Residual(S×Value))) (2)

[0041] Step 3 specifically includes: calculating the mean squared error loss, perceptual loss, and identity recognition loss between the reconstructed image and the input image, and training an encoding network composed of an encoder, a hidden layer self-attention module, an input mutual attention module, and an output mutual attention module.

[0042] Example 1

[0043] See Figure 1This invention combines an encoder, hidden layer self-attention, and input mutual attention modules to obtain hidden layer codewords. Then, an output mutual attention module is used to obtain style codewords, which are fed into the generator of a pre-trained StyleGAN model to obtain a reconstructed image. The invention trains an encoding network composed of an encoder, hidden layer self-attention, input mutual attention modules, and an output mutual attention module to improve the quality of the reconstructed image. Specifically, the invention includes the following steps:

[0044] S1: The input image is mapped to features F using an encoder. Information is extracted from the features by constructing a hidden layer self-attention and input mutual attention module to obtain the hidden layer codeword L.

[0045] See Figure 2 The specific steps are as follows:

[0046] Step 0: Construct hidden layer self-attention and input mutual attention modules using a neural network. The module is constructed by cross-connecting several hidden layer self-attention modules and several input mutual attention modules. The hidden layer self-attention module is composed of self-attention modules, and the input mutual attention module is composed of cross-attention modules.

[0047] Step 1: Construct an encoder using a convolutional neural network. Input an RGB three-channel input image with height H and width W into the encoder to obtain a feature map F, which serves as the key and value inputs for all input mutual attention modules.

[0048] Step 2: Randomly initialize a specified number of hidden codewords with dimensions B×M×C. Input these codewords into the StyleGAN style mapping network as learnable parameters to obtain parameterized initial hidden codewords. The number of hidden codewords M is a multiple of the number of style codewords N specified by the StyleGAN network. As one implementation, the number of hidden codewords M is set to 4 times the number of style codewords N in StyleGAN. The parameterized initial hidden codewords serve as the three inputs for the first-level hidden self-attention module: key, value, and query.

[0049] Step 3: Except for the first-level hidden layer self-attention module, the three inputs of the hidden layer self-attention module, namely Key, Value, and Query, are taken from the hidden layer codeword L updated by the previous input mutual attention module. The module updates the hidden layer codeword L through self-refinement of the hidden layer codeword.

[0050] Step 4: Input the query value of the mutual attention module. The query value is taken from the hidden codeword L updated by the previous hidden layer self-attention module. The key and value inputs are taken from the features F output by the encoder. The module uses the input image information extracted by the encoder to update the hidden codeword L.

[0051] S2: Extract information from the hidden codewords by constructing the output mutual attention module to obtain the style codeword W. Then, feed the output style codeword W into the pre-trained StyleGAN generator to obtain the reconstructed image.

[0052] See Figure 2 The specific steps are as follows:

[0053] Step 0: Construct the output mutual attention module using the cross attention module. The specific steps of the module are explained in Steps 3 and 4.

[0054] Step 1: Use the hidden codewords L output by the hidden layer self-attention and input mutual attention modules as two types of inputs: key and value. The dimension size is B×M×C, where B is the batch size of input data, M is the number of hidden codewords, and C is the feature dimension size.

[0055] Step 2: Randomly initialize style codewords with dimensions B×N×C, input them into the StyleGAN style mapping network as learnable parameters to obtain parameterized initial style codewords, and input the parameterized initial style codewords as the query value Query, where B and C are the same as in step 2-1, and N is the number of style codewords.

[0056] Step 3: Calculate the cross-similarity matrix S using the transpose of Query and Key according to Formula 1 below. The matrix S has dimensions B×N×M.

[0057] S = softmax(Query × Key) T (1)

[0058] Step 4: Following Formula 2 below, multiply S and Value, then pass through residual connection, normalization, and feedforward layer to obtain the updated style codeword W, with a dimension of B×N×C.

[0059] W=FFN(Norm(Residual(S×Value))) (2)

[0060] Step 5: Feed the style codeword W into the pre-trained StyleGAN generator to obtain the reconstructed image.

[0061] S3: Combine the encoder, hidden layer self-attention and input mutual attention modules, and output mutual attention module into a single encoding network. Fix the StyleGAN generator and train the encoding network by calculating the mean squared error loss, perceptual loss, and identity recognition loss between the reconstructed image and the input image, thereby improving the image reconstruction quality.

[0062] The above are merely preferred embodiments of the present invention and are not intended to limit the present invention. Any modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of the present invention should be included within the protection scope of the present invention.

Claims

1. A StyleGAN image reconstruction method based on input-output mutual attention and hidden layer self-attention, characterized in that, The method includes the following specific steps: Step 1: Map the input image to features using an encoder, and extract information from the features by constructing a hidden layer self-attention and input mutual attention module to obtain hidden layer codewords; Step 2: Extract information from the hidden codewords by constructing an output mutual attention module to obtain style codewords, and feed the output style codewords into the pre-trained StyleGAN generator to obtain the reconstructed image; Step 3: Combine the encoder, hidden layer self-attention and input mutual attention modules, and output mutual attention module into a single encoding network. Fix the StyleGAN generator and train the encoding network to improve image reconstruction quality; where: The hidden layer self-attention and input mutual attention modules are composed of several hidden layer self-attention modules and several input mutual attention modules connected in series. The hidden layer self-attention module is composed of self-attention modules and updates the hidden layer codewords through their own features. The input mutual attention module is composed of cross-attention modules and updates the initial hidden layer codewords with the input image features extracted by the encoder. Step 2, which involves constructing an output mutual attention module to extract information from the hidden codewords to obtain style codewords, specifically includes: Step 2-1: Output the mutual attention module, which is composed of cross attention modules; Step 2-2: Combine the hidden layer codewords output by the hidden layer self-attention and input mutual attention modules. As a key value With numerical values Two inputs, an input-to-output mutual attention module, with a dimension size of... , This refers to the batch size of the input data. It is the number of hidden codewords. It is the size of the feature dimension; Steps 2-3: Randomly initialize style code words, with a dimension size of... The initial style codewords are input into the StyleGAN style mapping network and used as learnable parameters to obtain parameterized initial style codewords; these parameterized initial style codewords are then used as query values. The input has a dimension size of , and Same as step 2-2, It refers to the number of words in the style code; Steps 2-4: According to the following formula (1), use and Calculate the cross-similarity matrix using the transpose Its dimension size is ; （1） In formula (1), T is the transpose of Key; Steps 2-5: According to the following formula (2), and After multiplication, the result is passed through residual join, normalization, and feedforward layer to obtain the updated style codeword. Its dimension size is ; （2） In formula (2), Residual(·) represents the residual connection, Norm(·) represents normalization, and FFN(·) represents the feedforward layer.

2. The image reconstruction method according to claim 1, characterized in that, The number of hidden codewords is not the same as the number of style codewords specified by the StyleGAN network, but a multiple of the number of style codewords, which can further extract image information.

Images