An image restoration method based on improved ViT-Unet

Abstract

The present invention discloses an image restoration method based on an improved ViT-Unet. The steps of the present invention are as follows: Phase 1, data collection and preprocessing, extracting a variety of image data sets. Phase 2, by combining the improved ViT module with the gated convolution of dynamic weights, efficient feature extraction and downsampling are achieved, serving as the encoder of the U-net structure; Phase 3, utilizing the synergy of the improved ViT module and PixelShuffle, feature upsampling and image reconstruction are performed, serving as the decoder of the U-net structure; Phase 4, combining the enhanced deep learning strategy, by adjusting the learning rate, the model is trained and optimized to ensure its adaptability and performance in different scenarios. 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 of the restoration task and the quality of the results, and is of great value to improving the technical level in the field of image processing.

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 fields of computer vision and image processing, aiming to restore damaged or missing areas in an image to provide a 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 (PDEs) and exemplar-based methods. These methods attempt 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, particularly convolutional neural networks (CNNs), can automatically learn high-level semantic features of images through large-scale data training and effectively handle complex restoration tasks. Through their hierarchical structure, convolutional neural networks capture image features at different scales, enabling them to better handle complex texture and structure restoration tasks. However, despite the success of CNNs in many applications, they still rely primarily on local features, which limits their ability to process global contextual information.

[0004] To overcome this limitation, researchers have gradually begun 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 sequential data. Subsequently, the Vision Transformer (ViT) was proposed in an attempt to apply this mechanism to the field of vision. However, directly applying the Transformer in image processing faces the problems of high computational complexity and high memory consumption, which is particularly evident when processing high-resolution images. These challenges have prompted researchers to explore more efficient Transformer variants to reduce the computational and memory burden while maintaining high performance.

[0005] Against this backdrop, many efficient Transformer architectures have emerged. By optimizing the self-attention mechanism and reducing redundant computation, these models significantly reduce memory usage and computational complexity, making them suitable for tasks that require processing large amounts of image data. Meanwhile, other research areas, such as the Visual State Space Model (VMamba), have also demonstrated potential for processing long sequences of data in vision tasks. The VSS block in VMamba, through its unique selective scanning mechanism, effectively captures contextual information in two-dimensional space and maintains excellent performance even under limited computational resources.

[0006] In summary, while deep learning and Transformers have made significant progress in image restoration, challenges remain in areas such as efficient utilization of computing resources and balancing global and local features. Current technological developments provide a solid foundation for further innovation and point the way for future research. Summary of the Invention

[0007] The present invention relates to an image restoration method based on an improved ViT-Unet, specifically a method for low-cost image restoration. Conventional convolutional neural networks (CNNs) currently have limitations in extracting visual features when processing high-resolution images, which can easily lead to loss of detail and blurred information. Therefore, the present invention proposes an image restoration method based on an improved ViT-Unet, combining the advantages of convolutional networks and the Vision Transformer (ViT) to improve the accuracy and efficiency of image restoration.

[0008] This paper specifically proposes a solution that utilizes an enhanced gated convolutional architecture to effectively fuse local and global image information. This method not only repairs images with large-area defects but also captures detailed image information at various scales, resulting in more natural and high-fidelity restorations. By dynamically adjusting convolutional features and gating mechanisms, this paper ensures the preservation of important features and suppresses noise, improving the model's generalization and 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 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 gated convolution with dynamic weights to achieve efficient feature extraction and downsampling, serving as the encoder of the U-net structure.

[0013] In the third stage, the improved ViT module and PixelShuffle are used 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 dataset: Using data from the CelebA-HQ, Places2, and Paris StreetView datasets as original images, and the PConv mask dataset as the mask dataset;

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

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

[0019] Improvements are made to each layer of the U-net encoder to better extract features from the image to be inpainted. This prevents pixels to be inpainted from being included as part of the features in ViT computation, which would otherwise reduce ViT's expressive power. By combining an improved ViT with gated convolution using dynamic weights, this approach extracts features while preserving image details. Gated convolution has been widely demonstrated to produce superior inpainting results compared to traditional convolution in image inpainting tasks. It avoids the drawback of traditional convolution operations, which treat every pixel as a valid pixel. The network not only dynamically adjusts convolution weights based on the input data but also selectively activates important features through a gating mechanism, suppressing noise or irrelevant information. The local processing capabilities of gated convolution effectively enhance the fusion of multi-scale features, addressing ViT's potential limitations in capturing local features. This improves ViT's performance in handling long-range dependencies, and reduces unnecessary parameters through its dynamic feature selection through a gating mechanism. The original ViT module employs a sandwich structure consisting of two FFN blocks wrapped around a cascaded attention layer. 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 gate 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-wise multiplication operation, and [:,:c,:,:] represents extracting the first c channels of a sample from the tensor.

[0024] Furthermore, a VSS block is introduced into 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 is composed 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 based on the relevance of each token in the sequence. By modeling the visual state space, it provides dynamic modeling capabilities for continuous image frames and global context. Its 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. The cascaded attention mechanism then splits and cascades 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] In summary, 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] Furthermore, the third stage is specifically implemented as follows:

[0029] Improvements are made to the U-net decoder. PixelShuffle is an efficient upsampling technique that increases the image resolution 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. This allows PixelShuffle to take advantage of the improved ViT module's ability to model global context while upsampling, making full use of the features of the input image. In this method, a 1×1 convolutional layer is first added to process the input feature map, converting the input feature map x with in_channels channels into a 1×1 convolutional layer. i , through the convolution operation, it becomes an intermediate feature map x with inter_channels channels i ', where inter_channels is calculated by a 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 dimensionality 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, the encoded features are upsampled layer by layer through the decoder to restore them to the original resolution of the input image (256×256×3 pixels). Each layer's decoder combines PixelShuffle upsampling with an improved ViT module, and integrates information from the encoder through a channel fusion module. The decoder restores features layer by layer to 64×64, 128×128, and 256×256. In the channel fusion module, these features are fused with the 64×64, 128×128, and 256×256 features from the encoder processed by improved ViT, respectively, to produce a high-resolution inpainted 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 preserving details and structural information during the image reconstruction process. Each layer of the U-Net structure has a jump connection that 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. 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 specific implementation of the fourth stage is 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 various scenarios.

[0039] The model uses a combination of multiple loss functions to comprehensively optimize the quality of image restoration. During training, we used a combination of reconstruction loss, adversarial loss, and other loss functions to enhance the model's ability to restore image details. The loss function formula 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, is the adversarial loss, and its weight λ a =0.1.

[0042] First, the reconstruction loss 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 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 restored image generated by the model, I g is the real target image.

[0045] Secondly, the perceptual loss uses the pre-trained VGG-19 network to extract features (using φ i denoted by ), measures the difference in high-level semantic features between the restored image and the real image. 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 in layer i, Represents the mathematical expectation operation.

[0048] In addition, the model also uses style loss to maintain consistency in the style features of 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. Its 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 present invention

[0058] The core technical feature of this invention lies in combining multiple techniques to improve the traditional ViT-Unet architecture, optimizing it for detail recovery in image restoration tasks. This combined approach effectively improves image feature extraction capabilities, enabling personalized adjustments to image content in different scenarios, ensuring high-quality restoration results.

[0059] Compared with existing patents and papers, this invention significantly improves image restoration while reducing video memory consumption and accelerating training speed, using the same training samples. Experimental results demonstrate that the invention outperforms existing technologies in multiple test scenarios, 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 of the present invention and other mainstream models.

[0062] Figure 3A detailed comparison of the improved ViT module used in this invention and the traditional ViT restoration 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 examples and drawings. The exemplary embodiments of the present invention and their descriptions are only used to explain the present invention and are not intended to limit the present invention.

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

[0065] Step s2: Define a class named CreateDataset to facilitate reading the training set images (CelebA-HQ dataset) and mask data (NVIDIA Irregular Mask dataset) and pre-set preprocessing hyperparameters. First, the make_dataset function parses the image data file path (img_file) and mask data file path (mask_file) respectively, checking whether the image and mask data sizes match. If the mask data size is insufficient, the mask path list is repeatedly expanded to cover the entire image dataset. Next, a function named get_transform is set to define the image preprocessing hyperparameters. The input image is resized to 256×256×3. The image intensities are normalized and linearly mapped to the range [0.0, 1.0]. Random rotations and random horizontal flips are added to enhance the data's transformation diversity. Next, in the getitem method, the image and mask data are loaded by index and the preprocessing transformations defined above are applied to each. 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 inpainted (i.e., the image containing the missing area) is first fed into the network's initial encoder module. The initial feature map is extracted through a 7×7 convolutional layer. After normalization and random affine transformation, a feature map X0 of size 256×256×16 is generated, which is input into the encoder portion of the U-net structure. During the encoding process, the feature map is fed into the encoder for layer-by-layer processing. Each layer consists of a gated convolution and an improved ViT module. The gated convolution dynamically activates valid pixels in the input features at each encoding layer, suppresses invalid features, and enhances the targeted nature of image inpainting. The improved ViT module, on the other hand, leverages its global modeling capabilities to capture long-range dependencies and establishes connections between features through an attention mechanism. First, the feature map X0 is input into the first layer of the encoder, and after gated convolution and improved ViT processing, the feature map X'0 is generated. The feature is output into the channel fusion module in the first layer of the decoder through skip connection for temporary storage. The feature is then downsampled and resized to a feature map X1 of 128×128×32. The feature is input into the second layer, and after gated convolution and improved ViT processing, the feature map X'1 is generated and input into the second layer of the decoder for temporary storage. The feature map X'1 is resized to a feature map X2 of 64×64×64 by downsampling again, and after gated convolution and improved ViT processing, the feature map X'2 is generated and input into the third layer of the decoder for temporary storage. The feature map X3 is resized to 32×32×128 by downsampling. Finally, the feature map X'3 is obtained after gated convolution and improved ViT processing. No downsampling is performed on the last layer.

[0067] The encoded features are upsampled layer by layer through the decoder, restoring them to the original resolution of the input image (256×256×3 pixels). In the decoder, the PixelShuffle upsampling method is combined with an improved ViT module. PixelShuffle upsampling improves image resolution by rearranging channel information while preserving detailed features. A channel fusion module fuses the PixelShuffle upsampled features with features from skip connections at each layer, ensuring that the reconstruction retains both detail and structure. First, the input feature map X'3 is obtained from the encoder. This feature map is processed by the improved ViT module and upsampled by Pixelshuffle, resulting in a resized feature map Y2 of 64×64×64. Y2 is then fused with feature map X'2 through the channel fusion module to generate feature map Y'2. Y'2 is again processed by ViT and upsampled by Pixelshuffle to generate feature map Y1 of 128×128×32. This is then fused with feature map X'1 to generate feature map Y'1. The same steps are followed to generate feature map Y'0 of 256×256×16. Finally, a 7×7 convolutional layer is used to supplement high-frequency signals and compress the resulting feature map channels, converting it into a restored image of 256×256×3, completing the image restoration. During training, the initial learning rate is set to 1e-4 for rapid convergence. Starting from the 10th epoch, the learning rate is reduced by 1e-5 from the previous epoch until it reaches 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 datasets (CelebA-HQ, Places2, Paris StreetView), the differences in the effects of different methods in processing missing image areas are demonstrated. It can be intuitively shown that our model performs more naturally and reasonably when repairing damaged textures of the eyes and mouth, and the lines of the hair are also smoother. This shows that our model performs well in repairing large-area damage.

[0069] based on Figure 3 The 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 comparison area, we can see that the improved ViT module can more clearly restore the striped pattern on the bed, 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, with obvious blurring and ghosting at the edges of the striped pattern, 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, further demonstrating the effectiveness of this method.

[0070] We also conducted a quantitative analysis, as shown in Tables 1 and 2. Our model is quantitatively evaluated on two datasets using the three metrics of FID, PSNR, and SSIM. We compare the results with those of commonly used models when the irregular mask ratio is between 20% and 40% and between 40% and 60%. Experimental results on the CelebA-HQ and Places2 datasets validate the effectiveness of our approach. As shown in the table, our model achieves excellent performance in terms of FID, PSNR, and SSIM. In particular, when the irregular mask ratio is between 20% and 40%, the FID for both datasets is 5.79 and 14.34, respectively, representing decreases of 0.37 and 0.6 compared to the next-best approach, demonstrating that the feature distribution of the images generated by our model is closer to the real image distribution. PSNR improves by 1.61 and 0.83, respectively, compared to the next-best approach, demonstrating that our model generates images with less distortion. On both datasets, SSIM improved by 0.008 and 0.016, respectively, compared to the next-best approach, demonstrating that the images generated by this model possess higher structural similarity. Experimental results on the CelebA-HQ and Places2 datasets validate the effectiveness of this method, demonstrating its leading position in restoration performance.

[0071] Table 1

[0072]

[0073] Table 2

[0074]

[0075] Those skilled in the art will appreciate that the embodiments of the present application can be provided as methods, systems, or computer program products. Therefore, the embodiments of the present application can be implemented in the form of a complete hardware embodiment, a complete software embodiment, or an embodiment combining software and hardware. Furthermore, the embodiments of the present application can be implemented in the form of a computer program product implemented on one or more computer-usable storage media (including but not limited to magnetic 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 produce a machine, so that the instructions executed by the processor of the computer or other programmable data processing device generate instructions for implementing the steps in the process. 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 that can direct a computer or other programmable data processing device to work in a specific manner, so that the instructions stored in the computer readable memory produce an article of manufacture comprising an instruction device, which implements the process Figure 1 a process or multiple processes and / or boxes Figure 1 The 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 operational steps are executed on the computer or other programmable device to produce a computer-implemented process, thereby providing the instructions executed on the computer or other programmable device for implementing the process. Figure 1 a process or multiple processes and / or boxes Figure 1 A step that specifies a function in one or more boxes.

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

[0080] Obviously, those skilled in the art can make various changes and modifications to the embodiments of the present application without departing from the spirit and scope of the embodiments of the present application. Thus, if these modifications and variations of the embodiments of the present application fall within the scope of the claims and their equivalents, 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 effectively integrate local and global image information. 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; In the third stage, the improved ViT module and PixelShuffle are used 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; 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. 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 gate 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-wise 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, forming a structure consisting of two VSS blocks wrapping a cascaded attention layer. 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 introducing the VSS block is as follows: the SS2D module in the VSS layer is used to scan and generate a sequence. The relevance of each token in the sequence effectively determines whether to propagate or discard information. By modeling the visual state space, it provides dynamic modeling capabilities for consecutive image frames and global context. The formula becomes: in, is the i-th visual state space layer, The VSS block is composed of multiple visual state space layers. The sandwich structure is flexible and adaptable, with two VSS blocks providing comprehensive modeling of visual information, weight adjustment, and knowledge fusion at the very beginning of the input. The cascaded attention mechanism then splits and cascades the channels to achieve efficient feature extraction and fusion. 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 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 fully utilize 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 them to the original resolution of the input image. The decoder of each layer is combined with the PixelShuffle upsampling and improved ViT modules, and a channel fusion module is used to integrate the information from the encoder. The decoder restores the features layer by layer to 64×64, 128×128, and 256×256, and fuses them with the 64×64, 128×128, and 256×256 features processed by the improved ViT from the encoder in the channel fusion module, and finally generates a high-resolution repaired image.

2. The image restoration method based on the improved ViT-Unet according to claim 1, characterized in that: The first stage is specifically implemented as follows: 1.1 Acquisition of training dataset: Using data from the CelebA-HQ, Places2, and Paris StreetView datasets as original images, and the PConv mask dataset as the mask dataset; 1.2 Adjust the original image size to 256×256×3 pixels, synthesize it with the PConv mask dataset to form the image set to be repaired, and preprocess it. Use the 7×7 convolution 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 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, thereby retaining details and structural information during image reconstruction. Each layer of the U-Net structure has a skip connection that 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. This skip 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.

4. The image restoration method based on the improved ViT-Unet according to claim 1, characterized in that: The fourth stage is specifically implemented as follows: The model is combined with an enhanced deep learning strategy, and the learning rate is adjusted to 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, is the adversarial loss, and 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 restored image generated by the model, I g is the real target image; Perceptual 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 in layer i, Indicates the mathematical expectation operation; 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 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; 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 real images. The formula is as follows: Where D is the discriminator network, whose output is a scalar, indicating 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