Image restoration method based on gated Fourier convolution residual error and multi-head attention mechanism

Abstract

The invention discloses an image restoration method based on a gated Fourier convolution residual error and a multi-head attention mechanism, and the method is characterized in that the method specifically comprises the steps: collecting a damaged image as a data set, carrying out the preprocessing, and dividing the data set into a training set and a test set according to a set proportion; constructing an LGFCDANet network, wherein the network comprises a coarse recovery network and a fine recovery network; carrying out adversarial training on the LGFCDANet network to obtain a trained LGFCDANet network, and testing the trained LGFCDANet network by adopting a test set; and carrying out image restoration by adopting the trained network. The invention aims to enhance the understanding and recovery capability of the model on image details, thereby improving the restoration effect and enhancing the dynamic adjustment capability on the color and texture of a special area.

Description

Technical Field

[0001] The present invention belongs to the field of image restoration technology, and in particular relates to an image restoration method based on gated Fourier convolution residual and multi-head attention mechanism. Background Art

[0002] In the field of mural image restoration, traditional image restoration algorithms often face challenges in processing complex textures and color restoration, especially in the reconstruction of large-scale damaged murals. With the rapid development of deep learning technology, convolutional neural networks (CNNs) have made significant progress in image reconstruction and restoration applications. Researchers have successively proposed a variety of improved algorithms, such as UNet structures and generative adversarial networks (GANs), to improve the quality and accuracy of restoration. However, these methods usually have difficulty capturing long-range dependencies and detailed information in images, especially in the case of high-scale occlusion, the model may cause information loss or reconstruction distortion.

[0003] To address these issues, recent research has enhanced the model's perceptual capabilities by introducing attention mechanisms. For example, edge-constrained attention mechanisms are used to dynamically weight image features to better capture important information, or methods for inpainting mural images utilize a bidirectional flow of information between texture and structure. However, existing techniques still fall short in simultaneously considering both spatial and channel-level features, limiting the comprehensiveness and stability of the restoration results. In particular, in color filling and feature fusion, the correct and effective integration of information from multiple feature sources remains a critical challenge.

[0004] In addition, in the process of repairing images with high damage rates, a single recovery strategy often fails to achieve ideal results. Summary of the Invention

[0005] Purpose of the invention: In order to solve the problems existing in the above-mentioned prior art, the present invention discloses an image restoration method based on gated Fourier convolution residual and multi-head attention mechanism.

[0006] Technical solution: The present invention provides an image restoration method based on gated Fourier convolution residual and multi-head attention mechanism, which specifically includes the following steps:

[0007] Step 1: Collect damaged images as a dataset and preprocess them, then divide them into training and test sets according to the set ratio;

[0008] Step 2: Construct an LGFCDANet network, which includes a coarse restoration network and a fine restoration network. The coarse restoration network includes a first fusion module and a first generator. The damaged image and the line drawing corresponding to the image are input into the fusion module for feature fusion, and the fused features are input into the first generator. The features output by the first generator are fused with the original damaged image and the features output by the fine restoration network, and then input into the fine restoration network. The fine restoration network includes a second fusion module, a second generator, and a discriminator. The feature map and mask map input into the fine restoration network are input into the second fusion module for feature fusion, and then the fused features are input into the second generator. The image output by the second generator is fused with the input image of the fine restoration network and then input into the discriminator for true and false discrimination;

[0009] Step 3: Perform adversarial training on the LGFCDANet network to obtain the trained LGFCDANet network, and use the test set to test the trained LGFCDANet network;

[0010] Step 4: Use the trained network to perform image restoration.

[0011] Furthermore, the first generator includes an initialization layer and a gated Fourier convolution residual network connected in sequence.

[0012] Furthermore, the initialization layer includes a reflection filling layer, a convolution layer, a regularization layer and a ReLU activation function connected in sequence.

[0013] Furthermore, the gated Fourier convolution residual network first downsamples through four consecutive linear gated convolution blocks, thereby reducing the input image to one-eighth of the original size, and then introduces four consecutive MFU modules. Finally, the features output by the MFU module are input into the UCG network for feature extraction in both spatial and channel dimensions; then, a continuous four-layer upsampling module is used for upsampling, and the upsampling module is jump-connected to the linear gated convolution block.

[0014] Furthermore, the first three linear gated convolution blocks have the same structure. First, the image input to the linear gated convolution block is convolved, and then the chunk function is used to split the convolution result into two parts, x and v. x is subjected to depthwise separable convolution and then input to the GeLU activation function. The result output by the GeLU activation function is multiplied by v, and the multiplication result is regularized, convolved, and Dropout processed in sequence. The fourth linear gated convolution block first performs a convolution operation on the feature image input to the linear gated convolution block, and then the chunk function is used to split the convolution result into two parts, x1 and v1. x1 is subjected to depthwise separable convolution and then input to the GeLU activation function, and the result output by the GeLU activation function is multiplied by v1.

[0015] Furthermore, the structures of the four MFU modules are the same. First, the feature map input to the MFU module is fast Fourier transformed, and the transformation result is separated into real and imaginary parts. Then, the stack operation is used to splice the real and imaginary parts to obtain the first tensor; the permute operation is used to place the dimensions of the real and imaginary parts immediately after the channel dimension, and then the view operation is used to splice the real and imaginary parts on the channel dimension to form the second tensor; the position code locMap is set, and the locMap is expanded to match the size of the current data. The expanded position code is spliced with the second tensor on the channel dimension, and the spliced tensor is The result is convolved, and the shift operation is performed on the convolved result to move the low-frequency component to the center of the spectrum, and then input into the ReLU activation function; the real and imaginary parts of the output of the ReLU activation function are separated, the locMap is shifted to align it with the current feature, the aligned locMap and the separated real and imaginary parts are spliced in the channel dimension to obtain a third tensor, the third tensor is convolved and shifted again, the convolution result is weightedly fused with the first tensor using the sigmod function, the fused features are reorganized into complex form, and an inverse fast Fourier transform is performed.

[0016] Furthermore, the UCG network first uses a position encoder to perform position encoding projection on the feature image input to the UCG network to generate spatial gating information h i :

[0017] h i =W L ReLU(W L-1 ...ReLU(W1P i +b1)+…)+b L

[0018] Among them, P i represents the Cartesian coordinates of the i-th pixel of the feature image input to the UCG network, W Lis the linear weight matrix of the L-th multilayer perceptron, b L is the bias of the L-th multilayer perceptron, and ReLU represents the activation function;

[0019] Then h i Input to the spatial gated convolution to obtain the spatial gated weight S i :

[0020] S i =W spatial *h i +b spatial

[0021] Among them, W spatial is the weight matrix of the spatial gated convolution, b spatial is the bias term, * represents the convolution operation;

[0022] The feature image input to the UCG network is input to the channel-gated convolution. The channel-gated convolution performs average pooling on the feature image input to the UCG network in the spatial dimension to generate the global mean of the input feature channel. These global means are then linearly transformed through 1x1 convolution to obtain the output C of the channel-gated convolution. i ;

[0023] S is transformed by sigmoid function i and C i Merge to form the final gating weight G i ; Set the gate weight G i Perform normalization operation and convert the normalized G i It is element-wise multiplied with the feature image input to the UCG network, and finally the multiplication result is convolved.

[0024] Furthermore, the second generator shown first uses four consecutive downsampling layers to downsample the feature map input to the second generator. The last downsampling layer outputs the feature image Y, [Y, Y] is input to the DAFA module, and then the output of the DAFA module is input to the residual module. Finally, the output of the residual module is upsampled for four consecutive layers, and the upsampling and downsampling are jump-connected. The residual module includes an MFU module and first to four residual blocks connected in sequence; the MFU module has the same structure as the MFU module in the first generator.

[0025] Furthermore, the DAFA module first splices the channels of the two feature images Y of size [1, 256, 32, 32], and reshapes the spliced tensor into a five-dimensional tensor of size [1, 2, 256, 32, 32]. The height dimension of the five-dimensional tensor is then summed to obtain a single feature of size [1, 256, 32, 32]. The single feature is flattened into a feature of size [1, 1024, 256]. The query vector Q, key vector K, and value vector V are then calculated through a linear layer. These three vectors are then input into the multi-head attention module, and the outputs of the multi-head attention module are merged and linear operations are performed again.

[0026] Furthermore, the discriminator includes five partial convolution modules connected in sequence, and the output of each partial convolution module is activated using the LeakyReLU activation function; the five partial convolution modules have the same structure, first dividing the input features into two parts according to the segmentation ratio, performing a convolution operation on the first part, and merging the result of the convolution operation with the second part.

[0027] Beneficial effects:

[0028] This paper proposes a new line-guided gated Fourier cascade dual attention network (LGFCDANet). This network comprises two stages: coarse and fine inpainting, each performing its own restoration function. This enhances the model's ability to understand and restore mural details, improving restoration results and enhancing the ability to dynamically adjust the color and texture of specific areas.

[0029] 2. This paper proposes a new gated Fourier convolution residual network, which includes a linear gated convolution block, an improved unbiased fast Fourier convolution block (MFU) module, and a novel lightweight coordinate gating module (UCG). Combining GLU and Fourier transform enhances feature extraction capabilities, especially when processing complex textures and structures. The UCG module enhances the utilization of spatial information through positional encoding, making the network more accurate in processing local details of the image.

[0030] 3. This paper proposes a new DynamicAttentionFeatureAggregator (DAFA) module, which strengthens the role of dynamic attention mechanism in feature aggregation by combining spatial and channel fields.

[0031] 4. The present invention uses partial convolution blocks to fill the 70×70 PatchGAN architecture, which reduces redundant calculations and further retains important features without changing the discriminator effect. BRIEF DESCRIPTION OF THE DRAWINGS

[0032] Figure 1The overall structure diagram of the LGFCDANet network provided by the present invention;

[0033] Figure 2 This is a structural diagram of the gated convolution block of generator 1 provided by the present invention;

[0034] Figure 3 A structural diagram of the MFU provided by the present invention;

[0035] Figure 4 A structural diagram of the UCG module provided by the present invention;

[0036] Figure 5 A structural diagram of the DAFA module provided by the present invention;

[0037] Figure 6 A structural diagram of the discriminator provided by the present invention;

[0038] Figure 7 This is a visual qualitative comparison result diagram of the LGFCDANet network provided by the present invention and other methods on the mural dataset;

[0039] Figure 8 This is a comparison of the PSNR values of the visualization results of the LGFCDANet network provided by the present invention and other methods on the mural dataset;

[0040] Figure 9 This is a comparison of the SSIM values of the visualization results of the LGFCDANet network provided by the present invention on the mural dataset with other methods; DETAILED DESCRIPTION

[0041] The accompanying drawings, which constitute a part of the present invention, are used to provide a further understanding of the present invention. The exemplary embodiments of the present invention and their descriptions are used to explain the present invention and do not constitute improper limitations on the present invention.

[0042] The main steps of the technical solution of the present invention are as follows:

[0043] Step 1: Dataset collection and image preprocessing. This example obtains a total of 6,661 mural images from different dynasties as a dataset. After preprocessing, the dataset is divided into a training set and a test set according to a set ratio.

[0044] Step 2: Build the LGFCDANet network, such as Figure 1 The LGFCDANet network shown is mainly divided into a coarse recovery network and a fine recovery network, where the coarse recovery network mainly includes generator 1, and the fine recovery network mainly includes generator 2 and a discriminator.

[0045] Step 3: Input the mural image obtained in step 1 into the LGFCDANet network in step 2 for training, calculate the loss function and update the network parameters to obtain the optimal parameter model.

[0046] Step 4: Feed the mural with real damage from step 1 into the best trained model and output the restored mural image.

[0047] The steps of step 1, data set collection and image preprocessing, are as follows:

[0048] Through sources such as the "Complete Collection of Dunhuang Murals in China," we collected a total of 5,472 mural images, covering authentic murals from sites such as the Mogao Grottoes in Dunhuang and the Yulin Grottoes in Xi'an. Since most authentic murals suffer varying degrees of damage, this example acquired 1,189 replicas to fill the dataset, resulting in a total of 6,661 mural images. After cropping all images to 720×720 pixels, this example divided the training, test, and validation sets into a 7:2:1 ratio.

[0049] The masks used in this embodiment are from the irregular mask dataset provided by Liu et al. All masks will be randomly assigned to the mural images during the training and testing processes.

[0050] The second step is to build a line-guided gated Fourier convolution residual and multi-head attention mechanism network, including the steps of coarse recovery network and fine recovery network:

[0051] Step 2.1: First, build a coarse restoration network. This network consists of three components: an input layer, an initialization layer, and a gated Fourier convolution residual network (GFCRNet). After feature fusion of the damaged mural image and the corresponding line drawing, the image is fed into the initialization layer for processing and then output to the GFCRNet for further processing. The initialization layer and GFCRNet together constitute Generator 1.

[0052] Step 2.2: Next, the feature map output by generator 1 and the damaged mural image are weighted fused and input into the fine restoration network.

[0053] Step 2.3: In the fine recovery network, the weighted fused feature map is first accepted, and the feature is fused with the mask map, and then input into the generator 2 for processing.

[0054] Step 2.4: The feature map processed in step 2.3 and the feature map output by generator 1 are weighted fused.

[0055] Step 2.5: Input the feature map generated in step 2.4 into the discriminator for true and false discrimination, and use the perceptual loss, L1 loss and histogram loss to jointly train the generator network and the discriminator network to achieve the final good restoration effect.

[0056] The main structures of the initialization layer and the gated Fourier convolution residual network GFCRNet in step 2.1 are as follows:

[0057] The initialization layer is the beginning of the network, which is used to receive image input and preprocess the image. The initialization layer includes a reflection padding layer, which fills the original image size [1, 3, 256, 256] to [1, 4, 262, 262]. Subsequently, after a 7×7 convolution layer, the image size is changed to [1, 32, 256, 256]. After that, after a batch normalization layer, the ReLU activation function is used to activate the features.

[0058] In GFCRNet, the image is first scaled down to one-eighth its original size through four consecutive downsampling linear gated convolution blocks. After the input features are passed through a convolutional layer, the feature map is split into two parts, x and v, using the chunk function. A depthwise separable convolution is performed on x, followed by a GeLU activation function, while v is directly applied to the gating mechanism. The processed features are multiplied by the gated part to achieve dynamic feature control. Next, regularization is performed before mapping them to the output feature space through another 1x1 convolutional layer and applying a dropout layer. This entire process combines convolution and gating mechanisms to further enhance the expressive power of features.

[0059] The linear gated convolution block structure is shown in the figure Figure 2 shown.

[0060] The first downsampling layer receives an image of size [1, 32, 256, 256]. After a 1×1 convolutional layer, the feature map is split into two parts, x and v. For the x part, a 3×3 depthwise separable convolution is applied, followed by a GeLU activation function. The activated result is multiplied by v to generate dynamic feature control. Next, a Drop layer is used for regularization, followed by a 1×1 convolution layer and a new Drop layer for final processing. The output image of the first gated convolution block is of size [1, 32, 256, 256]. Subsequently, a 4×4 convolution layer and a regularization layer are used to reduce the image size to [1, 64, 128, 128], and processing begins with the next gated convolution block.

[0061] The image sizes after processing by the subsequent three gated convolution blocks are [1, 128, 64, 64], [1, 256, 32, 32], and [1, 256, 32, 32], respectively. Note that only the first three gated convolution blocks are followed by 4×4 convolution and regularization layers, and the output of the fourth gated convolution block directly enters the MFU.

[0062] Subsequently, four consecutive MFU modules are introduced. The MFU module provides a unique feature processing method by introducing the Fourier transform into the neural network. Its workflow includes initializing the convolutional layer, batch normalization and activation function, and setting the position mapping and mixing parameters. During the forward propagation process, the module can selectively scale the input features and then apply the Fourier transform to convert the spatial data into the frequency domain. The formula is shown below:

[0063]

[0064] Where x(m,n) represents the pixel values of the input image; X(u,v) represents the complex form of the frequency domain representation, with M and N representing the image height and width, respectively. M represents the image height (the second dimension of the tensor), and N represents the image width (the third dimension of the tensor). u represents the coordinate in the frequency domain height direction. V represents the coordinate in the frequency domain width direction, and i represents the imaginary unit.

[0065] In the frequency domain, the features are separated into real and imaginary parts and concatenated with the dynamic position map. They are then processed through a convolutional layer, applying activation functions and spectral centering. The module then uses learnable parameters to mix the original and processed frequency domain features, and finally converts the data back to the spatial domain through an inverse Fourier transform, as shown in the following formula:

[0066]

[0067] The dynamic range is then adjusted to restore the original size.

[0068] The MFU structure diagram is as follows Figure 3 The specific processing flow of MFU is as follows:

[0069] Input feature map

[0070] Input: A feature map x of shape [B, C, H, W], where B is the batch size, C is the number of channels, and H and W are the height and width respectively.

[0071] Step 1: Fast Fourier Transform (FFT)

[0072] Perform real fast Fourier transform (rfftn) on the feature map to obtain the frequency domain representation ffted.

[0073] Outputs a complex tensor of shape [B, C, H', W' / 2+1], where W' / 2+1 is the frequency domain width halved due to the symmetry of the real FFT. H' is the new height, and W' is the new width.

[0074] Step 2: Separate the real and imaginary parts

[0075] Separate ffted into real and imaginary parts and concatenate them into a new four-dimensional tensor as follows:

[0076] The stack operation is used to combine the real and imaginary parts into a single tensor, forming a [B, C, H', W' / 2+1, 2] tensor. The permute operation moves the real and imaginary information from the last dimension to the position after the channel dimension, forming a new [B, C, 2, H', W' / 2+1] tensor. Finally, the view operation is used to combine the channels of the real and imaginary parts together to form a new tensor with a shape of [B, C × 2, H', W' / 2+1]: the first C channels correspond to the real part, and the last C channels correspond to the imaginary part.

[0077] Step 3: Splice the position code locMap

[0078] locMap is a positional encoding independently defined and initialized as a model parameter, with a shape of [input_shape, input_shape / / 2+1]. locMap is expanded to match the size of the current batch data [B, 1, H', W' / 2+1] and concatenated with the frequency domain features along the channel dimension to form a tensor of [B, C × 2+1, H', W' / 2+1].

[0079] Step 4: First convolution and activation

[0080] The concatenated features are processed using the first convolutional layer, and the output shape is [B, out_channels×2, H', W' / 2+1].

[0081] An FFT Shift operation is applied to the result to move the low-frequency components to the center of the spectrum, and the nonlinear expression capability is enhanced through the ReLU activation function.

[0082] Step 5: Second convolution aligned with positional encoding

[0083] Apply an FFT Shift operation to the locMap to align it with the current frequency domain features.

[0084] The real and imaginary parts of the features are separated and concatenated with the aligned locMap_shift in the channel dimension to form a tensor of [B, out_channels×2+1, H', W' / 2+1].

[0085] The second convolutional layer is used for processing, and the output shape is still [B, out_channels×2, H', W' / 2+1].

[0086] Then, an FFTShift operation is performed on the feature map processed by the second convolutional layer.

[0087] Step 6: Weighted Fusion

[0088] Use the learnable parameter lambda_base to control the fusion weight of the original frequency domain features and the processed features.

[0089] The contribution ratio between the two is dynamically adjusted according to lambda_base to obtain the fused feature map.

[0090] Step 7: Inverse Fourier Transform (IFFT)

[0091] The fused features are reorganized into complex form and inverse fast Fourier transform (irfftn) is performed to convert back to the spatial domain.

[0092] The out shape is restored to [B, out_channels, H', W'].

[0093] After passing through the MFU module, the features are input into the UnifiedCoordGate (UCG). By extracting features in both spatial and channel dimensions, UCG can capture the complexity and details of the input data more finely. The spatial gating mechanism enhances the spatial perception of the model, enabling it to dynamically adjust features based on location. In addition, the channel averaging and gating mechanism enable the model to adaptively adjust channel weights, improving adaptability to different image content. This combination of spatial and channel feature extraction significantly improves the model's expressiveness and performance, and solves the drawback of traditional coordinate gating that only extracts features in the spatial dimension.

[0094] The structure diagram of the UCG module is as follows Figure 4 As shown: UCG first includes a position encoder, which converts the two-dimensional coordinates (x i ,y i ) is mapped to a high-dimensional space through linear transformation and activation function, which helps to generate specific spatial features for each pixel location.

[0095] Assume that the Cartesian coordinates of the i-th pixel on the image are (x i ,y i ), where (x i ,y i ) is a coordinate in the linear interpolation range [-1,1]. For each input coordinate (xi ,y i )(Shape is a two-dimensional vector P i =[x i ,y i ]), this embodiment performs a series of linear transformations and activation function packaging, as shown below. The first-level linear transformation formula is shown below.

[0096]

[0097] Where W1∈R 64×2 , b1∈R 64 are the linear weight matrices and bias terms in the model. The input coordinates are converted to 64-dimensional features. The subsequent linear transformations for each layer are shown below.

[0098] h k+1 =ReLU(W k+1 h k +b k+1 )

[0099] Among them, k represents the kth layer, W k+1 ∈R 64×64 , b k+1 ∈R 64 After each layer of linear transformation, a 64-dimensional vector is generated. The final generated (h k+1 ) is the embedded high-dimensional feature vector. Therefore, the output of the position encoder can be expressed in a recursive form as shown below:

[0100]

[0101] Among them, W is the weight matrix in each layer of linear transformation, and b is the bias term.

[0102] Based on the position encoder, a 1x1 convolution operation is used to project the embedded position information onto the channel dimension to generate spatial gating weights, as shown below.

[0103] S i =W spatial *h i +b spatial

[0104] Among them, S i is the result of spatial gating, W spatial is the weight matrix of 1×1 convolution, b spatial is the bias term, and * represents a convolution operation. Channel gating performs average pooling on the tensor x input to the UCG module in the spatial dimension to generate the global mean of the input feature channels (i.e., calculate the average map). These global features are then linearly transformed using a 1x1 convolution, as shown below.

[0105] C i =W channel *Mean(x)+b channel

[0106] Among them, C i is the result of channel gating, Mean(x) is the channel mean of tensor x in the spatial dimension, W channel is the weight matrix of channel convolution, b channel is the bias term.

[0107] Spatial gating and channel gating are combined to form the final gating weights:

[0108] G i =σ(S i +C i )

[0109] Among them, G i is the gating weight, σ() represents the Sigmoid function, which normalizes the gating weight to the range of [0,1]. Subsequently, it acts on the input feature map through multiplication:

[0110]

[0111] in, is the final output result, X i is the input feature map, G i is the gate weight, W final Represents the weight of the last convolutional layer.

[0112] This feature extraction method that combines space and channels significantly improves the expressiveness and performance of the model, and solves the drawback of traditional coordinate gating that only extracts features in the spatial dimension.

[0113] Upsampling is performed through four consecutive linear gated convolutional units, during which skip connections are used to restore the feature map to its original size and retain more information.

[0114] The main structure of the generator 2 in step 2.3 is as follows:

[0115] The main structure of Generator 2 is similar to that of Generator 1, and it also uses the encoder-decoder and jump connection method. The difference is that in Generator 2, this embodiment uses traditional CNN blocks for downsampling and upsampling. At the same time, in the fine recovery network, this embodiment adopts a global residual strategy, that is, the input mural of Generator 2 is added to the output of Generator 2 to retain more original features of the image.

[0116] After four consecutive downsampling modules, this embodiment introduces DAFA. The advantage of this attention mechanism module is that it combines the powerful functions of multi-head self-attention and feature fusion. Multi-head self-attention allows the model to capture different feature representations through multiple parallel attention heads, thereby improving the model's expressiveness and stability. Each head focuses on a different feature subspace, which can carefully capture the diversity and complexity of the input data. In terms of feature fusion, the module integrates information from different feature sources by splicing multiple feature inputs and summing them in the spatial dimension. This integration method enhances the model's comprehensive perception ability. Subsequently, the self-attention mechanism is used to dynamically weight the integrated features, allowing the model to adaptively adjust the output features according to different input conditions, so that the model can achieve dynamic adjustment of color filling in image restoration.

[0117] The structure of DAFA in this embodiment is as follows Figure 5 As shown:

[0118] First, we receive two input tensors Y and Y (Y is the upsampled output), both of the same size [1, 256, 32, 32]. We then concatenate them channel-wise, with the concatenated size being [1, 512, 32, 32]. We then reshape them into a five-dimensional tensor of [B, height, C, H, W], now [1, 2, 256, 32, 32]. We then sum the height dimension to get a single feature, now of size [1, 256, 32, 32] (i.e., x'').

[0119] The image is then flattened to [B, N, C], where N = H * W = 32 * 32 = 1024, resulting in a flattened size of [1, 1024, 256]. A linear layer is then used to calculate Q, K, and V, resulting in a size of [1, 1024, 256]. A head-wise operation is then performed, splitting each vector into four heads, resulting in a size of [1, 4, 1024, 64]. The attention matrix for each head is then calculated and the softmax function is applied, resulting in a size of [1, 4, 1024, 1024]. Attention is then applied and the heads are merged. The weighted output shape is [1, 4, 1024, 64], and the output shape after merging multiple heads is [1, 1024, 256]. Finally, a linear layer is used to restore the image to its original size of [1, 256, 32, 32].

[0120] The five-dimensional expansion is performed because the input sources are not distinguished when the inputs are initially combined. Reshaping into a five-dimensional tensor allows the characteristics of each branch to be clearly distinguished. In other words, the original four-dimensional tensor [1,512,32,32] can be viewed as [[input1],[input2]] concatenated along the channel axis. Through reshaping, the channel axis (512) is split into height×C (2×256), clearly distinguishing the characteristics of each branch. Subsequently, the height dimension is summed and recombined into a four-dimensional tensor. Furthermore, this five-dimensional expansion avoids an explosion in the number of channels and reduces additional parameters.

[0121] The residual module is introduced, which mainly consists of an MFU block and four residual blocks.

[0122] The input size is [1, 256, 32, 32]. After the MFU, the size remains unchanged and is fed into the residual block. The residual block primarily consists of two convolutional layers. In the first layer, the image is first padded with a reflection padding of 1, followed by a 3×3 convolutional layer and normalization. The image is then fed into the second convolutional layer, which also padded with a reflection padding of 1, followed by a 3×3 convolution with a dilation rate of 1 and normalization. The output is then calculated by adding the input to the output of the convolutional block. At this point, the image size remains [1, 256, 32, 32].

[0123] Then, four layers of upsampling are performed continuously, and the specific process is the same as the downsampling process of generator 2, and the image size is finally adjusted to [1, 3, 256, 256]. The features are then input into the discriminator.

[0124] The discriminator structure is as follows Figure 6 As shown, the main structure is as follows:

[0125] First, an image input of size [1, 3, 256, 256] is received. This example uses the 70×70 PatchGAN architecture. The 70×70 PatchGAN determines whether overlapping image patches of size 70×70 are real. However, to reduce redundant computations and retain as many required features as possible for better discrimination, this example replaces all convolutional blocks in the original architecture with partial convolutional blocks. The discriminator's main structure includes partial convolution operations, spectral normalization, and the leakyReLU activation function.

[0126] After the first layer input passes through a partial convolution block, the input has 3 channels and a split ratio of 2. Therefore, the input is split into [1, 1.5, 256, 256] and [1, 1.5, 256, 256]. Since the number of channels is actually an integer, the size is resized to [1, 1, 256, 256] and [1, 2, 256, 256]. The output of the first part after passing through the 3×3 convolution layer remains [1, 1, 256, 256], while the output of the second part remains unchanged. The two parts are merged, and the output is resized to [1, 3, 256, 256]. The LeakyReLU activation function is applied for feature extraction.

[0127] Then, the above operation is repeated four times. Finally, the discriminator will output an output after all convolutional layers with a size of [1, 3, 256, 256].

[0128] The visual qualitative comparison results of LGFCDANet on the mural dataset provided by the embodiment and other methods are as follows: Figure 7 shown.

[0129] The comparison results of the PSNR values of the visualization results of LGFCDANet provided in the embodiment on the mural dataset and other methods are as follows: Figure 8 shown.

[0130] The SSIM value comparison results of the visualization results of LGFCDANet on the mural dataset provided by the embodiment and other methods are as follows: Figure 9 shown.

[0131] The quantitative comparison results of the LGFCDANet provided in the embodiment for repairing on the test set with different mask coverage ratios (1-10%, 10-20%, 20-30% and 30-40%) are shown in Table 1:

[0132] Table 1

[0133]

[0134]

[0135] According to the results, it can be seen that the LGFCDANet provided in this embodiment has the best performance under all masking conditions.

[0136] It should also be noted that the various specific technical features described in the above specific embodiments can be combined in any appropriate manner without contradiction. To avoid unnecessary repetition, the present invention will not further describe various possible combinations.

Claims

1. An image restoration method based on gated Fourier convolution residual and multi-head attention mechanism, characterized by: The specific steps include: Step 1: Collect damaged images as a dataset and preprocess them, then divide them into training and test sets according to the set ratio; Step 2: Construct an LGFCDANet network, which includes a coarse restoration network and a fine restoration network. The coarse restoration network includes a first fusion module and a first generator. The damaged image and the line drawing corresponding to the image are input into the fusion module for feature fusion, and the fused features are input into the first generator. The features output by the first generator are fused with the original damaged image and the features output by the fine restoration network, and then input into the fine restoration network. The fine restoration network includes a second fusion module, a second generator, and a discriminator. The feature map and mask map input into the fine restoration network are input into the second fusion module for feature fusion, and then the fused features are input into the second generator. The image output by the second generator is fused with the input image of the fine restoration network and then input into the discriminator for true and false discrimination; Step 3: Perform adversarial training on the LGFCDANet network to obtain the trained LGFCDANet network, and use the test set to test the trained LGFCDANet network; Step 4: Use the trained network to perform image restoration.

2. The image restoration method based on gated Fourier convolution residual and multi-head attention mechanism according to claim 1 is characterized in that: The first generator includes an initialization layer and a gated Fourier convolution residual network connected in sequence.

3. The image restoration method based on gated Fourier convolution residual and multi-head attention mechanism according to claim 2, characterized in that: The initialization layer includes a reflection filling layer, a convolution layer, a regularization layer and a ReLU activation function connected in sequence.

4. The image restoration method based on gated Fourier convolution residual and multi-head attention mechanism according to claim 2, characterized in that: The gated Fourier convolution residual network first downsamples through four consecutive linear gated convolution blocks, thereby reducing the input image to one-eighth of the original size, and then introduces four consecutive MFU modules. Finally, the features output by the MFU modules are input into the UCG network for feature extraction in both spatial and channel dimensions; then, a continuous four-layer upsampling module is used for upsampling, and the upsampling module is jump-connected to the linear gated convolution block.

5. The image restoration method based on gated Fourier convolution residual and multi-head attention mechanism according to claim 4, characterized in that: The first three linear gated convolution blocks have the same structure. First, the image input to the linear gated convolution block is convolved, and then the chunk function is used to split the convolution result into two parts, x and v. x is subjected to depthwise separable convolution and then input to the GeLU activation function. The result output by the GeLU activation function is multiplied by v, and the multiplication result is regularized, convolved, and Dropout processed in sequence. The fourth linear gated convolution block first performs a convolution operation on the feature image input to the linear gated convolution block, and then the chunk function is used to split the convolution result into two parts, x1 and v1. x1 is subjected to depthwise separable convolution and then input to the GeLU activation function, and the result output by the GeLU activation function is multiplied by v1.

6. The image restoration method based on gated Fourier convolution residual and multi-head attention mechanism according to claim 4, characterized in that: The four MFU modules have the same structure. First, the feature map input to the MFU module is fast Fourier transformed and the transformation result is separated into real and imaginary parts. Then, the real and imaginary parts are concatenated using the stack operation to obtain the first tensor. The permute operation is used to place the dimensions of the real and imaginary parts immediately after the channel dimension. Then, the view operation is used to concatenate the real and imaginary parts along the channel dimension to form the second tensor. Set the position code locMap, expand locMap to match the size of the current data, splice the expanded position code with the second tensor in the channel dimension, perform convolution on the spliced result, shift the convolved result, move the low-frequency component to the center of the spectrum, and then input it into the ReLU activation function; separate the real and imaginary parts of the output of the ReLU activation function, shift the locMap to align it with the current feature, splice the aligned locMap and the separated real and imaginary parts in the channel dimension to obtain a third tensor, convolve the third tensor and perform shift operation again, use the sigmod function to weightedly fuse the convolution result with the first tensor, reorganize the fused features into complex form, and perform inverse fast Fourier transform.

7. The image restoration method based on gated Fourier convolution residual and multi-head attention mechanism according to claim 4, characterized in that: The UCG network first uses a position encoder to perform position encoding projection on the feature image input to the UCG network to generate spatial gating information h i : h i =In L ReLU(W L-1 …ReLU(W1P i +b1)+…)+b L Among them, P i represents the Cartesian coordinates of the i-th pixel of the feature image input to the UCG network, W L is the linear weight matrix of the L-th multilayer perceptron, b L is the bias of the L-th multilayer perceptron, and ReLU represents the activation function; Then h i Input to the spatial gated convolution to obtain the spatial gating weight Si: S i =W spatial *h i +b spatial Among them, W spatial is the weight matrix of the spatial gated convolution, b spatial is the bias term, * represents the convolution operation; The feature image input to the UCG network is input to the channel-gated convolution. The channel-gated convolution performs average pooling on the feature image input to the UCG network in the spatial dimension to generate the global mean of the input feature channel. These global means are then linearly transformed through 1x1 convolution to obtain the output C of the channel-gated convolution. i ; S is transformed by sigmoid function i and C i Merge to form the final gating weight G i ; Set the gate weight G i Perform normalization operation and convert the normalized G i It is element-wise multiplied with the feature image input to the UCG network, and finally the multiplication result is convolved.

8. The image restoration method based on gated Fourier convolution residual and multi-head attention mechanism according to claim 4, characterized in that: The second generator shown first uses four consecutive downsampling layers to downsample the feature map input to the second generator. The last downsampling layer outputs the feature image Y, [Y, Y] is input to the DAFA module, and then the output of the DAFA module is input to the residual module. Finally, the output of the residual module is upsampled for four consecutive layers, and jump connections are made between the upsampling and downsampling. The residual module includes an MFU module and first to four residual blocks connected in sequence; the MFU module has the same structure as the MFU module in the first generator.

9. The image restoration method based on gated Fourier convolution residual and multi-head attention mechanism according to claim 8, characterized in that: The DAFA module first concatenates the channels of two feature images Y of size [1, 256, 32, 32], and reshapes the concatenated tensor into a five-dimensional tensor of size [1, 2, 256, 32, 32]. It then sums the height dimension of the five-dimensional tensor to obtain a single feature of size [1, 256, 32, 32]. The single feature is flattened into a feature of size [1, 1024, 256]. The query vector Q, key vector K, and value vector V are then calculated through a linear layer. These three vectors are then input into the multi-head attention module, and the outputs of the multi-head attention module are merged and linear operations are performed again.

10. The image restoration method based on gated Fourier convolution residual and multi-head attention mechanism according to claim 1, characterized in that: The discriminator includes five partial convolution modules connected in sequence, and the output of each partial convolution module is activated using the LeakyReLU activation function; the five partial convolution modules have the same structure, first dividing the input features into two parts according to the segmentation ratio, performing a convolution operation on the first part, and merging the result of the convolution operation with the second part.

Images