An image restoration method based on state space model

Abstract

This invention discloses an image restoration method based on a state-space model, comprising: a shallow feature extraction stage: extracting shallow features of the image to be restored through convolutional layers; a deep feature extraction stage: extracting deep features of the image to be restored based on the shallow features of the image to be restored using stacked residual Mamba group (RMG), nested sigma scanning strategy, and sequence shuffling attention module; and a reconstruction stage: based on the shallow and deep features of the image to be restored, introducing a pixel shuffling layer and a 3x3 convolutional layer for high-resolution image reconstruction, using a nested sigma scanning strategy to preserve the locality and continuity of the image, and using a sequence shuffling attention module to effectively aggregate contextual information of sequences in different directions. This method demonstrates excellent performance in various tasks such as super-resolution, denoising, deblurring, and dehazing, providing a new approach for efficient image restoration.

Description

Technical Field

[0001] This application relates to the field of image processing, and in particular to an image restoration method based on a state-space model. Background Technology

[0002] With the rapid development of digital image processing technology, image restoration has become an important research direction in the field of computer vision. Image restoration technology aims to recover high-quality clean images from given degraded images (such as noisy, blurred, or hazy images), and it has wide applications in fields such as medical image analysis, remote sensing image processing, and autonomous driving vision systems. In recent years, methods based on convolutional neural networks (CNNs) and Transformers have made significant progress in image restoration because they can effectively capture the inherent locality (i.e., fine-grained patterns and correlations within small regions) and continuity (i.e., smooth, gradual transitions within larger regions) in natural images. However, these methods still have limitations in capturing long-range dependencies. CNNs, due to their fixed-size receptive field, struggle to effectively model global relationships in images, while Transformer-based methods are limited by computational complexity and local window segmentation strategies, resulting in shortcomings in long-range modeling and limiting their application in a wider range of scenarios. Therefore, developing a method that can capture long-range dependencies while preserving the inherent locality and continuity of natural images is highly desirable.

[0003] State-space models (SSMs) are an emerging sequence modeling approach that efficiently models long sequence dependencies with near-linear complexity through selective scan operations. Mamba, a recent state-space model, has been widely adopted in the computer vision field in recent years due to its superior performance. However, since the core algorithm of state-space models, selective scan operations, is essentially designed for one-dimensional sequences, it cannot be directly applied to processing two-dimensional images. To address this issue, existing Mamba-based image restoration methods (such as MambaIR) typically involve a three-step process:

[0004] 1) Flatten a two-dimensional image into multiple one-dimensional sequences;

[0005] 2) Use selective scan operations to process each sequence independently;

[0006] 3) Aggregate the processed sequences to form the output two-dimensional image.

[0007] However, this paradigm still faces two drawbacks when processing images: First, when converting images into sequences, the images are unfolded into sequences row by row or column by column, which destroys the inherent locality and continuity of the images; second, it usually aggregates the processed sequences by summing at the pixel level, ignoring the different contexts between sequences unfolded in completely different ways. Summary of the Invention

[0008] To address the aforementioned shortcomings in existing technologies, this invention provides an image restoration method based on a state-space model. This method utilizes a nested S-shaped scanning strategy to preserve the locality and continuity of the image, and employs a sequence shuffling attention module to effectively aggregate contextual information from sequences in different directions. This solves the problem that existing unfolding methods neglect image locality and continuity during the unfolding stage of the visual state-space model.

[0009] To achieve the aforementioned objectives, the present invention employs the following technical solution: an image restoration method based on a state-space model, comprising:

[0010] Shallow feature extraction stage: shallow features of the image to be restored are extracted through convolutional layers;

[0011] Deep feature extraction stage: Based on the shallow features of the image to be restored, deep features of the image to be restored are extracted using stacked residual Mamba group (RMG), nested S-shaped scanning strategy and sequence mixing attention module.

[0012] Reconstruction stage: Based on the shallow and deep features of the image to be restored, a pixel shuffling layer and a 3x3 convolutional layer are introduced to reconstruct the high-resolution image.

[0013] Furthermore, each residual Mamba group RMG consists of several sequentially connected residual Mamba blocks RMB and a convolutional layer.

[0014] Furthermore, each residual Mamba block (RMB) includes a visual Mamba module (VNM).

[0015] Furthermore, each visual Mamba module (VNM) includes a MaIRM.

[0016] Furthermore, MaIRM processes the input two-dimensional features as follows:

[0017] The two-dimensional features are flattened into four one-dimensional sequences along four different directions by using a nested S-shaped scanning strategy;

[0018] Use selective scan operations to capture long-range dependencies of four one-dimensional sequences;

[0019] The output features are obtained by aggregating the one-dimensional sequence after capturing the dependencies through the sequence fusion attention module.

[0020] Furthermore, the nested S-shaped scanning strategy is used to extract sequences that maintain locality and continuity from input features. Specifically, the features are divided into multiple non-overlapping strips, and S-shaped scanning paths are used within and between the strips.

[0021] Furthermore, the nested S-shaped scan strategy also includes a strip sliding mechanism to maintain locality and continuity in the boundary region between adjacent strips; specifically including:

[0022] For two consecutive modules, the features of the first module are divided into non-overlapping strips, and a shift strip operation is used to set the width of the first and last strips of the first module to half of their original width, while the width of the other strips remains unchanged.

[0023] Furthermore, methods for using sequence hybrid attention modules to aggregate and capture one-dimensional sequences with captured dependencies to obtain output features include:

[0024] The one-dimensional sequence after capturing dependencies is spatially averaged and then concatenated to obtain the pooled sequence.

[0025] The pooled sequences are grouped, and the grouped sequences are shuffled and rearranged to obtain the rearranged sequences.

[0026] The channel attention weights of the rearranged sequence are obtained by grouping convolutions, and the rearranged sequence is rearranged back to the original order to obtain the channel attention weights of the original order.

[0027] The output feature is obtained by weighting and summing the one-dimensional sequence after capturing the dependencies with the channel attention weights of the original sequence.

[0028] The beneficial effects of this invention are as follows:

[0029] 1. A nested S-shaped scanning strategy was designed to maintain the locality and continuity within the strips of the image to be repaired, and to ensure the locality and continuity between strips through a strip sliding mechanism;

[0030] 2. A sequence shuffling attention module was designed to aggregate the processed sequences by calculating attention in the corresponding channels. This design enables it to capture the complex dependencies between different sequences, thereby making better use of complementary information from different scanning directions. Attached Figure Description

[0031] Figure 1 This is a schematic diagram of the state-space model used in the image restoration method based on the state-space model.

[0032] Figure 2 This is a schematic diagram of a nested S-shaped scanning strategy.

[0033] Figure 3 This is a schematic diagram of the sequence shuffling attention module.

[0034] Figure 4 The images show a comparison of the results of image restoration using the method proposed in this application. Detailed Implementation

[0035] The specific embodiments of the present invention are described below to enable those skilled in the art to understand the present invention. However, it should be understood that the present invention is not limited to the scope of the specific embodiments. For those skilled in the art, various changes are obvious as long as they are within the spirit and scope of the present invention as defined and determined by the appended claims. All inventions utilizing the concept of the present invention are protected.

[0036] like Figure 1 As shown in (a), in one embodiment of the present invention, an image restoration method based on a state-space model includes:

[0037] Shallow feature extraction stage: shallow features of the image to be restored are extracted through convolutional layers;

[0038] Deep feature extraction stage: Based on the shallow features of the image to be restored, deep features of the image to be restored are extracted using stacked residual Mamba group (RMG), nested S-shaped scanning strategy and sequence mixing attention module.

[0039] Reconstruction stage: Based on the shallow and deep features of the image to be restored, a pixel shuffling layer and a 3x3 convolutional layer are introduced to reconstruct the high-resolution image.

[0040] In the shallow feature extraction stage, for a given image to be restored Use a convolutional layer to extract shallow features ,in C , H and W These represent the number of model channels, height, and width, respectively.

[0041] In the deep feature extraction stage, shallow features Deep features are extracted by using several stacked residual Mamba groups (RMGs). ;

[0042] like Figure 1 (a) Figure 1 (b) and Figure 1As shown in (c), each residual Mamba group RMG includes several residual Mamba blocks RMB connected in sequence and a convolutional layer. Each residual Mamba block RMB includes a visual Mamba module VNM for capturing long-range dependencies. Furthermore, each visual Mamba module VNM includes a MaIRM (MaIR Module).

[0043] During the reconstruction phase, based on shallow features and deep features To reconstruct high-quality images, specifically for image super-resolution, we introduce a pixel shuffle layer and a 3x3 convolutional layer. To reconstruct high-resolution images This process can be represented as:

[0044]

[0045] In particular, for tasks that do not require upsampling (e.g., denoising, deblurring, and dehazing), we employ single convolutional layers with residual connections to build high-quality results.

[0046] MaIRM processes the input two-dimensional features as follows:

[0047] The two-dimensional features are flattened into four one-dimensional sequences along four different directions by using a nested S-shaped scanning strategy;

[0048] Use a selective scan operation to capture the long-range dependencies of four one-dimensional sequences;

[0049] The output features are obtained by aggregating the one-dimensional sequence after capturing the dependencies through the sequence fusion attention module.

[0050] The original state-space model was based on the formula To capture long-range dependencies, among which h ( t )yes t The hidden state of time yes h ( t ) is the derivative of . x yes t Input at any time A , B , C These are parameters that are independent of the input. y ( t )yes t Output at any moment;

[0051] This application is based on the RNN concept, and uses Mamba to analyze the input... x Learn (or reason) about the input A , B , C What would be a better amount? A , B , C The value is more suitable for the current input, and therefore the effect is better.

[0052] The nested S-shaped scanning strategy aims to extract sequences that maintain locality and continuity from input features, and its design is mainly based on two observations:

[0053] 1) If we want to maintain locality, we need to design a restricted scanning area;

[0054] 2) To avoid losing continuity, an S-shaped scanning path can be used to maintain continuity.

[0055] Therefore, as Figure 2 As shown in (a), this application relates to a nested S-shaped scanning strategy, the specific method of which is to divide the feature into multiple non-overlapping strips and use S-shaped scanning paths inside and between the strips to maintain locality and continuity;

[0056] In particular, to make better use of spatial information, we extracted the sequence in four different scanning directions: from the top left to the bottom right, from the bottom right to the top left, from the top right to the bottom left, and from the bottom left to the top right.

[0057] like Figure 2 As shown in (b), the nested S-shaped scan strategy also includes a strip sliding mechanism to maintain locality and continuity in the boundary region between adjacent strips; specifically including:

[0058] For two consecutive modules, the features of the first module are divided into non-overlapping strips. A shift strip operation is used to set the width of the first and last strips of the first module to half of their original width, while the width of the other strips remains unchanged. Therefore, the boundary region in the previous module will be completely covered by a single strip in this module, ensuring the locality and continuity of the boundary region between adjacent strips.

[0059] like Figure 3 As shown, Sequence Shuffle Attention (SSA) aggregates processed sequences by computing attention within corresponding channels. This design enables it to capture complex dependencies between different sequences, thereby better utilizing complementary information from different scanning directions. Methods for the sequence shuffle attention module to aggregate and capture the one-dimensional sequences with captured dependencies to obtain output features include:

[0060] The one-dimensional sequence after capturing dependencies is spatially averaged and then concatenated to obtain the pooled sequence. This process can be represented as follows:

[0061]

[0062] in, Represents the pooling sequence. This indicates a splicing operation. This represents the spatial average pooling operation, used to reduce the size of the feature spatial region, thereby reducing computational cost. This represents a one-dimensional sequence after capturing dependencies. It is the first k The first sequence d Pooling features in each channel D It is the number of channels in MaIRM. k ={1,2,3,4};

[0063] The pooling sequences are grouped, and the grouped sequences are shuffled and rearranged to obtain the rearranged sequences. The process is represented as follows:

[0064]

[0065] in, Indicates a rearranged sequence. This indicates a sequence shuffling operation;

[0066] The channel attention weights of the rearranged sequence are obtained by grouping convolutions, and the rearranged sequence is then rearranged back to its original order to obtain the channel attention weights in the original order. The process is as follows:

[0067]

[0068] in, Channel attention weights representing the original order. This represents a sequence rearrangement operation; This indicates a grouped convolution operation. The channel attention weights represent the rearranged sequence;

[0069] The output feature is obtained by weighted summation of the one-dimensional sequence after capturing dependencies and the channel attention weights of the original sequence. This process is represented as follows:

[0070]

[0071] in, Y This represents the output of the sequence mixing attention module. This represents the fusion weight of the sequences after block division. This indicates a block operation.

[0072] Regarding the loss function, for image super-resolution algorithms, this application adopts the L1 loss function, namely:

[0073]

[0074] in, L For loss function, y For the target image, Represents the L1 norm;

[0075] For tasks such as image denoising, deblurring, and dehazing, we use Charbonnier Loss to train the model, that is:

[0076]

[0077] in, Describes the 2-norm. This is a hyperparameter.

[0078] Benefiting from the preservation of image locality and continuity as well as attention to sequence differences, our method achieves superior performance on multiple public datasets.

[0079] Specifically, as shown in Table 1, in image super-resolution tasks, our method (MaIR) achieves state-of-the-art results compared to limiting methods such as MambaIR in super-resolution tasks of 2x, 3x, and 4x.

[0080] Table 1 Comparison of Experimental Results of Double Super-Resolution

[0081]

[0082] Similar results can also be seen in tasks such as image denoising, as shown in Table 2, where our method achieves superior performance.

[0083] Table 2 Comparison of image denoising experimental results, where 15, 25, and 50 represent different noise intensities.

[0084]

[0085] The same applies to visual effects, such as... Figure 4 As shown, taking the results on image super-resolution as an example, our method can recover a visual effect that is closer to the reference image (GT), that is, our texture is clearer and more accurate.

[0086] The above embodiments are only used to illustrate the technical solutions of this application, and are not intended to limit them. Although this application has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that modifications can still be made to the technical solutions described in the foregoing embodiments, or equivalent substitutions can be made to some of the technical features. Such modifications or substitutions do not cause the essence of the corresponding technical solutions to deviate from the spirit and scope of the technical solutions of the embodiments of this application.

Claims

1. An image restoration method based on a state-space model, characterized in that, include: Shallow feature extraction stage: shallow features of the image to be restored are extracted through convolutional layers; Deep feature extraction stage: Based on the shallow features of the image to be restored, deep features of the image to be restored are extracted using stacked residual Mamba group (RMG), nested S-shaped scanning strategy and sequence mixing attention module. Reconstruction stage: Based on the shallow and deep features of the image to be restored, a pixel shuffling layer and a 3x3 convolutional layer are introduced to reconstruct the high-resolution image. Each residual Mamba group RMG consists of several sequentially connected residual Mamba blocks RMB and a convolutional layer; Each residual Mamba block RMB includes a visual Mamba module VNM; Each visual Mamba module (VNM) includes a MaIRM; MaIRM processes the input two-dimensional features as follows: The two-dimensional features are flattened into four one-dimensional sequences along four different directions by using a nested S-shaped scanning strategy; Use selective scan operations to capture long-range dependencies of four one-dimensional sequences; The output features are obtained by aggregating the one-dimensional sequence after capturing the dependencies through the sequence fusion attention module.

2. The image restoration method based on a state-space model according to claim 1, characterized in that, The nested S-shaped scan strategy is used to extract sequences that maintain locality and continuity from input features. Specifically, the features are divided into multiple non-overlapping strips, and S-shaped scan paths are used within and between the strips.

3. The image restoration method based on a state-space model according to claim 2, characterized in that, The nested S-shaped scan strategy also includes a strip sliding mechanism to maintain locality and continuity in the boundary region between adjacent strips; specifically including: For two consecutive modules, the features of the first module are divided into non-overlapping strips, and a shift strip operation is used to set the width of the first and last strips of the first module to half of their original width, while the width of the other strips remains unchanged.

4. The image restoration method based on a state-space model according to claim 2, characterized in that, Methods for using sequence hybrid attention modules to aggregate and capture one-dimensional sequences with dependencies to obtain output features include: The one-dimensional sequence after capturing dependencies is spatially averaged and then concatenated to obtain the pooled sequence. The pooled sequences are grouped, and the grouped sequences are shuffled and rearranged to obtain the rearranged sequences. The channel attention weights of the rearranged sequence are obtained by grouping convolutions, and the rearranged sequence is rearranged back to the original order to obtain the channel attention weights of the original order. The output feature is obtained by weighting and summing the one-dimensional sequence after capturing the dependencies with the channel attention weights of the original sequence.

Images