Image inpainting method, device and medium based on mask perception mamba

Abstract

The present application relates to the technical field of image inpainting, and particularly relates to an image inpainting method, device and medium based on mask perception Mamba. In the present application, the encoder and the decoder of the image inpainting model both include MDI modules, the global flow of the MDI modules adopts a mask perception Mamba combined with a geometry-guided unidirectional selective scan (GUSS) strategy, global dependence modeling with linear complexity is realized, and the local flow is extracted by progressive texture feature extraction from coarse to fine, so that the capturing ability for texture details is enhanced.

Description

Technical Field

[0001] This invention relates to the field of image restoration technology, and more specifically to image restoration methods, devices, and media based on mask-aware Mamba. Background Technology

[0002] In the current field of computer vision, State Space Models (SSMs) have gradually become a research hotspot due to their efficient modeling capabilities for long sequence dependencies with linear time complexity, providing a new research path to overcome the computational bottleneck of the traditional Transformer architecture. These models exhibit outstanding efficiency advantages in natural language processing tasks, and subsequent variants further integrate selection mechanisms, enabling dynamic filtering and control of key information based on input data, thereby reducing computational overhead while improving overall modeling efficiency.

[0003] With the continuous development of related technologies, state-space models have been gradually introduced into the field of image inpainting, and various improved architectures for vision tasks have emerged, demonstrating superior computational efficiency compared to traditional Transformers in long-range dependency modeling. Meanwhile, frequency domain analysis has also gradually become an important research direction in image inpainting. Frequency domain transformation can decompose image content into different frequency components, exhibiting significant advantages in texture generation and detail preservation. Related improvement schemes have also achieved remarkable results in scenarios such as repetitive texture generation.

[0004] Existing image inpainting models based on state-space models generally neglect the full utilization of high-frequency information in the frequency domain, making it difficult to meet the actual needs of image inpainting for fine texture restoration. Furthermore, current work on frequency-domain image inpainting mostly introduces frequency-domain processing only in the network architecture or loss function, failing to achieve effective synergy between architecture and loss constraints. This results in insufficient utilization of frequency-domain information, limiting further improvements in image inpainting quality. Summary of the Invention

[0005] To improve image restoration, this invention provides an image restoration method, apparatus, and medium based on mask-aware Mamba.

[0006] A first aspect of the present invention provides an image inpainting method based on mask-aware Mamba, comprising the following steps: The image to be restored is input into the image restoration model, and the image restoration model outputs the image restoration result. The training method of the image restoration model includes the following steps: Obtain the damaged image and its corresponding binary mask; The damaged image and the binary mask are concatenated along the channel dimension and mapped to a high-dimensional feature space to obtain the embedded features; The embedded features are encoded using an encoder; The encoded embedded features are decoded using a decoder to obtain the restored image; Both the encoder and decoder include an MDI (Mamba-based Dual-stream Interaction) module. In the MDI module, the input features are divided into a global stream and a local stream. For the global stream, after mask-aware gating, geometric transformation is performed, followed by processing through an efficient visual scanning module, and then inverse geometric transformation is performed to obtain the global stream features. For the local stream, after activation by an activation function, progressive texture feature extraction from coarse to fine is performed, followed by instance normalization, dynamic feature modulation through a gating mechanism, and local stream features are obtained through residual connections. The global stream features and local stream features are fused to obtain the output features of the MDI module.

[0007] In this invention, both the encoder and decoder of the image restoration model include an MDI module. The global flow of the MDI module adopts a mask-aware Mamba combined with a geometry-guided unidirectional selective scan (GUSS) strategy to achieve global dependency modeling with linear complexity. The local flow enhances the ability to capture texture details through progressive texture feature extraction from coarse to fine.

[0008] In one possible implementation, the efficient visual scanning module is the EfficientSS2D model (efficient one-way selective scanning state space model). In the EfficientSS2D model, the input features are first projected through a linear layer and divided into a content branch and a gated branch. The content branch is then processed sequentially through depthwise separable convolution (DWConv), SiLU activation, S6 selective scanning, and LN layer normalization. Finally, it is multiplied element-wise with the SiLU activation result of the gated branch, and the feature dimension is restored through a linear layer.

[0009] In one possible implementation, the global flow characteristics are obtained by the process shown in the following equation:

[0010] in, Represents global flow characteristics. This indicates the input characteristics of the MDI module. This represents the corresponding binary mask. Represents element-wise product. Representation layer normalization, Represents geometric transformation, This indicates the processing of the EfficientSS2D model. This represents the inverse of a geometric transformation.

[0011] In one possible implementation, the local flow characteristics are obtained by the process shown in the following equation:

[0012] in, Indicates local flow characteristics, This indicates the input characteristics of the MDI module. Represents the activation function of the exponential linear unit. Represents element-wise product. This represents the Sigmoid function. and Represents the features after instance normalization modulation Features obtained by uniformly dividing along the channel, including features after instance normalization modulation. The process shown in the following formula is used to obtain:

[0013] in, , , These represent the features obtained by applying 1×1 convolution, batch normalization, SiLU activation, and uniform partitioning along the channels to the input features of the MDI module. , , This indicates a 3×3 depth-separable multi-scale bottleneck operation. This indicates a 5×5 depth-separable multi-scale bottleneck operation. This indicates a 7×7 depth-separable multi-scale bottleneck operation. Represents a 1×1 convolution. Indicates channel splicing; This indicates instance normalization.

[0014] In one possible implementation, fusing global flow features and local flow features includes: The global flow features are processed by 1×1 convolution and activated by the Sigmoid function to obtain the global gated signal; The local flow features are processed by 1×1 convolution and activated by the Sigmoid function to obtain the local gate control signal; The local flow features are modulated using global gating signals, and the global flow features are modulated using local gating signals. The modulated dual-path features are then spliced ​​along the channel and fused using 1×1 convolution to obtain bidirectional cross-gating fused features.

[0015] In one possible implementation, the MDI module further includes the following steps: The bidirectional cross-gated fusion features are residually connected with the input features, then layer normalization is performed, and then processed through a frequency domain enhanced feedforward network (FDEFN). In the frequency domain enhancement feedforward network, 1×1 convolution, 3×3 channel-wise convolution, effective information filtering using a gating mechanism, and 1×1 convolution are performed sequentially. Then, the network is divided into blocks, and each block is subjected to a 2-Dimensional Real-input Fast Fourier Transform (RFFT2) to enhance high-frequency components in the frequency domain. Finally, a 2-Dimensional Inverse Real-input Fast Fourier Transform is performed to restore each block to a feature map.

[0016] In one embodiment, the encoder further includes a MAD (Mask-Guided Attention Downsampling) module, which processes the output features of the MDI module. In the MAD module, the input feature map is downsampled to obtain a query matrix, which is then mapped to a key matrix and a value matrix via 1×1 convolution. Local window blocks are extracted using a sliding window unpacking operation. The scaled dot product attention is calculated based on the query matrix and the local window blocks. The dot product attention score of the damaged region is set to negative infinity based on the mask feature map corresponding to the input feature map. After Softmax normalization, the aggregated feature map is weighted and aggregated with the local window blocks of the value matrix to obtain the aggregated feature map.

[0017] In one possible implementation, in the training method of the image inpainting model, the parameters of the image inpainting model are updated according to the dynamic high-pass frequency domain loss (DHFL), wherein the dynamic high-pass frequency domain loss is calculated as follows:

[0018] in, Indicates dynamic high-pass frequency domain loss. Represents the set of all frequency points in the frequency domain. Indicates the total number of frequency points. Representing frequency point The Euclidean distance between the restored image and the real image in the high-pass spectrum. This represents the Focal dynamic weight matrix.

[0019] In one possible implementation, the parameters of the image inpainting model are updated based on reconstruction loss, perceptual loss, style loss, and adversarial loss during the training method of the image inpainting model.

[0020] A second aspect of the present invention provides an electronic device including a memory and a processor, the memory storing a computer program, the processor executing the computer program to implement the steps of the mask-aware Mamba-based image restoration method.

[0021] A third aspect of the present invention provides a computer-readable storage medium storing a computer program that, when executed by a processor, implements the steps of the mask-aware Mamba-based image restoration method.

[0022] In this invention, the image inpainting method based on mask-aware Mamba has at least the following beneficial effects: (1) An MDI module is proposed. To address the difficulty in coordinating global structure reconstruction and local texture generation, the MDI module achieves linear complexity global awareness through mask-aware Mamba and GUSS strategies in the global stream, enhances texture capture capabilities through multi-scale cascaded convolutions in the local stream, and utilizes a cross-gating mechanism to achieve complementary fusion of features from both streams. Simultaneously, a frequency domain feedforward network is introduced at the end of the branches to further enhance frequency domain representation, thereby improving the texture realism of the restoration results.

[0023] (2) Design the MAD module. To address the issues of spatial feature misalignment and hole noise leakage caused by traditional convolutional downsampling, the MAD module uses each spatial location in the low-resolution feature map as the query, performs retrieval and dynamically aggregates effective information within the corresponding high-resolution local window in the original resolution feature map, and uses mask injection and gating mechanisms to shield damaged areas and achieve downsampling operations that preserve spatial coherence and semantics.

[0024] (3) Introducing Dynamic High-Pass Frequency Domain Loss (DHFL). To address the problem of blurred edges and loss of high-frequency details in the restoration result caused by pixel-level loss functions, DHFL converts the image to the frequency domain through global 2D-FFT, designs a fixed threshold high-pass filter to suppress low-frequency components, and calculates weighted Euclidean distance based on the Focal mechanism to dynamically adjust the spectral loss weights, guiding the network to effectively achieve high-frequency detail synthesis of texture and edges. Attached Figure Description

[0025] Figure 1 This is a schematic diagram of the overall architecture of an image restoration model in one embodiment of the present invention.

[0026] Figure 2 This is a schematic diagram of the structure of the MDI module in one embodiment of the present invention.

[0027] Figure 3 This is a schematic diagram of the global stream processing process of the MDI module in one embodiment of the present invention, wherein... Figure 3 (a) in the diagram represents the global stream processing procedure. Figure 3 (b) in the diagram represents the processing procedure of the efficient visual scanning module during the global stream processing.

[0028] Figure 4 This is a schematic diagram of the local flow processing process of the MDI module in one embodiment of the present invention.

[0029] Figure 5 This is a schematic diagram of the structure of FDEFN in one embodiment of the present invention.

[0030] Figure 6 This is a schematic diagram of the MAD module in one embodiment of the present invention.

[0031] Figure 7 This is one of the visualization comparisons of the repair results of different models on Places2 (first and second rows), Paris Street View (third and fourth rows), CelebA-HQ (fifth and sixth rows), and Dunhuang Grottoes dataset (seventh and eighth rows) with a mask ratio of (0.3, 0.4].

[0032] Figure 8 The second image shows a comparison of the visualization repair results of different models on Places2 (rows 1 and 2), Paris Street View (rows 3 and 4), CelebA-HQ (rows 5 and 6), and Dunhuang Grottoes dataset (rows 7 and 8) with a mask ratio of (0.3, 0.4].

[0033] Figure 9 The third image shows a comparison of the visualization repair results of different models on Places2 (rows 1 and 2), Paris Street View (rows 3 and 4), CelebA-HQ (rows 5 and 6), and Dunhuang Grottoes dataset (rows 7 and 8) with a mask ratio of (0.3, 0.4]. Detailed Implementation

[0034] To make the objectives, features, and advantages of this invention more apparent and understandable, the technical solutions of the embodiments of this invention will be described below with reference to the accompanying drawings. Obviously, the described embodiments are merely some embodiments of this invention, and not all embodiments. All other embodiments obtained by those skilled in the art based on the embodiments of this invention without creative effort are within the scope of protection of this invention.

[0035] In recent years, deep learning-based image inpainting methods have made significant progress. However, existing image inpainting methods still face key challenges such as insufficient co-modeling of global structure and local texture, the tendency of downsampling to lead to feature misalignment and hole noise leakage, and the limited ability of pixel-level loss functions to constrain high-frequency details.

[0036] To improve image restoration, one embodiment of the present invention provides an image restoration method based on mask-aware Mamba. This method restores damaged images using a trained image restoration model and outputs the restored image.

[0037] It should be noted that the damaged image is the image to be repaired, and its corresponding binary mask can reflect the damaged area of ​​the image. This can be achieved using... Indicates a damaged image. ,in, These represent the height and width of the image, respectively. You can use... Represents a binary mask. Where 1 represents a damaged area and 0 represents an undamaged area. It can be used... This represents the embedded features mapped to a high-dimensional feature space. It can be used... Indicates image repair. .

[0038] like Figure 1 As shown, first, input the damaged image. With binary mask The data is concatenated along the channel dimension and mapped to a high-dimensional feature space through a gated embedding layer. The gated embedding layer uses 3×3 convolutions to generate dual-channel branches, one activated by a Gaussian Error Linear Unit (GELU) and the other as a gated branch. These two branches are multiplied element-wise to initially suppress invalid responses in the masked region. This process preserves the texture features of the effective region while reducing noise interference in the damaged area, thus outputting high-quality embedded features. .

[0039] Subsequently, embedded features XThe data is sequentially processed through three encoder stages, each consisting of a Mamba-based dual-stream interaction (MDI) module and a mask-guided attention downsampling (MAD) module. Specifically, the MDI module first extracts multi-scale features using both mask-aware global and local streams; then, it performs gated fusion through bidirectional interaction units; finally, it introduces a frequency-domain enhanced feedforward network (FDEFN) to enhance the frequency domain representation of the fused features, achieving global structure awareness and local texture enhancement. The global stream employs mask-aware Mamba combined with hard gating and a geometrically transformed one-way selective scanning (GUSS) strategy to suppress hole noise while maintaining global modeling capabilities with linear complexity; the local stream enhances the representation of texture features at different scales through multi-scale cascaded convolutions. Furthermore, the MAD module achieves spatial downsampling through stride convolutions of queries, while dynamically aggregating effective features within local windows using a mask-guided attention mechanism. In the MAD module, the attention score for damaged regions is set to negative infinity, and then its weights are normalized to zero using Softmax, thus preventing features from damaged regions from leaking into the downsampled features. At the end of the encoder, the feature resolution is reduced to 1 / 8 of the input, and deep semantic aggregation is performed through an additional MDI module (bottleneck layer). The decoder has a symmetrical structure with the encoder, containing three levels of upsampling, skip connections, and an MDI module to progressively repair spatial details. During this process, binary masks are used. The feature resolutions were downsampled to 1 / 2, 1 / 4, and 1 / 8 respectively, and then input into the corresponding modules to ensure that the network always accurately perceives the location of the damaged area.

[0040] Finally, the decoder's output features The feature channels are compressed to 3 using a 3×3 convolutional layer, and then activated using the hyperbolic tangent activation function. Activation, with a value range of [-1, 1], is then linearly scaled and mapped to the effective pixel value range of [0, 1] to obtain the final repaired image. The specific operation is shown in the following formula:

[0041] in, This represents the output characteristics of the decoder. This represents a 3×3 convolution. This represents the hyperbolic tangent activation function. This indicates that the image has been repaired.

[0042] The encoder-decoder network described above constitutes the generator. Additionally, a discriminator is set up. During training, the damaged image and a binary mask are input into the generator. By generator Output the restored image; input the real image and the restored image into the discriminator respectively. By the discriminator Output a local confidence map; based on this, construct a least-squares adversarial loss, and combine it with the reconstruction loss, perceptual loss, style loss, and DHFL to alternately update the discriminator. and generator When testing or performing image inpainting, only the trained generator is used. (i.e., image inpainting model), generated by the generator Output a repaired image based on the damaged image.

[0043] The following sections will provide a detailed introduction to each part of the image restoration model.

[0044] (1) Encoder The core task of the encoder is to progressively compress the spatial resolution and extract semantic features while preserving effective region information, and at the same time avoid interference from mask holes in feature learning. The encoder consists of three levels (Level 1 to Level 3), each consisting of... It consists of an MDI module and a MAD module connected in series, where It is a positive integer, and its value ranges from 1 to 4.

[0045] The recursive relationship of encoder levels can be uniformly expressed as follows: Encoder level 1 The input features of the layer are The corresponding binary mask is (when When =1, For gated embedding layer output, (A binary mask for the original input image). For each encoder level, first through... Each MDI module performs feature transformation, followed by a downsampling operation via the MAD module. The overall forward propagation operation can be represented as:

[0046] in, Indicates the first The downsampling features of the layer, where Indicates batch size; , That is, the feature width and height are halved; This means doubling the number of channels to compensate for the loss of spatial resolution; This represents the binary mask passed to the next layer after synchronous downsampling; express The sequential stacking of MDI modules; This represents a mask-guided attention downsampling module. After three levels of encoding by the encoder, the feature resolution is gradually reduced to the input. .

[0047] At the end of the encoder, the deepest features are... Its binary mask Input to The bottleneck layer (fourth level), composed of MDI modules, performs deep semantic aggregation and outputs the final features of the encoder. The entire encoder, through the dual-stream collaborative mechanism of the MDI module and the mask-aware dynamic aggregation strategy of the MAD module, effectively suppresses semantic ambiguity at mask boundaries while protecting the structural integrity of undamaged areas.

[0048] (2) Decoder The decoder aims to extract deep semantic features from the encoder output. The spatial resolution and texture details are gradually restored. The decoder also consists of three levels, symmetrical to the encoder. Each level of the decoder includes an upsampling module, a skip connection (concatenated with the output features of the corresponding encoder level), and... The MDI module effectively fuses spatial details from various encoder levels using skip connections throughout the decoding process. Combined with the dual-stream interaction mechanism of the MDI module, it achieves consistent restoration of semantic features and texture information.

[0049] The upsampling module uses subpixel convolution to double the resolution. First, it expands the number of channels through 3×3 convolution, and then doubles the spatial size through pixel recombination to obtain the upsampled features. Subsequently, the upsampled features The output features corresponding to the encoder layer are concatenated along the channel dimension, and the number of channels is compressed using a 1×1 convolution to match the input dimension of subsequent modules (this compression step is not needed in the lowest resolution decoder layer if the number of channels matches), resulting in the fused features. Finally, the fused features Binary mask of corresponding scale Sent in together Each MDI module performs dual-stream feature interaction and refined feature reconstruction. The above process is shown in the following equation:

[0050] in, This represents the input features of the upsampling module at the current decoder level, i.e., the output or bottleneck features of the previous level encoder; Represents a subpixel convolution operator; Indicates a channel splicing operation; Indicates the encoder's first... Intermediate features of the layer output; This represents a 1×1 convolution operation. This represents a 3×3 convolution operation; This indicates the input features of this layer and the binary mask. They are used together as initial values ​​and then sequentially connected in series within this layer. One MDI module; This represents the output feature of this layer. In the above formula, the 1×1 convolution preprocessing operation is performed only once before entering the first MDI module of this layer, and is used to generate the aggregated input features.

[0051] After three levels of decoding, the feature resolution is successively restored to 1 / 4, 1 / 2, and the original input resolution. The entire decoding process effectively fuses the spatial details of each level of the encoder through skip connections, and combined with the dual-stream interaction mechanism of the MDI module, achieves consistent restoration of semantic features and texture information.

[0052] (3) MDI module like Figure 2 As shown, given input features and the corresponding binary mask The MDI module first captures long-range structural dependencies and multi-scale local details through parallel global and local streams. Then, the two branches' features undergo adaptive interaction and complementary fusion via bidirectional cross-gating. Finally, the fused features are selectively enhanced in the frequency domain using FDEFN, outputting refined features. .

[0053] Global Flow, with mask-aware Mamba as its core, combines hard gating and GUSS strategies to achieve linear complexity global dependency modeling while suppressing hole noise interference, providing reasonable structural priors for large-area damaged regions. For example... Figure 3 As shown in (a) above, for a given input feature Mask-aware gating is used to set the feature responses corresponding to the damaged areas to zero, thereby blocking the further propagation of void noise, and layer normalization is performed to stabilize the training process. Specifically, firstly, based on the current stacking index of the MDI module... (The range of values ​​is) , (The total number of MDI modules in the generator), and apply specific geometric transformations (including spatial flipping or transpose) to the layer-normalized features to simulate multi-directional sequence scanning. Subsequently, the geometrically transformed features... The input is fed into the EfficientSS2D model (efficient visual scanning module), and global context inference is performed along the flattened spatial sequence. After scanning, the original spatial arrangement is restored through corresponding inverse transformations to obtain global flow features. The above process is shown in the following formula:

[0054] in, Indicates input features, This represents a binary mask corresponding to the scale. Represents element-wise product. Representation layer normalization, This indicates that the GUSS strategy is based on the stacked index. Performed geometric transformations (flip or transpose). This represents the inverse transform of the corresponding geometric transformation. This represents the EfficientSS2D model.

[0055] like Figure 3 As shown in (b), the EfficientSS2D model performs unidirectional selective scanning on a one-dimensional sequence after GUSS transformation. The input features are first projected through a linear layer and divided into content branches and gated branches. The content branches are sequentially processed by depthwise separable convolution (DWConv), SiLU activation, and sequence flattening. Selective scanning parameters are then generated by the linear mapping, and state space recursion is completed through S6 selective scanning. The scan output is normalized by an LN layer, multiplied element-wise with the SiLU activation result of the gated branch, and finally restored to the original feature dimension through a linear layer before undergoing inverse GUSS transformation. The EfficientSS2D model uses only a single set of SSM parameters, combined with the geometric transformation of GUSS, to achieve multi-directional global context aggregation.

[0056] This design enables the global flow to decompose the multi-directional scanning process into cross-layer progressive sub-steps by leveraging the geometric transformation diversity of different layer indices, using only a single set of SSM parameters. Furthermore, it achieves the aggregation of contextual information in different directions through inter-layer residual connections, effectively improving the global structure awareness capability while reducing the number of parameters.

[0057] Local Flow focuses on capturing multi-scale texture details and enhances the local consistency of binary mask boundary regions. Unlike Global Flow, which emphasizes structural priors, Local Flow employs multi-scale cascaded convolutions, aggregating texture information from different receptive fields through heterogeneous convolutional kernels of 3×3, 5×5, and 7×7, thereby improving its ability to capture both microscopic and macroscopic details. For example... Figure 4 As shown, firstly, the input features are... With binary mask The concatenation is performed along the channel dimension, and then compressed back to the original number of channels using a 1×1 convolution. Batch normalization and SiLU activation are performed to obtain the initial embedding features. Subsequently, these embedding features are uniformly divided into three groups along the channels. , , The number of channels is as follows: , , The features are input into depthwise separable convolutional branches with kernel sizes of 3×3, 5×5, and 7×7, respectively. To enhance the effective fusion of multi-scale features, a cascaded additive structure is adopted between branches, where the input of the large convolutional kernel branch is the superposition of the output of the previous branch and the features of the current group, thereby achieving progressive texture extraction from coarse to fine. The three output features are concatenated and then modulated using 1×1 convolution and instance normalization. The multi-scale feature extraction and fusion process is shown in the following equation:

[0058] in, This represents the bottleneck transform unit, which is formed by expanding the channels using a 1×1 convolution. Bottleneck transform unit consisting of channel-wise convolution and 1×1 convolution projection; This represents the output of the first path (3×3 depthwise separable convolution). ; This represents the output of the second path (5×5 depthwise separable convolution). ; Indicates channel splicing; This indicates instance normalization.

[0059] Finally, the features after instance normalization modulation are... Divide the channel evenly into content features With gate weights The system consists of two parts, which use a gating mechanism to achieve dynamic feature modulation and are then compared with the input features. The residuals are summed to obtain the local flow characteristics. As shown in the following formula:

[0060] in, Represents the activation function of the exponential linear unit. This represents the Sigmoid function. This gating mechanism can adaptively suppress background noise and enhance the response intensity of textured regions, enabling local flow to effectively improve the ability to repair complex textures while maintaining computational efficiency.

[0061] To promote the complementary fusion of global structural information and local texture details, a bidirectional cross-gated fusion unit is designed as a bridge connecting the global and local flows. Unlike traditional simple addition or concatenation operations, this mechanism enables the two branches to adaptively generate gating signals based on each other's content, selectively modulating their own features, thereby achieving bidirectional guidance between structural priors and texture features. Specifically, such as... Figure 2 As shown, firstly, global flow features With local flow characteristics Gated signals that are modulated with each other are generated by 1×1 convolution and the sigmoid function, respectively. Used to modulate local flow characteristics This is used to modulate global flow features. Subsequently, the modulated dual-path features are concatenated along the channels and fused using a 1×1 convolution to obtain bidirectional cross-gated fused features. As shown in the following formula:

[0062] in, Represents element-wise product. This indicates a channel splicing operation. This indicates a bidirectional cross-gating fusion feature.

[0063] In the MDI module, the global stream guides the structural layout of the local stream, while the local stream supplements the texture details of the global stream; the gating signal is dynamically generated by the content of the other branch, rather than being fixed or modulated in a single direction. Therefore, the MDI module can adaptively adjust its contribution weight according to the content of each other, thereby maintaining structural coherence in the damaged area while improving texture richness.

[0064] To further enhance the high-frequency detail representation capability of the fused features, the MDI module introduces FDEFN at the end. FDEFN, through explicit frequency domain modeling, compensates for the insufficient representation of high-frequency components in the spatial domain, enabling the network to generate restoration results with sharper edges and richer textures. Bidirectional cross-gated fusion features. Input features Perform residual connections and use a learnable scaling factor. (Initially set to 0 to stabilize the initial training phase) Adjust the fusion strength. Then, apply layer normalization to the superimposed features and input them into FDEFN for processing. Specifically, as... Figure 5 As shown, within FDEFN, channel mixing and spatial neighborhood aggregation are first performed using 1×1 convolution and 3×3 channel-wise convolution, and a gating mechanism is used to filter effective information. Next, the feature map is divided into 8×8 local blocks, and a two-dimensional real-number fast Fourier transform (RFFT2) is performed on each block. Learnable weighted enhancement is applied in the frequency domain to enhance high-frequency detail representation. Finally, the spatial domain is restored using a two-dimensional inverse real-number fast Fourier transform (IRFFT2). This process is illustrated in the following equation:

[0065] in, This indicates FDEFN processing. The final output characteristics of FDEFN. Low-frequency core information is preserved through dual residual connections, while high-frequency details enhanced by FDEFN are superimposed.

[0066] FDEFN achieves efficient global frequency domain awareness through block-based RFFT2 without altering the spatial resolution of features. Simultaneously, the learnable frequency domain mask within this module adaptively adjusts the enhancement level for different frequency components based on the repair task, enabling the network to suppress low-frequency noise while strengthening the high-frequency components corresponding to edges and textures. Through explicit frequency domain modeling, this module effectively compensates for the smoothing effect of pixel-level loss functions on high-frequency textures, guiding the network to generate sharper boundaries and richer texture details.

[0067] (4) MAD module The MAD module uses each spatial location in the low-resolution query feature map obtained by convolutional downsampling as a query vector, and extracts the corresponding location from the original resolution feature map for that location. A local window serves as the space carrying keys and values. Subsequently, through mask injection and gating mechanisms, invalid noise is effectively prevented from propagating deeper, thus achieving reliable downsampling operations. For example... Figure 6 As shown, given the encoder's first... Intermediate features of the layer and the corresponding binary mask In the MAD module, intermediate features are first processed by a 3×3 convolution with a stride of 2. Downsampling yields a low-resolution query matrix. ( , Meanwhile, intermediate features Mapped to a key matrix via 1×1 convolution AND-value matrix And through a step size of 2 Slide window expand operation to extract local window block , Similarly, binary masks Also launched simultaneously Subsequently, the low-resolution query matrix was... Each spatial location is used as a query vector, and in its corresponding... Within a local window and Calculate the scaled dot product attention and utilize Set the attention score corresponding to the damaged area to 0. After Softmax normalization, Perform weighted aggregation. Aggregated feature map. First, gated weights are generated using a 1×1 convolution and a sigmoid function, then the number of channels is expanded using another 1×1 convolution. to obtain downsampling features The process is shown in the following formula:

[0068]

[0069] in, This represents the attention score for the scaled dot product. Indicates an indicator function, This represents matrix multiplication. The attention weights are multiplied by the expanded feature values, and then... Dimension summation aggregates features within a local window to generate gating weights, thereby filtering out effective information.

[0070] Meanwhile, to ensure that subsequent layers of the encoder can continuously sense the damaged area, the MAD module uses a binary mask. Perform a max pooling operation with a step size of 2 to obtain a binary mask. As shown in the following formula:

[0071] in, This represents 2×2 max pooling with a step size of 2.

[0072] The MAD module ensures that as long as there are holes in the local window, the corresponding position after downsampling is still marked as broken, effectively preventing excessive shrinkage of the mask after multiple downsamplings and guaranteeing the robust transmission of mask information. In summary, the MAD module, through content-aware attention aggregation and an effective design of robust mask downsampling, effectively preserves key spatial information while reducing feature resolution and suppressing noise propagation, providing higher-quality feature input for semantic modeling in subsequent encoder layers.

[0073] (5) Loss function During the training of the image inpainting model, to guide the model to explicitly focus on high-frequency components in the frequency domain, DHFL is used in conjunction with reconstruction loss, perceptual loss, style loss, and adversarial loss to jointly guide model optimization. DHFL will restore the image... With real images The network is transformed to the frequency domain using a standard two-dimensional Fast Fourier Transform (FFT2). A dynamic high-pass filter is then used to suppress low-frequency flat regions, thereby highlighting high-frequency components. Simultaneously, a Focal weighting mechanism based on spectral error dynamically adjusts the loss contribution of each frequency component, guiding the network to focus on high-frequency textures and structural edges that are difficult to reconstruct.

[0074] Specifically, the first step is to repair the image. With real images The frequencies are transformed to the frequency domain using FFT2 to obtain the corresponding complex spectra. For ease of subsequent calculations, a real-valued spectrum representation function is defined. By stacking the real and imaginary parts of the complex spectrum along the channel dimension, we obtain the real-valued spectrum representation, as shown in the following equation:

[0075] in, Representing an image The real-valued spectrum of ; The input image is represented by the repaired image, which is then substituted into the input image when calculating the loss. With real images ; This represents the standard two-dimensional Fast Fourier Transform; This indicates a stacking operation along the channel dimension; and These represent the operations of extracting the real part and the imaginary part, respectively.

[0076] The design threshold is A high-pass filter, based on each frequency point to DC component normalized distance Generate a high-pass binary mask As shown in the following formula:

[0077]

[0078] in, This indicates a minimum value operation, which compares each element and selects the smaller value. Indicates the frequency index in the horizontal direction. Indicates the frequency index in the vertical direction; and To handle the periodicity of FFT2, the frequency point is taken to the DC component. The shortest equivalent distance; This represents the equivalent distance from the frequency point to the DC component. Indicates the high-pass filter threshold; Represents a high-pass binary mask; 1 indicates when At this frequency, the frequency point is far from the DC component and belongs to high-frequency retention; 0 indicates that when At this frequency, the frequency point is relatively close to the DC component, which is considered low-frequency suppression.

[0079] Then, the high-pass binary mask With real-valued spectrum Element-wise multiplication is performed to filter out low-frequency components, resulting in the high-pass filtered spectrum. Subsequently, the high-pass filter spectrum of the restored image was calculated. High-pass filter spectrum of real image Euclidean distance at each frequency point :

[0080] in, This represents the real-valued spectrum after high-pass filtering, with subscripts indicating the frequency response. , These represent the real and imaginary parts of the spectrum, respectively. Normalization and exponentiation are performed to construct a dynamic spectrum weight matrix. :

[0081] in, This indicates the operation of taking the maximum value. This dynamic spectrum weight matrix... and The positive correlation means that frequencies with larger errors have higher weights in the loss, thus guiding the network to focus on optimizing frequency components with poor reconstruction quality and reducing the influence of well-reconstructed components. (Parameters) The degree of emphasis on large error frequency components is controlled. Finally, DHFL is defined as the mean of the weighted spectral distance after high-pass filtering, i.e.:

[0082] in, Indicates dynamic high-pass frequency domain loss. Represents the set of all frequency points in the frequency domain. This represents the total number of frequency points. DHFL, through the synergistic effect of high-pass filtering and Focal weighting, effectively guides the network to prioritize the recovery of high-frequency texture details that are difficult to reconstruct, complementing the pixel spatial loss.

[0083] In addition to DHFL, reconstruction loss, perceptual loss, style loss, and adversarial loss are also used to guide the training process of the image inpainting model.

[0084] The reconstruction loss is used to constrain the pixel-by-pixel consistency between the restored image and the real image in the RGB space. The reconstruction loss employed is... Defined as:

[0085] in, Indicates image repair; Represents a real image (undamaged image); Indicates the input binary mask; This represents the average value of the absolute error for each element. This represents the mask mean, used to normalize the loss, where , Represents a binary mask The Middle Line 1 The elements of the column.

[0086] Perceptual loss is used to enhance the semantic consistency of the generated images, and a pre-trained VGG-19 network is used to extract features of the restored image and the real image at specific layers. The features of the restored image and the real image at corresponding feature layers are then calculated. Distance. Perceptual loss used. Defined as:

[0087] in, Indicates the first in the VGG-19 network Activation features of layers; sets It contains multi-level features ranging from shallow texture to deep semantics; Indicates the first i Layer weights.

[0088] Style loss is used to maintain style consistency between the repaired and undamaged regions, and is achieved by calculating the feature Gram matrix of each layer. Distance is used to measure the overall texture statistics of an image. A style loss method is employed. Defined as:

[0089] in, This represents the function for calculating the Gram matrix. We selected mid-to-high-level features in VGG-19 that emphasize style expression.

[0090] Adversarial loss is used to enhance the realism of the restored image; the discriminator is trained using least-squares adversarial loss. With generator Discriminator The PatchGAN architecture is used to output a two-dimensional confidence map to capture the fidelity of local image patches, and the discriminator loss... Defined as:

[0091] Generators against loss Defined as:

[0092] in, Indicates the discriminator, Indicates in real image Take the expected value. This indicates the output image of the generator. Take the expected value.

[0093] Total loss of image inpainting model Defined as the weighted sum of the losses mentioned above, as shown in the following formula:

[0094] in, These represent the weighting coefficients for each loss, and their values ​​are primarily set based on their magnitude in the total loss and their importance in the optimization. In this embodiment of the invention, =25, =0.5, =175, =0.012, =10. This total loss function effectively enhances the texture details and realism of the generated image while ensuring the accuracy of the repaired pixels.

[0095] To evaluate the effectiveness and generalization of the image inpainting model in this embodiment, extensive experiments and analyses were conducted on four publicly available image inpainting datasets (Places2, CelebA-HQ, Paris Street View, and Dunhuang Grottoes). PSNR (Peak Signal-to-Noise Ratio), SSIM (Structural Similarity Index Measure), and MAE (Mean Absolute Error) were used as evaluation metrics. During the training phase, the batch size (Batch Size, ...) was... The value is set to 2, using the Adam optimizer, whose first-order moment decay coefficient is... β 1 = 0.5, second-order moment attenuation coefficient β 2 = 0.999, initial learning rate set to 5 × 10 -5 All input images were resized to 256×256 pixels.

[0096] To determine the impact of the state dimension in Mamba on image inpainting performance, a controlled experiment was designed on the Places2 dataset to optimize the state dimension parameters. In the image inpainting model of this embodiment, This represents the hidden state dimension of the SSM within the MDI module, and its size directly determines the amount of contextual information the model can remember during long program sequence scans. Therefore, under identical training conditions, three representative sets were compared on the Places2 dataset. The values ​​are shown in Table 1.

[0097] Table 1 Different SSM state dimensions ( Image inpainting performance comparison on the Places2 dataset (bold highlighting of the best metric).

[0098] As shown in Table 1, when processing sequences of length 65,536 (i.e., 256×256), excessively high state dimensions can easily cause numerical instability in the state transition matrix, resulting in a small number of abnormal pixel outputs and negatively impacting the PSNR, which is sensitive to mean square error. However, neither SSIM nor MAE showed significant degradation. Therefore, based on a trade-off between training stability and restoration accuracy, the final choice was... =32 is used as the configuration for subsequent experiments.

[0099] To determine the optimal number of MDI modules at each level of the encoder, further comparative experiments were conducted. The number of MDI modules at the four levels of the encoder is denoted as […]. N 1, N 2, N 3, N [4], and under the premise of keeping the other conditions consistent on the Places2 dataset, the experimental results of six representative configurations were compared, as shown in Table 2.

[0100] Table 2 Comparison of repair performance on the Places2 dataset with different MDI module count configurations (optimal metrics are highlighted in bold).

[0101] As shown in Table 2, comparing configurations [1, 1, 1, 1] and [1, 1, 1, 2], increasing the number of MDI modules in the bottleneck layer (fourth level) from 1 to 2 improves the PSNR by 0.04 dB, but further increasing it to 3 results in a performance decrease. This indicates that deepening the bottleneck layer is beneficial for global structure modeling, but excessive stacking can lead to feature redundancy. Further comparing [1, 1, 2, 2] and [1, 2, 1, 2], with similar parameter counts, allocating additional MDI modules to the middle layer (second level) and deeper layers (third level) is more effective. This suggests that dense interaction on the 1 / 2 scale feature map is more effective in coordinating texture details and structural priors. Ultimately, the configuration [1, 2, 1, 2] achieved the best performance (PSNR of 35.27dB, SSIM of 0.9684, and MAE of 0.366%). Therefore, [1, 2, 1, 2] was used as the default number configuration for the MDI modules in subsequent experiments.

[0102] To verify the effectiveness and superiority of the image inpainting model in this embodiment, it was compared with various current mainstream image inpainting models (SEM-Net, APLRL, TSGDAM, LSG-Net, TSBGNet, HINT, C2FGAN, WEMamba-MDFL) on six mask intervals (from the low interval (0.01, 0.1] to the high interval (0.5, 0.6]). The results are shown in Tables 3-6. All experiments were conducted in a unified hardware and software environment: the operating system was 64-bit Ubuntu 22.04.5 LTS, the deep learning framework was PyTorch 2.2.0 (CUDA 12.1), the GPU was NVIDIA GeForce RTX 4090, and the Python version was 3.10.19.

[0103] Table 3 shows the quantitative comparison results of different models on the Places2 dataset with different mask rates. ↑ indicates higher is better, ↓ indicates lower is better, and the optimal result is marked in bold.

[0104]

[0105] Table 3 shows the performance differences between the image inpainting model of this invention and various comparative models at different mask rates on the Places2 dataset. It can be seen that the image inpainting model of this invention achieves optimal results in all ranges, and its advantage becomes more pronounced as the mask rate increases.

[0106] Table 4. Quantitative comparison results of different models on the Paris Street View dataset with different mask ratios. ↑ indicates higher is better, ↓ indicates lower is better, and the optimal result is marked in bold.

[0107]

[0108] As shown in Table 4, on the Paris Street View dataset, the image restoration model of this invention achieves the best results in PSNR, SSIM and MAE metrics, demonstrating stable and superior restoration performance.

[0109] Table 5. Quantitative comparison results of different models on the CelebA-HQ dataset with different mask ratios. ↑ indicates higher is better, ↓ indicates lower is better, and the optimal result is marked in bold.

[0110]

[0111] Table 5 presents the quantitative comparison results of the image inpainting model of this invention with various contrast models under six different mask ratios on the CelebA-HQ dataset. It can be seen that the image inpainting model of this invention achieves optimal performance across all mask ranges, especially demonstrating a significant and stable advantage when facing the highly structured priors and fine texture requirements unique to face images.

[0112] Table 6. Quantitative comparison results of different models on the Dunhuang Grottoes dataset with different mask rates. ↑ indicates higher is better, ↓ indicates lower is better, and the optimal result is marked in bold.

[0113]

[0114] Unlike natural scenes and facial images, Dunhuang murals possess complex texture degradation patterns and highly structured characteristics. Experimental results on the Dunhuang Grottoes dataset demonstrate that even with slight deficiencies in extremely small damaged areas, the image restoration model of this invention still exhibits significant advantages in pixel-level accuracy under conditions of large-area defects. This is of great importance for the digital restoration of large-area peeling areas in the murals.

[0115] To further verify the effectiveness of the image restoration model of this invention in image restoration tasks from a visual perception perspective, Figures 7-9 The visualization compares the restoration results of the image restoration model of this embodiment and the contrasting model within the mask interval (0.3, 0.4]. It can be seen that the image restoration model of this embodiment exhibits significant image restoration performance across various image types, including natural scenes, architectural street scenes, facial semantics, and artistic murals. The image restoration model of this embodiment not only generates restoration results with fewer artifacts, more reasonable structure, and more realistic textures from a visual perspective, but also fully demonstrates the robustness and effectiveness of the image restoration model in handling large-area damaged image restoration tasks.

[0116] To systematically validate the effectiveness of the designed core components, ablation experiments were conducted on the Dunhuang Grottoes dataset. Using the complete model as a baseline, nine variants were designed, and the contribution of each component to the repair performance was evaluated by progressively removing or replacing key modules. The nine variants are designed as follows: 1. w / ConvDown: Replace the MAD module with a standard 3×3 convolution with a stride of 2 to verify the effectiveness of the MAD module; 2. w / o DHFL: Remove DHFL, retaining only reconstruction loss, perceptual loss, style loss, and adversarial loss; 3. w / o Local: Remove the local flow branch in the MDI module and retain only the global flow branch to verify the contribution of local texture modeling; 4. w / o Global: Remove the global flow branch in the MDI module and keep only the local flow branch to verify the necessity of global awareness; 5. w / o FDEFN: Remove FDEFN from the MDI module to verify the effect of frequency domain feature enhancement; 6. w / o GUSS: Disable the GUSS strategy, that is, cancel the geometric flip and transpose operation of the input image, and only retain the unidirectional scan along the spatial unfolding sequence to verify the impact of multi-directional state space modeling on the repair performance. 7. w / Add: Replace the cross-gating interactive operation in the MDI module with a simple element-wise addition of feature maps to verify the superiority of the gating fusion mechanism; 8. w / o Mask: Disables masking in the attention calculation of the MAD module (attention score is set to...). This allows the model to aggregate invalid information from the damaged area to verify the necessity of masking; 9. w / o HPF: In DHFL, the frequency domain loss is preserved but the high-pass filter is disabled, and only the original frequency domain is constrained without emphasizing high-frequency components, to verify the effect of high-frequency enhancement.

[0117] The results of the ablation experiments are shown in Table 7. It can be seen that the complete model achieved the best performance across all mask ratios, and the performance degradation trend of each variant provides strong evidence for the effectiveness of each component.

[0118] Table 7 compares the ablation experimental results on the Dunhuang Grottoes dataset. ↑ indicates higher is better, ↓ indicates lower is better, and bold indicates optimal.

[0119]

[0120] Those skilled in the art will understand that all or part of the processes in the above methods can be implemented by a computer program instructing related hardware. The computer program can be stored in a computer-readable storage medium, and when executed, it can include the processes of the above methods. The computer-readable storage medium can include non-volatile and / or volatile memory. Non-volatile memory can include read-only memory (ROM), programmable ROM (PROM), electrically programmable ROM (EPROM), electrically erasable programmable ROM (EEPROM), or flash memory. Volatile memory can include random access memory (RAM) or external cache memory.

[0121] The above description is merely a specific embodiment of the present invention, but the scope of protection of the present invention is not limited thereto. Any variations or substitutions that can be easily conceived by those skilled in the art within the technical scope disclosed in the present invention should be included within the scope of protection of the present invention. Therefore, the scope of protection of the present invention should be determined by the scope of the claims.

Claims

1. An image inpainting method based on mask-aware Mamba, characterized in that, Includes the following steps: The image to be restored is input into the image restoration model, and the image restoration model outputs the image restoration result. The training method of the image restoration model includes the following steps: Obtain the damaged image and its corresponding binary mask; The damaged image and the binary mask are concatenated along the channel dimension and mapped to a high-dimensional feature space to obtain the embedded features; The embedded features are encoded using an encoder; The encoded embedded features are decoded using a decoder to obtain the restored image; Both the encoder and decoder include an MDI module. In the MDI module, the input features are divided into a global stream and a local stream. For the global stream, after mask-aware gating, a geometric transformation is performed, followed by processing through an efficient visual scanning module, and then an inverse geometric transformation is performed to obtain the global stream features. For the local stream, after activation by an activation function, progressive texture feature extraction from coarse to fine is performed, followed by instance normalization, dynamic feature modulation through a gating mechanism, and local stream features are obtained through residual connections. The global stream features and local stream features are fused to obtain the output features of the MDI module.

2. The method according to claim 1, characterized in that, The efficient visual scanning module is the EfficientSS2D model. In the EfficientSS2D model, the input features are first projected through a linear layer and divided into a content branch and a gated branch. The content branch is then processed sequentially through depthwise separable convolution, SiLU activation, S6 selective scanning, and LN layer normalization. Finally, it is multiplied element-wise with the SiLU activation result of the gated branch, and the feature dimension is restored through a linear layer.

3. The method according to claim 2, characterized in that, The global flow characteristics are obtained through the process shown in the following equation: ； in, Represents global flow characteristics. This indicates the input characteristics of the MDI module. This represents the corresponding binary mask. Represents element-wise product. Representation layer normalization, Represents geometric transformation, This indicates the processing of the EfficientSS2D model. This represents the inverse transformation of a geometric transformation.

4. The method according to claim 3, characterized in that, The local flow characteristics are obtained by the process shown in the following equation: ； in, Indicates local flow characteristics, This indicates the input characteristics of the MDI module. Represents the activation function of the exponential linear unit. Represents element-wise product. This represents the Sigmoid function. and Represents the features after instance normalization modulation Features obtained by uniformly dividing along the channel, including features after instance normalization modulation. The process shown in the following formula is used to obtain: ； in, , , These represent the features obtained by applying 1×1 convolution, batch normalization, SiLU activation, and uniform partitioning along the channels to the input features of the MDI module. , , This indicates a 3×3 depth-separable multi-scale bottleneck operation. This indicates a 5×5 depth-separable multi-scale bottleneck operation. This indicates a 7×7 depth-separable multi-scale bottleneck operation. Represents a 1×1 convolution. Indicates channel splicing; This indicates instance normalization.

5. The method according to claim 4, characterized in that, The fusion of global flow features and local flow features includes: The global flow features are processed by 1×1 convolution and activated by the Sigmoid function to obtain the global gated signal; The local flow features are processed by 1×1 convolution and activated by the Sigmoid function to obtain the local gate control signal; The local flow features are modulated using global gating signals, and the global flow features are modulated using local gating signals. The modulated dual-path features are then spliced ​​along the channel and fused using 1×1 convolution to obtain bidirectional cross-gating fused features.

6. The method according to claim 1, characterized in that, The MDI module also includes the following steps: The bidirectional cross-gated fusion features are residually connected with the input features, then layer normalization is performed, and then processed through a frequency domain enhanced feedforward network. In the frequency domain enhancement feedforward network, 1×1 convolution, 3×3 channel-wise convolution, effective information filtering using a gating mechanism, and 1×1 convolution are performed sequentially. Then, the network is divided into blocks, and a two-dimensional real-number fast Fourier transform is performed on each block to enhance the high-frequency components in the frequency domain. Finally, a two-dimensional inverse real-number fast Fourier transform is performed to restore each block to a feature map.

7. The method according to claim 1, characterized in that, The encoder also includes a MAD module, which processes the output features of the MDI module. In the MAD module, the input feature map is downsampled to obtain a query matrix, and then mapped to a key matrix and a value matrix through 1×1 convolution. Local window blocks are extracted using a sliding window unpacking operation. The scaled dot product attention is calculated based on the query matrix and the local window blocks. The dot product attention score of the damaged region is set to negative infinity based on the mask feature map corresponding to the input feature map. After Softmax normalization, it is weighted and aggregated with the local window blocks of the value matrix to obtain the aggregated feature map.

8. The method according to claim 1, characterized in that, In the training method of the image inpainting model, the parameters of the image inpainting model are updated according to the dynamic high-pass frequency domain loss, where the dynamic high-pass frequency domain loss is calculated as follows: ； in, Indicates dynamic high-pass frequency domain loss. Represents the set of all frequency points in the frequency domain. Indicates the total number of frequency points. Representing frequency point The Euclidean distance between the restored image and the real image in the high-pass spectrum. This represents the Focal dynamic weight matrix.

9. An electronic device, characterized in that, It includes a memory and a processor, the memory storing a computer program, and the processor executing the computer program to implement the steps of the image inpainting method based on mask-aware Mamba as described in any one of claims 1-8.

10. A computer-readable storage medium, characterized in that, The device contains a computer program that, when executed by a processor, implements the steps of the image inpainting method based on mask-aware Mamba as described in any one of claims 1-8.

Images