Image restoration method based on improved ViT-Unet

Abstract

The invention discloses an image restoration method based on an improved ViT-Unet. The method comprises the following steps of: 1, collecting and preprocessing data, and extracting a diversified image data set; stage 2, through combination of an improved ViT module and dynamic weight gating convolution, efficient feature extraction and down-sampling are realized, and the method is used as an encoder of a U-net structure; stage 3, performing feature up-sampling and image reconstruction by utilizing the synergistic effect of an improved ViT module and Pixel Shuffle to serve as a decoder of a U-net structure; and stage 4, in combination with an enhanced deep learning strategy, the model is trained and optimized by adjusting the learning rate, and the adaptability and performance of the model in different scenes are ensured. The method is wide in application scene, is suitable for repairing medical images, satellite images and daily shot pictures, can remarkably improve the processing speed of a repairing task and the quality of a result, and has an important value for improving the technical level of the image processing field.

Description

Technical Field

[0001] The present invention relates to the technical field of computer vision and image processing, and in particular to an image restoration method based on an improved ViT-Unet. Background Art

[0002] Image inpainting is an important research direction in the field of computer vision and image processing, which aims to restore damaged or missing areas in images to provide complete and visually coherent image output. Traditional image inpainting methods are usually based on mathematical models or image statistical properties, such as partial differential equations (PDE) and exemplar-based methods. These methods try to fill in the missing areas with appropriate content by leveraging local features and prior knowledge of the image. However, as the complexity of image content increases, these traditional methods are often unable to cope with complex scenes, and are prone to discontinuous boundaries or unnatural inpainting effects.

[0003] In recent years, with the rise of deep learning technology, the field of image restoration has undergone significant changes. Deep learning models, especially convolutional neural networks (CNNs), can automatically learn high-level semantic features of images and effectively handle complex restoration tasks through large-scale data training. Convolutional neural networks capture image features at different scales through their hierarchical structure, which can better handle complex texture and structure restoration tasks. However, despite the success of CNNs in many applications, they still rely mainly on local features, which limits their ability to process global contextual information.

[0004] To overcome this limitation, researchers gradually began to introduce the Transformer architecture into image restoration tasks. The Transformer model initially performed well in the field of natural language processing, relying on the Self-Attention Mechanism to capture global dependencies in sequence data. Subsequently, the Vision Transformer (ViT) was proposed in an attempt to apply this mechanism to the visual field. However, directly applying the Transformer in image processing faces the problems of high computational complexity and high memory consumption, especially when processing high-resolution images. These challenges prompted researchers to explore more efficient Transformer variants in order to reduce the computational and memory burden while maintaining high performance.

[0005] Based on this background, many efficient Transformer architectures have emerged. Such models significantly reduce memory usage and computational complexity by optimizing the self-attention mechanism and reducing redundant calculations, and are suitable for tasks that require processing large-scale image data. At the same time, other research directions such as the Visual State Space Model (VMamba) have also shown the potential for processing long sequence data in visual tasks. The VSS block in VMamba can effectively capture contextual information in two-dimensional space through its unique selective scanning mechanism, and still maintain excellent performance under conditions of limited computing resources.

[0006] In summary, although deep learning and Transformer have made important progress in the field of image restoration, they still face challenges in terms of efficient use of computing resources and balanced processing of global and local features. The development of current technology provides a good foundation for further innovation and also points out the direction of future research. Summary of the invention

[0007] The present invention relates to an image restoration method based on an improved ViT-Unet, in particular to a method for low-cost image restoration. When processing high-resolution images, traditional convolutional neural networks (CNNs) have limitations in the process of visual feature extraction, which easily leads to loss of details and blurred information. Therefore, the present invention proposes an image restoration method based on an improved ViT-Unet, which combines the advantages of convolutional networks and Vision Transformer (ViT) to improve the accuracy and efficiency of image restoration.

[0008] The present invention specifically proposes a solution that uses an enhanced gated convolution structure to achieve effective fusion of local and global information of the image. This method can not only repair images with large-area defects, but also capture detailed information of the image at different scales, thereby generating more natural and high-fidelity repair results. By dynamically adjusting the convolution features and the gating mechanism, the present invention ensures the retention of important features and suppresses noise, thereby improving the generalization ability of the model and the ability to overcome overfitting.

[0009] In addition, the present invention has a wide range of application scenarios and is suitable for the restoration of medical images, satellite images, and daily photos. It can significantly improve the processing speed and result quality of restoration tasks, and is of great value in improving the technical level in the field of image processing.

[0010] Specifically, the implementation steps of the present invention include the following stages:

[0011] Phase 1: Data collection and preprocessing, extracting diverse image datasets.

[0012] In the second stage, the improved ViT module is combined with the gated convolution with dynamic weights to achieve efficient feature extraction and downsampling as the encoder of the U-net structure;

[0013] Phase 3: Using the synergy of the improved ViT module and PixelShuffle to perform feature upsampling and image reconstruction as the decoder of the U-net structure;

[0014] Phase 4: Combined with enhanced deep learning strategies, the model is trained and optimized by adjusting the learning rate to ensure its adaptability and performance in different scenarios.

[0015] Furthermore, the first stage is specifically implemented as follows:

[0016] 1.1 Acquisition of training data set: using data from CelebA-HQ, Places2 and Paris StreetView data sets as original images, and PConv’s mask data set as mask data set;

[0017] 1.2 Adjust the original image size to 256×256×3 pixels, synthesize it with the PConv mask data set to form the image to be repaired, and preprocess it. Use a 7×7 convolutional layer to extract the initial feature map to obtain the initial feature map X 0 ∈R C×H×W .

[0018] Furthermore, the second stage is specifically implemented as follows:

[0019] The structure of each layer of the U-net encoder is improved. In order to better extract features from the image to be repaired and prevent the pixels to be repaired from being used as part of the features in ViT for calculation, thereby reducing the feature expression ability of ViT, an improved ViT is combined with a gated convolution with dynamic weights to extract features of the image while retaining the image details. Gated convolution has been widely proven in image repair tasks to produce better repair effects than traditional convolutions. It avoids the disadvantage that traditional convolution operations treat every pixel as a valid pixel during the calculation process. The network can not only dynamically adjust the convolution weights according to the input data, but also selectively activate important features through the gating mechanism to suppress noise or irrelevant information. The local processing capability of gated convolution can effectively enhance the fusion of multi-scale features, make up for the possible shortcomings of ViT in capturing local features, improve ViT's performance in processing long-distance dependencies, and dynamically select features through the gating mechanism to reduce unnecessary parameters. The original ViT module adopts a sandwich structure, which consists of two FFN blocks wrapping a cascaded attention layer structure. The original formula is as follows:

[0020]

[0021] The formula after adding gated convolution is as follows:

[0022]

[0023] Among them, x i is the input feature map, x i+1 is the output feature map, g is the gating signal, is the fully connected layer, is a fully connected block composed of multiple fully connected layers. is a cascaded attention layer, ⊙ is an element-by-element multiplication operation, and [:,:c,:,:] represents extracting the first c channels of a sample from the tensor.

[0024] Furthermore, a VSS block is introduced in the ViT module to replace the FFN block, forming a structure consisting of two VSS blocks wrapping a cascaded attention layer, and each VSS block consists of multiple VSS layers. The VSS layer integrates the selection mechanism into its state model space and enables it to be applied to the visual field. The specific implementation of the improved ViT module after the introduction of the VSS block is as follows: the SS2D module in the VSS layer is used to scan and generate a sequence, and the information is effectively decided to be propagated or discarded according to the relevance of each token in the sequence. By modeling the visual state space, the dynamic modeling capability of continuous image frames and global context is provided. The formula becomes:

[0025]

[0026] in, is the i-th visual state space layer, It is a VSS block composed of multiple visual state space layers. The sandwich structure is provided with flexible adaptability by two VSS blocks. It provides comprehensive modeling of visual information, adjusts weights and performs knowledge fusion at the very beginning of input, and then uses the cascaded attention mechanism to split and cascade channels to achieve efficient feature extraction and fusion. The integration of each level of attention can significantly improve the diversity and depth of information.

[0027] To summarize, in the encoder part, the encoding features are reduced layer by layer from 256×256 to 128×128, 64×64, and 32×32, respectively retaining global and local detail information. Each layer of the encoder consists of a gated convolution and an improved ViT module.

[0028] Further, the stage three is specifically implemented as follows:

[0029] Improvements are made to the U-net decoder. PixelShuffle is an efficient upsampling technique that can increase the resolution of an image by rearranging the channels in the feature map. By adding a convolutional layer, PixelShuffle upsampling is better combined with the attention mechanism of the improved ViT module, so that PixelShuffle can take advantage of the modeling ability of the improved ViT module for global context while upsampling, and fully utilize the features of the input image. In this method, a 1×1 convolutional layer is first added to process the input feature map, and the input feature map x with in_channels channels is converted to i , through the convolution operation, it becomes an intermediate feature map x with inter_channels channels i ', the inter_channels is calculated by the preset upsampling factor (upscale_factor), and then PixelShuffle upsampling is performed. This method allows the U-net decoder to maintain feature information while controlling the dimension of the data. This channel transformation ensures that the information in the feature map can be fully extracted and optimized before being rearranged by PixelShuffle. The specific implementation is as follows:

[0030] x i ' = Conv2D(x i ,in_channels→inter_channels,kernel_size=1,stride

[0031] =1,padding=0,bias=False)

[0032] in:

[0033]

[0034] y = PixelShuffle(x i ', upscale_factor = 2)

[0035] Among them, x i is the input feature map, in_channels is the number of input channels, x i ' is the intermediate feature map, inter_channels is the number of intermediate channels, upscale_factor is the upsampling factor, kernel_size is the convolution kernel size, stride is the step size, padding is the padding, bias is the bias term, PixelShuffle() is the upsampling method, and y is the upsampling result.

[0036] In the decoder part, the encoded features are upsampled layer by layer through the decoder to restore to the original resolution of the input image (256×256×3 pixels). The decoder of each layer is combined with the PixelShuffle upsampling and the improved ViT module, and the information from the encoder is integrated through a channel fusion module. The decoder restores the features to 64×64, 128×128, and 256×256 layer by layer, and is fused with the 64×64, 128×128, and 256×256 features from the encoder processed by the improved ViT in the channel fusion module, and finally generates a high-resolution repaired image.

[0037] The channel fusion module is designed to efficiently fuse the features from the corresponding layer of the encoder with the low-level information that has been upsampled, thereby retaining details and structural information during image reconstruction. Each layer of the U-Net structure has a jump connection that outputs the processed features in the encoder to the channel fusion module, and the size of the feature is the same as the upsampled feature of the next layer in the decoder. This jump connection design enables the model to effectively retain the detailed features extracted during the encoding process and perform feature fusion through the channel fusion module during the decoding process. This not only enhances the contextual information of the decoder, but also effectively improves the quality of the final output image, ensuring better performance in image reconstruction and segmentation tasks. Further, the stage four is specifically implemented as follows:

[0038] The model is combined with an enhanced deep learning strategy. By adjusting the learning rate, the model is trained and optimized to ensure its adaptability and performance in different scenarios. During the model training phase, we adopted a multi-stage learning rate scheduling strategy to improve the convergence and performance of the model. In the initial stage, we set the learning rate to a larger value (1e -4 ) to quickly approach the minimum of the loss function. As training progresses, the learning rate will gradually decrease and eventually decrease to (2e -5 ) to fine-tune the model and ensure its stable performance in a variety of scenarios.

[0039] The model uses a combination of multiple loss functions in the design of the loss function to comprehensively optimize the quality of image restoration. We used multiple loss functions such as reconstruction loss and adversarial loss in training to enhance the model's ability to restore image details. The formula of the loss function is as follows:

[0040]

[0041] in, is the reconstruction loss, and its weight λ r =1, is the perceptual loss, and its weight λp =1, is the style loss, and its weight λ s =20, To counter the loss, its weight λ a =0.1.

[0042] The first is the reconstruction loss, which uses the L1 norm to measure the difference between the generated image and the real image at the pixel level. The choice of L1 loss helps to reduce the blurring effect of the image, making the restored image clearer and sharper. Its formula is as follows:

[0043]

[0044] Among them, I out is the inpainted image generated by the model, I g is the real target image.

[0045] Second, the perceptual loss uses the pre-trained VGG-19 network to extract features (using φ i denoted by ), which measures the difference between the restored image and the real image in high-level semantic features. The introduction of perceptual loss aims to improve the visual effect of the image, making the restored result closer to the real image in visual perception. Its formula is as follows:

[0046]

[0047] Among them, φ i is the activation function of the i-th layer of the VGG network, N i is the number of features at the i-th layer, Represents the mathematical expectation operation.

[0048] In addition, the model also uses style loss to maintain the consistency of style features between the generated image and the real image. Style loss measures the style features of the image by calculating the difference in the Gram matrix, which helps to generate realistic texture details. Its formula is as follows:

[0049]

[0050] in, Represents the Gram matrix calculated from the activation map of the jth layer. The formula of the Gram matrix is ​​as follows:

[0051]

[0052] Among them, F j (x) is the activation feature map of x in layer j, F j (x)T is the transposed matrix of the activation feature map of x at layer j, C j is the number of channels in the jth layer, H j W j is the spatial resolution of each channel in the jth layer.

[0053] Finally, the model introduces adversarial loss, using the loss function of the least squares generative adversarial network (LSGAN) to make the generated image distribution closer to the distribution of real images. The use of adversarial loss enhances the model's ability to generate realistic images and reduces artifacts and unnatural transition effects. The formula is as follows:

[0054]

[0055] Where D is the discriminator network, and its output is a scalar representing the probability of the input image being judged as real or generated. Represents the real image I g The expected operation, Represents the generated image I out expected operation.

[0056] By combining these loss functions, the model can simultaneously consider pixel-level accuracy, perceptual quality, style consistency, and image authenticity during the optimization process, thereby comprehensively improving the restoration effect.

[0057] Beneficial Effects of the Invention

[0058] The core technical feature of the present invention is to improve the traditional ViT-Unet structure by combining multiple technologies, and optimize the detail recovery problem in image restoration tasks. Through this combination method, the present invention can effectively improve the ability of image feature extraction, perform personalized adjustments for image content in different scenes, and ensure high-quality restoration effects.

[0059] Compared with existing patents and papers, the present invention can significantly improve the effect of image restoration while reducing the consumption of video memory resources and speeding up training when using the same training samples. Experimental results show that in multiple test scenarios, the restoration effect of the present invention is better than the existing technology, demonstrating its practicality and innovation in the field of image processing. BRIEF DESCRIPTION OF THE DRAWINGS

[0060] Figure 1 It is a structural diagram of the present invention.

[0061] Figure 2 This is a comparison chart between the present invention and other mainstream models.

[0062] Figure 3A detailed comparison of the improved ViT module used in the present invention and the restoration using traditional ViT DETAILED DESCRIPTION

[0063] In order to make the objectives, technical solutions and advantages of the present invention more clearly understood, the present invention is further described in detail below in conjunction with embodiments and drawings. The exemplary embodiments of the present invention and their description are only used to explain the present invention and are not intended to limit the present invention.

[0064] Step s1: The training datasets used are the CelebA-HQ dataset and the NVIDIA Irregular Mask dataset, which are used to synthesize the image to be repaired. The original image of the CelebA-HQ dataset is stored in the file path img_file, and the mask data of the NVIDIA Irregular Mask dataset is stored in the file path mask_file. At the same time, the pre-trained weight file of VGG19 is used as a feature extractor when calculating the perceptual loss to measure the difference between the generated image and the real image in high-level semantic features.

[0065] Step s2: Define a class named CreateDataset to assist in reading the above training set images (CelebA-HQ dataset) and mask data (NVIDIA Irregular Mask dataset), and pre-set the preprocessing hyperparameters. First, the make_dataset function is used to parse the image data file path (img_file) and the mask data file path (mask_file), and check whether the data scale of the image and mask matches; when the amount of mask data is insufficient, the mask path list is repeatedly expanded to cover the entire image dataset. Then, a function named get_transform is set to define the image preprocessing hyperparameters and adjust the input image size to 256×256×3. The image intensity values ​​are normalized and linearly mapped to the range [0.0,1.0]. Random rotation and random horizontal flip operations are added to enhance the transformation diversity of the data. After that, in the getitem method, the image and mask data are loaded by index and the preprocessing transformations defined above are applied respectively. Finally, the image and its corresponding mask are stored in the form of a dictionary, where the image is stored as the keyword "img" and the mask is stored as the keyword "mask".

[0066] Step s3: To generate training input, the image to be repaired (i.e., the image containing the missing area) is first fed into the initial encoder module of the network. The initial feature map is extracted through a 7×7 convolutional layer, and after normalization and random affine transformation, a feature map X with a size of 256×256×16 is generated. 0, and input it into the encoder part of the U-net structure. During the encoding process, the feature map is sent to the encoder for layer-by-layer processing. Each layer consists of a gated convolution and an improved ViT module. The gated convolution dynamically activates the valid pixels in the input features in each layer of encoding, suppresses invalid features, and enhances the pertinence of image restoration. The improved ViT module uses its global modeling ability to capture long-distance dependencies and establishes connections between features through the attention mechanism. First, the feature map X 0 The input is fed into the first layer of the encoder and processed by gated convolution and improved ViT to generate the feature map X' 0 , the feature is output to the channel fusion module in the first layer of the decoder through skip connection for temporary storage, and then the feature is resized to a feature map X of 128×128×32 through downsampling 1 , the feature is input into the second layer and processed by gated convolution and improved ViT to generate feature map X' 1 And input into the second layer of the decoder for temporary storage, and then downsample the feature map X' 1 Resize to 64×64×64 feature map X 2 After gated convolution and improved ViT processing, the feature map X' is generated 2 And input into the third layer of the decoder for temporary storage, and resize to obtain a 32×32×128 feature map X by downsampling 3 Finally, after gated convolution and improved ViT processing, the feature map X' is obtained 3 , the last layer is no longer downsampled.

[0067] The encoded features are upsampled layer by layer through the decoder to restore the original resolution of the input image (256×256×3 pixels). In the decoder part, the PixelShuffle upsampling method is used in conjunction with the improved ViT module. PixelShuffle upsampling improves image resolution by rearranging channel information while retaining detail features. The channel fusion module fuses the features sampled by PixelShuffle upsampling with the features from the jump connection at each layer to ensure that the reconstruction result has both detail and structure information. First, the input feature map X' is obtained from the encoder 3 , the feature map is processed by the improved ViT module and upsampled by Pixelshuffle, and resized to a 64×64×64 feature map Y 2 , Y 2 Will be combined with the feature map X' through the channel fusion module 2 Perform feature fusion to generate feature map Y' 2 , Y' 2It is processed again by ViT and upsampled by Pixelshuffle to generate a feature map Y of size 128×128×32 1 , and with the feature map X' 1 Fusion generates feature map Y' 1 , and the same steps are followed to generate a feature map Y' of size 256×256×16 0 Finally, a 7×7 convolutional layer is used to supplement the high-frequency signal and compress the generated feature map channel to convert it into a repaired image of size 256×256×3, thus completing the image repair. During the training process, the initial learning rate is set to 1e-4 to achieve rapid convergence. Starting from the 10th epoch of model training, the learning rate of each round will be reduced by 1e-5 based on the previous round until the learning rate drops to 2e-5.

[0068] Step s4: After completing the network training, we used the trained model to repair the images in the test set, generated the corresponding repair results, and performed a qualitative analysis on them. The test results are as follows: Figure 2 As shown in the figure, the superiority of this method in image restoration tasks is intuitively demonstrated by comparing it with other mainstream image restoration methods. These methods include EdgeConnect (EC), Gated Convolution (GC), Partial Convolution (PIC), Recurrent Feature Reasoning (RFR), and Diverse Transformer Supervision (DTS). By comparing these models on different data sets (CelebA-HQ, Places2, Paris StreetView), the difference in the effects of different methods in processing missing areas of the image is demonstrated. It can be intuitively shown that our model is more natural and reasonable when repairing the damaged textures of the eyes and mouth, and the lines of the hair are smoother. It shows that our model performs well in repairing large areas of damage.

[0069] based on Figure 3The advantages of the improved ViT module used in this method are analyzed. First, in the window and bed parts of the image, the improved ViT module can better preserve details and textures. In the enlarged contrast area, we can see that the improved ViT module can restore the striped pattern on the bed more clearly, showing sharper edges and more natural color transitions. In contrast, the restoration effect of the traditional ViT module in the same area is slightly blurred, the edges of the striped pattern show obvious blurring and ghosting, and the color transition is not smooth enough. Secondly, in the bedside and wall parts of the image, the improved ViT module can more accurately restore the color and texture of the wall, and can maintain the clarity of the edges of the furniture. This performance shows that the improved ViT module has achieved a better balance in restoring the global structure and local details of the image. The traditional ViT module shows a greater sense of blur in the same area, which further reflects the effectiveness of this method.

[0070] At the same time, we also conducted quantitative analysis. As shown in Tables 1 and 2, we used FID, PSNR and SSIM to quantitatively evaluate our model on two datasets. We compared the evaluation results with the commonly used models when the irregular mask ratio was 20% to 40% and 40% to 60%. The experimental results on the CelebA-HQ and Places2 datasets verified the effectiveness of the method. As can be seen from the data in the table, the model of this study performed well in indicators such as FID, PSNR, and SSIM. Especially in the experiment with an irregular mask ratio of 20% to 40%, the FID indicators of the two datasets were 5.79 and 14.34, respectively, which were 0.37 and 0.6 lower than the suboptimal method, respectively, reflecting that the feature distribution of the image generated by this model is closer to the real image distribution. PSNR was improved by 1.61 and 0.83 on the two datasets compared with the suboptimal method, respectively, reflecting that the image generated by this model has less distortion. The SSIM on the two datasets is improved by 0.008 and 0.016 respectively compared with the suboptimal method, which shows that the images generated by this model have higher structural similarity. The experimental results on the CelebA-HQ and Places2 datasets verify the effectiveness of this method, which is in a leading position in restoration effect.

[0071] Table 1

[0072]

[0073] Table 2

[0074]

[0075] Those skilled in the art will appreciate that the embodiments in the present application may be provided as methods, systems, or computer program products. Therefore, the present application may adopt a complete hardware embodiment, a complete software embodiment, or a form of an embodiment combining software and hardware. Moreover, the present application may adopt a form of a computer program product implemented on one or more computer-usable storage media (including but not limited to disk storage, CD-ROM, optical storage, etc.) containing computer-usable program code.

[0076] The embodiments of the present application are described with reference to the flowcharts and / or block diagrams of the methods, devices (systems), and computer program products according to the embodiments of the present application. It should be understood that each process and / or box in the flowchart and / or block diagram, as well as the combination of the processes and / or boxes in the flowchart and / or block diagram, can be implemented by computer program instructions. These computer program instructions can be provided to a processor of a general-purpose computer, a special-purpose computer, an embedded processor, or other programmable data processing device to generate a machine, so that the instructions executed by the processor of the computer or other programmable data processing device generate instructions for implementing the processes in the flowchart. Figure 1 A process or multiple processes and / or boxes Figure 1 A device that provides the functions specified in a block or multiple blocks.

[0077] These computer program instructions may also be stored in a computer-readable memory capable of directing a computer or other programmable data processing device to operate in a specific manner, so that the instructions stored in the computer-readable memory produce an article of manufacture including an instruction device, which implements the process Figure 1 A process or multiple processes and / or boxes Figure 1 A function specified in one or more boxes.

[0078] These computer program instructions can also be loaded onto a computer or other programmable data processing device so that a series of operating steps are executed on the computer or other programmable device to produce a computer-implemented process, thereby providing instructions for implementing the process. Figure 1 A process or multiple processes and / or boxes Figure 1 The steps for the functions specified in one or more boxes.

[0079] Although the preferred embodiments of the present application have been described, those skilled in the art may make additional changes and modifications to these embodiments once they have learned the basic creative concept. Therefore, the appended claims are intended to be interpreted as including the preferred embodiments and all changes and modifications falling within the scope of the embodiments of the present application.

[0080] Obviously, those skilled in the art can make various changes and modifications to the embodiments in the present application without departing from the spirit and scope of the embodiments in the present application. Thus, if these modifications and variations of the embodiments in the present application fall within the scope of the claims and their equivalents in the embodiments of the present application, the embodiments of the present application are also intended to include these modifications and variations.

Claims

1. An image restoration method based on improved ViT-Unet, characterized in that: The enhanced gated convolution structure is used to achieve effective fusion of local and global information of the image. By dynamically adjusting the convolution features and gating mechanism, important features are retained while noise is suppressed. The specific steps include the following stages: Phase 1: Data collection and preprocessing, extracting diverse image datasets; In the second stage, the improved ViT module is combined with the gated convolution with dynamic weights as the encoder of the U-net structure to achieve efficient feature extraction and downsampling; Phase 3: Using the synergy of the improved ViT module and PixelShuffle to perform feature upsampling and image reconstruction as the decoder of the U-net structure; Phase 4: Combined with enhanced deep learning strategies, the model is trained and optimized by adjusting the learning rate to ensure its adaptability and performance in different scenarios.

2. The image restoration method based on the improved ViT-Unet according to claim 1, characterized in that: The specific implementation of the first stage is as follows: 1.1 Acquisition of training data set: using data from CelebA-HQ, Places2 and Paris StreetView data sets as original images, and PConv’s mask data set as mask data set; 1.2 Adjust the original image size to 256×256×3 pixels, synthesize it with the PConv mask data set to form the image set to be repaired, and preprocess it. Use a 7×7 convolutional layer to extract the initial feature map to obtain the initial feature map X0∈R C×H×W .

3. The image restoration method based on the improved ViT-Unet according to claim 1, characterized in that: The second stage is specifically implemented as follows: 2-1. Improve the structure of each layer of the U-net encoder. The original ViT module adopts a sandwich structure, which consists of two FFN blocks wrapping a cascaded attention layer structure. The original formula is as follows: The formula after adding gated convolution is as follows: Among them, x i is the input feature map, x i+1 is the output feature map, g is the gating signal, is the fully connected layer, is a fully connected block composed of multiple fully connected layers. is a cascaded attention layer, ⊙ is an element-by-element multiplication operation, and [:,:c,:,:] represents extracting the first c channels of a sample from the tensor; 2-2. In the ViT module, a VSS block is introduced to replace the FFN block to form a structure consisting of two VSS blocks wrapping a cascaded attention layer, and each VSS block consists of multiple VSS layers; the VSS layer integrates the selection mechanism into its state model space and enables it to be applied to the visual field. The specific implementation of the improved ViT module after the introduction of the VSS block is as follows: the SS2D module in the VSS layer is used to scan and generate a sequence, and the information is effectively decided to be propagated or discarded according to the relevance of each token in the sequence. By modeling the visual state space, the dynamic modeling capability of continuous image frames and global context is provided, and its formula becomes: in, is the i-th visual state space layer, It is a VSS block composed of multiple visual state space layers. The sandwich structure is provided with flexible adaptability by two VSS blocks. It provides comprehensive modeling of visual information, adjusts weights and performs knowledge fusion at the very beginning of input, and then uses the cascaded attention mechanism to split and cascade channels to achieve efficient feature extraction and fusion.

4. The image restoration method based on the improved ViT-Unet according to claim 3 is characterized in that: The specific implementation of the third stage is as follows: Improvements are made to the U-net decoder by adding a convolutional layer to better combine the Pixelshuffle upsampling with the attention mechanism of the improved ViT module, so that PixelShuffle can use the improved ViT module's ability to model the global context while upsampling, and make full use of the features of the input image; the specific implementation is as follows: First, a 1×1 convolutional layer is added to process the input feature map, and the input feature map x with in_channels channels is converted to i , through the convolution operation, it becomes an intermediate feature map x with inter_channels channels i ', the inter_channels is calculated by the preset upsampling factor, and then PixelShuffle upsampling is performed. The specific implementation is as follows: x i '=Conv2D(x i ,in_channels→inter_channels,kernel_size=1,stride =1,padding=0,bias=False) y=PixelShuffle(x i ',upscale_factor=2) Among them, x i is the input feature map, in_channels is the number of input channels, x i ' is the intermediate feature map, inter_channels is the number of intermediate channels, upscale_factor is the upsampling factor, kernel_size is the convolution kernel size, stride is the step size, padding is the padding, bias is the bias term, PixelShuffle() is the upsampling method, and y is the upsampling result; The encoded features are upsampled layer by layer through the decoder to restore to the original resolution of the input image. The decoder of each layer is jointly used by the PixelShuffle upsampling and the improved ViT module, and the information from the encoder is integrated through a channel fusion module. The decoder restores the features to 64×64, 128×128, and 256×256 layer by layer, and is fused with the 64×64, 128×128, and 256×256 features from the encoder processed by the improved ViT in the channel fusion module to finally generate a high-resolution repaired image.

5. The image restoration method based on the improved ViT-Unet according to claim 4, characterized in that: The channel fusion module is used to fuse the features from the corresponding layer of the encoder with the low-level information that has been upsampled, so as to retain the details and structural information during the image reconstruction process; each layer of the U-Net structure has a jump connection, which outputs the processed features in the encoder to the channel fusion module. The size of the feature is the same as the upsampled feature of the next layer in the decoder; the design of this jump connection enables the model to effectively retain the detailed features extracted during the encoding process and perform feature fusion through the channel fusion module during the decoding process.

6. The image restoration method based on the improved ViT-Unet according to claim 4, characterized in that: The fourth stage is specifically implemented as follows: Combine the model with enhanced deep learning strategies, adjust the learning rate, train and optimize the model to ensure its adaptability and performance in different scenarios; In the design of the loss function, this model uses a combination of multiple loss functions to comprehensively optimize the quality of image restoration. The formula of the loss function is as follows: in, is the reconstruction loss, and its weight λ r =1, is the perceptual loss, and its weight λ p =1, is the style loss, and its weight λ s =20, To counter the loss, its weight λ a =0.1; Reconstruction losses The L1 norm is used to measure the difference between the generated image and the real image at the pixel level. The formula is as follows: Among them, I out is the inpainted image generated by the model, I g is the real target image; Perceived loss The pre-trained VGG-19 network is used to extract features and measure the difference between the restored image and the real image in high-level semantic features. The formula is as follows: Among them, φ i is the activation function of the i-th layer of the VGG network, N i is the number of features at the i-th layer, It represents the operation of finding mathematical expectation; Style Loss Used to maintain the consistency of style features between the generated image and the real image, the style loss measures the style features of the image by calculating the difference in the Gram matrix. The formula is as follows: in, Represents the Gram matrix calculated from the activation map of the jth layer. The formula of the Gram matrix is ​​as follows: Among them, F j (x) is the activated feature map of x in layer j, F j (x) T is the transposed matrix of the activation feature map of x at layer j, C j is the number of channels in the jth layer, H j W j is the spatial resolution of each channel in the jth layer; The model introduces adversarial loss The loss function of the least squares generative adversarial network is used to make the generated image distribution closer to the distribution of the real image. The formula is as follows: Where D is the discriminator network, and its output is a scalar, which represents the probability of the input image being judged as real or generated; Represents the real image I g The expected operation, Represents the generated image I out expected operation.

Images