Image restoration method based on state space model

Abstract

The invention discloses an image restoration method based on a state space model. The method comprises the following steps: a shallow feature extraction stage: extracting shallow features of an image to be restored through a convolutional layer; a deep feature extraction stage: based on the shallow features of the to-be-restored image, extracting the deep features of the to-be-restored image by using a stacked residual Mama group RMG, a nested S-shaped scanning strategy and a sequence mixed attention module; and a reconstruction stage: based on the shallow features of the image to be restored and the deep features of the image to be restored, introducing a pixel shuffling layer and a convolution layer with the convolution kernel of 3 * 3 to perform high-resolution image reconstruction, and retaining the locality and continuity of the image by using a nested S-shaped scanning strategy, and effective context information aggregation is carried out on sequences in different directions through the sequence shuffling attention module, excellent performance is shown in multiple tasks such as super-resolution, denoising, deblurring and defogging, and a new thought is provided for efficient image restoration.

Description

Technical Field

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

[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 restore high-quality clean images from given degraded images (such as noisy, blurred or hazy images). It has wide applications in 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 in small areas) and continuity (i.e., smooth and gradual transitions in larger areas) in natural images. However, these methods still have limitations in capturing long-range dependencies. Convolutional neural networks are difficult to effectively model the global relationship of images due to their fixed-size receptive field, while Transformer-based methods are limited by computational complexity and local window segmentation strategies, and have shortcomings in long-range modeling, limiting their application in a wider range of scenarios. Therefore, it is highly desirable to develop a method that can capture long-range dependencies while maintaining the inherent locality and continuity of natural images.

[0003] The State Space Model (SSM) is an emerging sequence modeling method that can efficiently model long sequence dependencies with approximately linear complexity through the Selective Scan Operation. Mamba, a new state-space model, has been widely introduced into the visual field in recent years due to its superior performance. However, since the core algorithm of the state-space model, the selective scanning operation, is essentially designed for one-dimensional sequences, it cannot be directly applied to process two-dimensional images. To solve this problem, existing Mamba-based image restoration methods (such as MambaIR) usually involve a three-step process: 1) Flatten the two-dimensional image into multiple one-dimensional sequences; 2) Process each sequence independently using a selective scan operation; 3) Aggregate the processed sequences to form an output 2D image.

[0004] However, this paradigm still faces two shortcomings 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 image; second, it usually aggregates the processed sequences by pixel-level summation, ignoring the different contexts between sequences unfolded in completely different ways. Summary of the invention

[0005] In view of the above-mentioned deficiencies in the prior art, the present invention provides an image restoration method based on a state-space model, which utilizes a nested S-shaped scanning strategy to retain the locality and continuity of the image, and effectively aggregates contextual information of sequences in different directions through a sequence shuffling attention module, thereby solving the problem that the existing expansion methods ignore the locality and continuity of the image during the expansion stage of the visual state-space model.

[0006] In order to achieve the above-mentioned object of the invention, the technical solution adopted by the present invention is: an image restoration method based on a state space model, comprising: Shallow feature extraction stage: extract shallow features of the image to be restored through the convolution layer; Deep feature extraction stage: Based on the shallow features of the image to be restored, the stacked residual Mamba group RMG, nested S-shaped scanning strategy and sequential mixed attention module are used to extract the deep features of the image to be restored; Reconstruction stage: Based on the shallow features and deep features of the image to be restored, a pixel shuffling layer and a convolution layer with a convolution kernel of 3x3 are introduced to perform high-resolution image reconstruction.

[0007] Furthermore: each residual Mamba group RMG includes a number of sequentially connected residual Mamba blocks RMB and a convolutional layer.

[0008] Furthermore: each residual Mamba block RMB includes a visual Mamba module VNM.

[0009] Furthermore: each visual Mamba module VNM includes MaIRM.

[0010] Further: the method of processing the input two-dimensional features by MaIRM is: The two-dimensional features are flattened into four one-dimensional sequences along four different directions by nesting S-shaped scanning strategies; The selective scanning operation is used to capture the long-range dependencies of the four 1D sequences; The one-dimensional sequence after capturing the dependencies is aggregated through the sequence mixed attention module to obtain the output features.

[0011] Furthermore: A nested S-shaped scanning strategy is used to extract a sequence that maintains locality and continuity from the input features, and the specific method is: the features are divided into multiple non-overlapping strips, and an S-shaped scanning path is used within and between the strips.

[0012] Furthermore: the nested S-shaped scanning strategy also includes a strip sliding mechanism for maintaining locality and continuity in the boundary area between adjacent strips; specifically including: For two consecutive modules, the features of the first module are divided into non-overlapping strips, and the shift strip operation is used to set the width of the first and last strips of the first module to 1 / 2 of the original width, while the widths of other strips remain unchanged.

[0013] Further: the sequence hybrid attention module aggregates the one-dimensional sequence after capturing the dependency, and the method of obtaining the output feature includes: The one-dimensional sequence after capturing the dependency is spatially averaged and concatenated to obtain a pooled sequence; The pooled sequences are grouped, and the grouped sequences are shuffled and rearranged to obtain rearranged sequences; Obtain the channel attention weights of the rearranged sequence through grouped convolution, and rearrange the rearranged sequence back to the original order to obtain the channel attention weights of the original order; The one-dimensional sequence after capturing the dependency is weighted summed with the channel attention weights of the original order to obtain the output feature.

[0014] The beneficial effects of the present invention are: 1. A nested S-shaped scanning strategy is designed to maintain the locality and continuity within the image strip to be repaired, and the locality and continuity between strips are ensured through the strip sliding mechanism; 2. A sequence shuffle attention module is designed to aggregate the processed sequences by calculating attention within the corresponding channels. This design enables it to capture the complex dependencies between different sequences, thereby better utilizing the complementary information from different scanning directions. BRIEF DESCRIPTION OF THE DRAWINGS

[0015] Figure 1 Schematic diagram of the state space model used in the image restoration method based on the state space model.

[0016] Figure 2 Schematic diagram of the nested S-shaped scanning strategy.

[0017] Figure 3 Schematic diagram of the sequence shuffle attention module.

[0018] Figure 4 A comparison chart of image restoration results including the method proposed in this application. DETAILED DESCRIPTION

[0019] The specific implementation modes of the present invention are described below so that those skilled in the art can understand the present invention. However, it should be clear that the present invention is not limited to the scope of the specific implementation modes. For those of ordinary skill in the art, as long as various changes are within the spirit and scope of the present invention as defined and determined by the attached claims, these changes are obvious, and all inventions and creations utilizing the concept of the present invention are protected.

[0020] like Figure 1 As shown in (a), in one embodiment of the present invention, an image restoration method based on a state space model comprises: Shallow feature extraction stage: extract shallow features of the image to be restored through the convolution layer; Deep feature extraction stage: Based on the shallow features of the image to be restored, the stacked residual Mamba group RMG, nested S-shaped scanning strategy and sequential mixed attention module are used to extract the deep features of the image to be restored; Reconstruction stage: Based on the shallow features and deep features of the image to be restored, a pixel shuffling layer and a convolution layer with a convolution kernel of 3x3 are introduced to perform high-resolution image reconstruction.

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

[0022] In the deep feature extraction stage, shallow features After several stacked residual Mamba groups RMG, deep features are extracted ; like Figure 1 (a) Figure 1 (b) and Figure 1 As 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-distance dependencies; further, each visual Mamba module VNM includes MaIRM (MaIR Module).

[0023] In the reconstruction phase, based on shallow features and deep features Reconstruct high-quality images; specifically, for image super-resolution, we introduce a pixel shuffle layer and a convolution layer with a convolution kernel of 3x3. To reconstruct a high-resolution image , this process can be expressed as:

[0024] In particular, for tasks that do not require upsampling (e.g., denoising, deblurring, and dehazing), we adopt a single convolutional layer with residual connections to achieve high-quality results.

[0025] The method of MaIRM to process the input two-dimensional features is: The two-dimensional features are flattened into four one-dimensional sequences along four different directions by nesting S-shaped scanning strategies; Use Selective Scan Operation to capture the long-range dependencies of four one-dimensional sequences; The one-dimensional sequence after capturing the dependencies is aggregated through the sequence mixed attention module to obtain the output features.

[0026] The original state space model is expressed by the formula To capture long-range dependencies, h ( t )yes t The hidden state at the moment, yes h ( t ) is the derivative of, x yes t Input at the moment, A , B , C is a parameter that has nothing to do with the input. y ( t )yes t Output at the moment; This application is based on the idea of ​​RNN, through Mamba according to the input x , learn (or infer) the input A , B , C How much should be better? In this case, A , B , C The value is more suitable for the current input and the effect is better.

[0027] The nested S-shaped scanning strategy aims to extract sequences that maintain locality and continuity from the input features. Its design is mainly based on two observations: 1) If we want to maintain locality, we need to design a limited scanning area; 2) To avoid loss of continuity, continuity can be maintained through an S-shaped scanning path.

[0028] Therefore, if Figure 2 As shown in (a), the present application involves a nested S-shaped scanning strategy, and the specific method is: dividing the feature into multiple non-overlapping strips, and using an S-shaped scanning path within and between the strips to maintain locality and continuity; In particular, to better utilize the spatial information, we extract sequences in four different scanning directions: from top-left to bottom-right, from bottom-right to top-left, from top-right to bottom-left, and from bottom-left to top-right.

[0029] like Figure 2 As shown in (b), the nested S-shaped scanning strategy also includes a strip sliding mechanism for maintaining locality and continuity in the boundary area between adjacent strips; specifically, it includes: For two consecutive modules, the features of the first module are divided into non-overlapping strips, and the shift strip operation is used to set the width of the first and last strips of the first module to 1 / 2 of the original, and the width of other strips remains unchanged. Therefore, the boundary area in the previous module will be completely covered by a single strip in this module, ensuring the locality and continuity of the boundary area between adjacent strips.

[0030] like Figure 3 As shown in the figure, Sequence Shuffle Attention (SSA) aggregates the processed sequences by calculating attention in the corresponding channels. This design enables it to capture the complex dependencies between different sequences, thereby better utilizing the complementary information from different scanning directions. The sequence hybrid attention module aggregates the one-dimensional sequence after capturing the dependencies, and the method for obtaining the output features includes: The one-dimensional sequence after capturing the dependency is spatially averaged and concatenated to obtain a pooled sequence. The process can be expressed as:

[0031] in, represents the pooling sequence, Represents a splicing operation, Represents the spatial average pooling operation, which is used to reduce the feature spatial domain size to reduce the computational cost. represents a one-dimensional sequence after capturing dependencies, It is k The first d The pooled features in channels, D is the number of channels in MailM,k ={1,2,3,4}; The pooled sequences are grouped, and the grouped sequences are shuffled and rearranged to obtain the rearranged sequences. The process is expressed as:

[0032] in, Represents a rearranged sequence, Indicates a sequence shuffle operation; The channel attention weights of the rearranged sequence are obtained through grouped convolution, and the rearranged sequence is rearranged back to the original order to obtain the channel attention weights of the original order. The process is expressed as:

[0033] in, represents the channel attention weights in the original order, Represents a sequence rearrangement operation; represents the grouped convolution operation, represents the channel attention weight of the rearranged sequence; The one-dimensional sequence after capturing the dependency is weighted summed with the channel attention weights of the original order to obtain the output feature. The process is expressed as:

[0034] in, Y represents the output of the sequence mixed attention module, represents the fusion weight of each sequence after block division, Indicates a block operation.

[0035] In terms of loss function, for the image super-resolution algorithm, this application adopts the L1 loss function, namely:

[0036] in, L is the loss function, y is the target image, represents the L1 norm; For image denoising, deblurring, defogging and other tasks, we use Charbonnier Loss to train the model, namely:

[0037] in, represents the 2-norm, is a hyperparameter.

[0038] Benefiting from the preservation of image locality and continuity and the focus on sequence differences, our method produces better performance than the original ones on multiple public datasets.

[0039] Specifically, as shown in Table 1, in the image super-resolution task, our method (MaIR) has achieved the current best results in 2, 3, and 4 times super-resolution tasks compared with extreme methods such as MambaIR.

[0040] Table 1 Comparison of double super-resolution experimental results

[0041] Similar results can also be confirmed in tasks such as image denoising. As shown in Table 2, our method has achieved the current best results.

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

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

[0044] The above embodiments are only used to illustrate the technical solutions of the present application, rather than to limit them. Although the present application has been described in detail with reference to the aforementioned embodiments, those skilled in the art should understand that they can still modify the technical solutions described in the aforementioned embodiments, or make equivalent replacements for some of the technical features therein. However, these modifications or replacements do not deviate the essence of the corresponding technical solutions from the spirit and scope of the technical solutions of the embodiments of the present application.

Claims

1. An image restoration method based on a state space model, characterized in that: include: Shallow feature extraction stage: extract shallow features of the image to be restored through the convolution layer; Deep feature extraction stage: Based on the shallow features of the image to be restored, the stacked residual Mamba group RMG, nested S-shaped scanning strategy and sequential mixed attention module are used to extract the deep features of the image to be restored; Reconstruction stage: Based on the shallow features and deep features of the image to be restored, a pixel shuffling layer and a convolution layer with a convolution kernel of 3x3 are introduced to perform high-resolution image reconstruction.

2. The image restoration method based on the state space model according to claim 1, characterized in that: Each residual Mamba group RMG consists of several sequentially connected residual Mamba blocks RMB and a convolutional layer.

3. The image restoration method based on the state space model according to claim 2, characterized in that: Each residual Mamba block RMB includes a visual Mamba module VNM.

4. The image restoration method based on the state space model according to claim 3, characterized in that: MaIRM is included in every Visual Mamba module VNM.

5. The image restoration method based on the state space model according to claim 4, characterized in that: The method of MaIRM to process the input two-dimensional features is: The two-dimensional features are flattened into four one-dimensional sequences along four different directions by nesting S-shaped scanning strategies; The selective scanning operation is used to capture the long-range dependencies of the four 1D sequences; The one-dimensional sequence after capturing the dependencies is aggregated through the sequence mixed attention module to obtain the output features.

6. The image restoration method based on the state space model according to claim 5, characterized in that: The nested S-shaped scanning strategy is used to extract a sequence that maintains locality and continuity from the input features. The specific method is: the features are divided into multiple non-overlapping strips and S-shaped scanning paths are used within and between strips.

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

8. The image restoration method based on the state space model according to claim 5, characterized in that: The sequence hybrid attention module aggregates the one-dimensional sequence after capturing the dependency, and the method of obtaining the output features includes: The one-dimensional sequence after capturing the dependency is spatially averaged and concatenated to obtain a pooled sequence; The pooled sequences are grouped, and the grouped sequences are shuffled and rearranged to obtain rearranged sequences; Obtain the channel attention weights of the rearranged sequence through grouped convolution, and rearrange the rearranged sequence back to the original order to obtain the channel attention weights of the original order; The one-dimensional sequence after capturing the dependency is weighted summed with the channel attention weights of the original order to obtain the output feature.

Images