An image inpainting method based on a transformer and a convolutional neural network

Abstract

The present application belongs to the field of machine learning, computer vision and image inpainting, and provides an image inpainting method based on a Transformer module and a convolutional neural network. The image inpainting model of the present application uses a hybrid end-to-end network combining a Transformer module and a convolutional neural network to obtain a better image inpainting result after edge extraction of an input image to obtain structural prior, while maintaining high precision and also maintaining the speed of image inpainting. The present application designs two new types of Transformer modules, which are respectively used for extraction and processing of global features and local features of an image, fully utilizes effective information of image pixels and feature maps, and effectively improves the effect of image inpainting. At the same time, a Canny operator and a convolutional neural network are used to obtain image edges, which provide a type of structural prior for the image inpainting network to improve the image inpainting effect.

Description

TECHNICAL FIELD

[0001] The application belongs to the field of machine learning, computer vision and image inpainting, and relates to a Transformer module, a convolutional neural network, an edge detection algorithm and a generative adversarial network, in particular to an image inpainting method based on a Transformer and a convolutional neural network. BACKGROUND

[0002] Image inpainting refers to a technology of restoring images distorted due to various external influences and reconstructing them by using background information, which can generate missing or damaged parts of image data by using complex algorithms. As an important computer vision research problem, in recent years, image inpainting technology based on deep learning has attracted more and more researchers' attention. With the development of convolutional neural networks, in 2016, Deepak Pathak proposed Context Encoders, which opened up the use of convolutional neural networks for image inpainting. The network uses an encoder-decoder structure, and uses a full connection layer based on channel features in the middle layer to connect the encoder and the decoder. At the same time, a generative adversarial network structure is used for training to improve the network performance. On this basis, in 2018, Yu Jiahui et al. proposed the Deepfill series, Yan Zhaoyi et al. proposed the Shift-Net network structure, which further accelerated the development of image inpainting networks. The above three methods focus more on fixed position image inpainting, and the repair results of arbitrary shape area have large artifacts, and cannot fully utilize the pixel information and structural information of the image. As a branch of computer vision research with strong practical application, image inpainting technology can solve more practical image problems in the future.

[0003] The Transformer module is a popular module structure in natural language processing research, which has been gradually applied to computer vision tasks such as image classification, object detection, semantic segmentation and deep completion in recent years. Compared with traditional convolutional neural networks that focus on internal information of convolution kernels, the Transformer module can model the whole image and obtain the whole image features through the self-information mechanism. The Transformer module is composed of a multi-head attention layer, a linear layer and a normalization layer, and the final output is connected in a skip layer manner. The global modeling capability of the Transformer is realized by the multi-head attention layer, which can integrate global features and process useful information in them.

[0004] In the image inpainting problem, the perception of the whole image is crucial, so the whole image feature extraction capability of the Transformer module can better obtain the global information of the image. The convolutional neural network itself has an inductive bias that is beneficial to the generation task, but it lacks a receptive field in the whole image range. Therefore, effectively combining the two will greatly help improve the performance of the image inpainting task. SUMMARY

[0005] The present application aims to provide an image inpainting method based on Transformer module and convolutional neural network, which solves the problem that existing image inpainting algorithms cannot simultaneously utilize image global information and inductive bias, and simultaneously utilizes image edge information as image structure prior, further enhancing the effect of image inpainting.

[0006] The method of the present application can be deployed on desktop computer host or smart phone and other devices for image inpainting work.

[0007] The technical solution of the present application is:

[0008] An image inpainting method based on Transformer and convolutional neural network, the steps are as follows:

[0009] Step 1: Collect image data sets and mask data sets for model training;

[0010] Step 2: Use the image to be repaired and the Canny edge generation operator to generate the edge information map of the image to be repaired;

[0011] Step 3: Input the edge information map generated in step 2 into the edge generator to obtain the completed input image edge map;

[0012] Step 4: Input the image to be repaired and the edge information map generated in step 3 into the image inpainting network for repair to obtain the final inpainting image; Specifically as follows:

[0013] The image inpainting network mainly consists of a convolutional neural network and a Transformer module. In the training of the image inpainting network, a discriminator is added to perform adversarial training on the generated results of the image inpainting network to enhance the generation performance of the image inpainting network. The basic module of the convolutional neural network mainly refers to Edge Connect, which uses stacked convolutional layers and activation layers to increase the depth of the network. Two convolutional layers with a step size of 2 are used in the convolutional neural network to construct a 2-fold and 4-fold down-sampling feature pyramid encoder. The decoder uses a symmetrical structure to the encoder to ensure that the output image has the same size as the input image. Eight residual block structures are used in the middle layer to fill the confirmed area using residual structures. At the same time, the feature maps of the encoder are sent to the corresponding decoder as additional information for up-sampling by using a global Transformer module to obtain global information at each down-sampling size. In addition, an additional down-sampling network is designed, which inputs the input image and the completed edge into the down-sampling network, and outputs the features after down-sampling using a local Transformer module to process local information. The obtained local information is input into the decoder to obtain the local detail information of the image.

[0014] The global information obtained by the global Transformer from the feature maps of each layer of the encoder includes the following steps:

[0015] (1) Preprocessing: The input image to be inpainted is divided and embedded, and a pixel classification flag is added. The two are connected in channel order, and the obtained image feature sequence is then position embedded. In the position embedding process, the Dropout layer is removed to retain all image feature information and ensure that there is no information loss during the image inpainting process.

[0016] (2) The image feature sequence processed in step (1) is input into the global Transformer module to extract global self-information. Specifically as follows:

[0017] (2.1) The features obtained in step (1) are normalized and decomposed to obtain K sequences, Q sequences, and V sequences for calculating image self-information. These three sequences are input into the multi-head attention mechanism for self-attention calculation. The calculation formula is as follows:

[0018]

[0019] In the formula, is the feature dimension of each head in the multi-head attention mechanism;

[0020] (2.2) The self-information obtained in step (2.1) is connected to the input feature X by a skip layer, and then added; after that, the result is used for pixel restoration by a multi-layer perception (MLP); the structure of the MLP is composed of a hierarchical normalization layer, a fully connected layer, a GeLU activation layer and another fully connected layer, and the final output is connected to the input by a skip layer;

[0021] (2.3) Steps (2.1) and (2.2) are repeated N times to obtain the final global self-information;

[0022] (3) The global self-information is output through an MLP head structure, and then the output is passed through a reshaping module to restore the output size to the same size as the input size, and finally an up-sampling convolution is used to arrange the output channel number to the same number as the input channel number.

[0023] The local Transformer obtains image detail information mainly including the following steps:

[0024] (1) Preprocessing: input the down-sampled image feature into three different sequence extraction convolutional networks to extract specific K, Q and V sequences, and the calculation formula is as follows:

[0025] K = Reshape (f (X)), Q = Reshape (g (X)), V = Reshape (h (X))

[0026] Wherein, f (·), g (·), h (·) are three different convolution calculations;

[0027] (2) The K, Q and V sequences obtained are used for self-information calculation in the convolution kernel, and the formula is as follows:

[0028]

[0029] In the formula, is the feature dimension of each head in the multi-head attention mechanism;

[0030] (3) The local self-information obtained is integrated through a multi-head attention layer;

[0031] (4) Steps (2) and (3) are repeated N times;

[0032] (5) The local self-information obtained in step (4) is connected to the input feature by a skip layer, and the final local information output is obtained after passing through a convolution layer.

[0033] The beneficial effects of the present application are:

[0034] (1) A Transformer and convolutional neural network combined image inpainting network is designed, which fully utilizes the inductive bias of the convolutional neural network and the performance of the Transformer module to improve the image inpainting performance. At the same time, this network is the first end-to-end training network combining the two.

[0035] (2) Global and local Transformer modules are designed to fully utilize global and local image features, solve the problem of lack of global receptive field in ordinary convolutional neural networks, enhance image inpainting performance, and effectively remove artifacts in image inpainting results. BRIEF DESCRIPTION OF DRAWINGS

[0036] Figure 1 Figure 1 is a schematic diagram of an image inpainting network structure.

[0037] Figure 2 Figure 3 is a schematic diagram of a global Transformer module structure.

[0038] Figure 3 Figure 4 is a schematic diagram of a local Transformer module structure. DETAILED DESCRIPTION

[0039] The specific embodiments of the present application will be further described below in combination with the drawings and technical solutions.

[0040] Figure 1 Figure 1 is a schematic diagram of an image inpainting network structure. Before entering the image inpainting network, first obtain the completion edge map of the input image as the structure priori, which is input into the image inpainting network together with the image to be inpainted. The basic convolution module of the encoder and decoder of the image inpainting network is composed of a convolution layer, an instance normalization layer (IN layer), and a ReLU activation function; the middle layer residual block is composed of two reflection padding layers, two convolution layers, a spectral normalization layer (SN layer), and an instance normalization layer. The network input is down-sampled by two convolution modules with a step of 2 to obtain a 4-fold down-sampled feature map. The feature map is further processed by 8 residual blocks in the middle layer for feature and pixel filling. In addition, the features of each layer in the down-sampling encoder are input into the corresponding global Transformer module to extract the full image information of the image, and the information is input into the corresponding up-sampling decoder. In addition, a local Transformer module is separately set up for the 4-fold down-sampling branch, which is composed of four convolution modules. The local information after 8 local Transformer modules is input into the decoder as additional information for up-sampling to restore the image.

[0041] Figure 2Figure 1 is a structural diagram of a global Transformer module. The module is used as an intermediate layer connecting the encoder and the decoder symmetric position feature map. The module uses a self-attention mechanism to calculate the global self-information of the input features as additional prior information input into the up-sampling decoder in addition to the convolutional features, providing additional global self-information for image inpainting.

[0042] Figure 3 Figure 2 is a structural diagram of a local Transformer module. The module is used as an image local detail extraction layer, and a separate down-sampling branch is designed to extract image detail features. The extracted details are then calculated using a local attention mechanism to achieve the extraction of local detail attention in the entire image. In the extraction of local information, a convolutional layer is used to obtain an attention calculation sequence, increasing the attention to local information in the image and further improving the performance of image inpainting.

[0043] The training set of the image inpainting network consists of two parts: image data set and mask data set. The image data set is selected according to different inpainting tasks. For scene inpainting, Places2 data set is selected, and for face inpainting, CelabA data set is selected. The mask data set consists of NVIDI A-ALDR data set and hand-drawn QD-IMD data set to adapt to various image covering methods.

[0044] The network optimization method selects AdamW optimizer, the initial learning rate of the network is set to 0.001, the learning rate of the discriminator is 1 / 10 of the network learning rate, the Places2 data set is trained for 2 epochs, the CelebA data set is trained for 40 epochs, the learning rate is gradually decayed during the training process, and the input image size of the network during the inference process is 512x512.

[0045] The structure of the image inpainting network is as follows:

[0046]

[0047]

Claims

1. An image inpainting method based on Transformer and Convolutional Neural Network, characterized in that, The steps are as follows: Step 1: Collect the image dataset and the mask dataset for occlusion region for model training; Step 2: Using the image to be repaired and the Canny edge generation operator, generate an edge information map of the image to be repaired; Step 3: Input the edge information map generated in Step 2 into the edge generator to obtain the completed input image edge map; Step 4: Input the image to be repaired and the edge map of the completed input image generated in Step 3 into the image inpainting network for repair, and obtain the final repaired image; details are as follows: The image inpainting network consists of a convolutional neural network (CNN) and a Transformer module. A discriminator is added during the training of the image inpainting network to perform adversarial training on the generated results, thereby enhancing its performance. The basic module of the CNN follows the Edge Connect architecture, using stacked convolutional and activation layers to increase network depth. Two convolutional layers with a stride of 2 are used to construct a feature pyramid encoder with 2x and 4x downsampling. The decoder adopts a symmetrical structure with the encoder to ensure that the output image is the same size as the input image. The intermediate layers use an 8-residual-block structure to fill in missing regions. Simultaneously, for the encoder's feature map, global information at each downsampling size is obtained through the global Transformer module and sent to the corresponding decoder using a skip-layer structure as additional upsampling information. In addition, an additional downsampling network is designed. The input image and the edge map of the completed input image are input into the downsampling network, and the downsampled feature output is processed by the local Transformer module to process local information. The obtained local information is then input into the decoder to obtain local detail information of the image.

2. The image inpainting method based on Transformer and convolutional neural networks according to claim 1, characterized in that, The global Transformer module obtains global information about the feature maps of each layer of the encoder through the following steps: (1) Preprocessing: The input image to be repaired is divided into blocks and embedded, and a pixel classification label is added. The two are connected in channel order, and the resulting image feature sequence is then embedded in position. In the position embedding process, the Dropout layer is removed to retain all image feature information and ensure that no information is lost during the image repair process. (2) The image feature sequence processed in step (1) is input into the global Transformer module to extract global self-information; as follows: (2.1) The features obtained in step (1) are subjected to hierarchical normalization to decompose them into K-sequence, Q-sequence, and V-sequence for calculating image self-information. These three sequences are then input into a multi-head attention mechanism for self-attention calculation. The calculation formula is as follows: ； In the formula, For each head in the multi-head attention mechanism; (2.2) The self-information obtained in step (2.1) is added to the input image feature sequence by performing a skip-layer connection; then the result is used to perform pixel recovery using a multilayer perceptron (MLP); the structure of the MLP consists of a hierarchical normalization layer, a fully connected layer, a GeLU activation layer and another fully connected layer, and the final output is then connected to the input by a skip-layer connection. (2.3) Repeat steps (2.1) and (2.2) N times to obtain the final global self-information; (3) The global self-information is output through an MLP header structure, and then the output is processed through an shaping module to resize the output to the same size as the input. Finally, an upsampling convolution is used to resize the number of output channels to the same number of input channels.

3. The image inpainting method based on Transformer and convolutional neural networks according to claim 1, characterized in that, The local Transformer obtains image detail information through the following steps: (1) Preprocessing: The downsampled image features are input into three different sequence extraction convolutional networks to extract specific K, Q, and V sequences. The calculation formula is as follows: ； in, Calculate for three different convolutions. for ; (2) Calculate the self-information within the convolution kernel of the obtained K, Q, and V sequences, using the following formula: ； In the formula, For each head in the multi-head attention mechanism; (3) The obtained local self-information is integrated through a multi-head attention layer; (4) Repeat steps (2) and (3) N times; (5) Connect the local self-information obtained in step (4) with the input features by skipping layers, and then obtain the final local information output after passing through the convolutional layer.

Images