An optical flow estimation method, program, device and storage medium based on a state space model

Abstract

This invention belongs to the field of computer vision technology, specifically relating to an optical flow estimation method, program, device, and storage medium based on a state-space model. The invention designs an image feature enhancement network based on a CNN-Mamba hybrid module and a Mamba feature cross-enhancement module, and designs a neural network with a CNN-Mamba hybrid recurrent regression framework. This invention utilizes the CNN-Mamba hybrid module's image feature enhancement network to efficiently extract global information from the input image, not only exploring long-range dependencies between pixels but also maintaining linear complexity. This invention uses the Mamba feature cross-enhancement module to obtain global motion features, aggregates global texture information within the motion boundary, and enhances the texture features, ultimately improving the accuracy of optical flow estimation.

Description

Technical Field

[0001] This invention belongs to the field of computer vision technology, specifically relating to an optical flow estimation method, program, device, and storage medium based on a state-space model. Background Technology

[0002] Optical flow estimation across consecutive video frames is a long-standing, fundamental, and challenging problem in computer vision. It is crucial for building higher-level cognitive abilities in scene understanding, such as object recognition, object tracking, action recognition, and scene segmentation. With the development of deep learning, using neural network models to solve the optical flow estimation problem has gradually become the mainstream approach. PWC-Net and RAFT are representative recent methods based on CNN (Convolutional Neural Network). PWC-Net constructs hierarchical local cost volumes with warp features and progressively estimates the flow of these local costs. RAFT proposes an iterative framework that constructs a multi-scale four-dimensional correlation volume for all pixel pairs and deploys gated recurrent units (GRUs) to iteratively update the predicted optical flow. It also makes numerous iterative improvements, convolving different local cost volumes at different iteration stages to progressively reach a near-global search space, achieving excellent performance on standard benchmarks. However, these methods are limited by the size of the CNN convolution window and have not achieved better accuracy.

[0003] Recently, with the rise of Mamba in vision tasks, its excellent input adaptability and global information modeling capabilities, state space, and Mamba architecture have achieved success in multiple vision subtasks. Summary of the Invention

[0004] The purpose of this invention is to provide an optical flow estimation method, program, device, and storage medium based on a state-space model.

[0005] An optical flow estimation method based on a state-space model, which uses two consecutive frames of images I t and I t+1 The input is fed into the pre-trained neural network of the CNN-Mamba hybrid recurrent regression framework;

[0006] The neural network of the CNN-Mamba hybrid recurrent regression framework includes an image feature enhancement network module, a texture feature extraction network module, a motion feature encoding module, a Mamba feature cross enhancement module, and a recurrent iterative refinement module; the image feature enhancement network module includes a CNN-based convolutional network module and a Mamba module; the Mamba feature cross enhancement module includes a Mamba motion enhancement module and a texture feature enhancement module;

[0007] The CNN-based convolutional network module processes the input image It and I t+1 Local features are extracted and image features are downsampled simultaneously to obtain image features F. t and F t+1 Image features F t and F t+1 Image feature enhancement is performed using the Mamba module to obtain the enhanced image features. and The texture feature extraction network module extracts the texture from the input image I. t Extracting texture features F c ; and The global correlation volume is obtained by performing pixel-by-pixel dot product operations, and then the motion features F are obtained through the motion feature encoding module. m Motion characteristics F m and texture features F c Global motion features are obtained using the Mamba motion enhancement module. Global motion characteristics and motion texture features F c Texture enhancement features are obtained through the texture feature enhancement module. Global motion characteristics and texture enhancement features The input is fed into the iterative refinement module to calculate the optical flow estimate.

[0008] Furthermore, the image feature enhancement network includes multiple sets of Mamba modules. Each Mamba module captures the long-range dependencies between each pixel in the input image features while maintaining linear complexity and reducing computational overhead. The operation of the Mamba module is represented as follows:

[0009]

[0010] Among them, F t k and This represents the input and output of the k-th Mamba module; F represents t k Regularized image features LN represents the normalization layer; f linear DWConV represents a linear function; DWConV represents a depthwise separable convolutional layer; SiLU represents an activation function. SS2D represents a pixel-wise dot product operation; SS2D represents a 2D selective scan operation.

[0011] Furthermore, the SS2D selective scanning operation is specifically described as follows:

[0012] Δ=Δ(x)=wΔ x, B=B(x)=w B x, C = C(x) = w c x

[0013]

[0014] y = [y1, y2, ..., y t ,…,y L ]

[0015] Where x represents the input feature in the SS2D operation, x = [x1, x2, ..., x t ,…,x L ]; Δ represents the time scale parameter; B and C represent the linear projections of the input features, respectively; w Δ ,w B ,w C Representing three different parameter matrices; A represents the state matrix. and These are the discrete parameters of A and B, respectively; D represents the skip link, h t Let y be the intermediate state at time t, and y represent the output characteristics of the SS2D operation.

[0016] Furthermore, the motion feature F m and texture features F c The long-range dependencies between pixels are explored using the Mamba motion enhancement module, and the learned weights are then aggregated into the motion feature F. m Then, global motion features are obtained through skip links. The specific operation is represented as follows:

[0017]

[0018]

[0019] in, and represents the regularized motion features and texture features, respectively; Z represents the linear projection of the input features; GS represents global feature search.

[0020] Furthermore, the global motion features and motion texture features F c The texture feature enhancement module predicts a set of scale weights from the global motion features, and then uses these weights to achieve a global adjustment of the entire dense displacement, ultimately obtaining the enhanced texture features. The specific operation is represented as follows:

[0021]

[0022] Where Sigmoid represents normalization; Conv represents convolution; ReLU represents activation function; and AVG represents average pooling.

[0023] Furthermore, the iterative refinement module first refines the motion features F m Global motion characteristics and texture enhancement features The optical flow is stitched along the channel dimension and then fed into a gated loop unit to iteratively update the predicted optical flow. Subsequently, the corresponding optical flow is upsampled to obtain the original resolution optical flow map.

[0024] Furthermore, the texture feature extraction network module is constructed in the same way as the image feature enhancement network module.

[0025] A computer device / apparatus / system includes a memory, a processor, and a computer program stored in the memory, wherein the processor executes the computer program to implement the steps of the above-described optical flow estimation method based on a state-space model.

[0026] A computer-readable storage medium having a computer program / instructions stored thereon, which, when executed by a processor, implements the steps of the aforementioned optical flow estimation method based on a state-space model.

[0027] A computer program product includes a computer program / instructions that, when executed by a processor, implement the steps of the aforementioned optical flow estimation method based on a state-space model.

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

[0029] This invention designs an image feature enhancement network based on a CNN-Mamba hybrid module and a Mamba feature cross-enhancement module, and designs a neural network with a CNN-Mamba hybrid recurrent regression framework. This invention utilizes the CNN-Mamba hybrid module's image feature enhancement network to efficiently extract global information from the input image, not only exploring long-range dependencies between pixels but also maintaining linear complexity. This invention uses the Mamba feature cross-enhancement module to obtain global motion features, aggregates global texture information within the motion boundary, and enhances the texture features, ultimately improving the accuracy of optical flow estimation. Attached Figure Description

[0030] Figure 1 This is a flowchart illustrating the overall concept of the present invention.

[0031] Figure 2 This is a diagram showing the overall structure of the neural network in the CNN-Mamba hybrid recurrent regression framework of this invention.

[0032] Figure 3 This is a structural diagram of the CNN-Mamba hybrid module in this invention.

[0033] Figure 4 This is a structural diagram of the Mamba feature cross-enhancement module in this invention.

[0034] Figure 5 This is a structural diagram of the 2D selection search (SS2D) in this invention. Detailed Implementation

[0035] The present invention will now be further described with reference to the accompanying drawings.

[0036] This invention proposes an optical flow estimation method based on a state-space model, employing a neural network within a CNN-Mamba hybrid recurrent regression framework. The invention also constructs an image feature enhancement network based on a CNN-Mamba hybrid module. This network consists of a feature extraction convolutional network and an image feature enhancement module composed of Mamba modules, which not only efficiently extracts global information from the input image but also maintains linear complexity and reduces computational overhead. To better aggregate texture information within motion boundaries, this invention constructs a Mamba feature cross-enhancement module. This module comprises a Mamba motion enhancement module and a texture feature enhancement module. The Mamba motion enhancement module utilizes a global feature search mechanism to explore long-term dependencies between pixels in motion features and motion texture features, and performs effective feature fusion to obtain global motion features. The texture feature enhancement module predicts a set of scale weights from the global motion features and then uses them to enhance the information of the texture features.

[0037] This invention constructs a neural network based on a CNN-Mamba hybrid recurrent regression framework, including an image feature enhancement network module based on a CNN-Mamba hybrid module, a texture feature extraction network module based on a CNN-Mamba hybrid module, a motion feature encoding module, a Mamba feature cross enhancement module, and a recurrent iterative refinement module. The Mamba feature cross enhancement module includes a Mamba motion enhancement module and a texture feature enhancement module.

[0038] Two consecutive frames of images I t and I t+1 The input is fed into a pre-trained neural network based on a CNN-Mamba hybrid recurrent regression framework. The image feature enhancement network, based on the CNN-Mamba hybrid module, first extracts local features from the input image using a feature extraction convolutional network while simultaneously downsampling the image features, resulting in image features F. t and F t+1 Next, F t and F t+1The input image feature enhancement module contains N Mamba modules. Each Mamba module captures the long-range dependencies between each pixel in the input image features while maintaining linear complexity and reducing computational overhead. The last Mamba module outputs the enhanced image features. and

[0039] Will and The global correlation volume is obtained by performing pixel-by-pixel dot product operations, and then the motion features F are obtained through the motion feature encoding module. m The texture feature extraction network module based on the CNN-Mamba hybrid module extracts the input image I. t Extracting texture features F c Motion characteristics F m and texture features F c Global motion features are obtained using the Mamba motion enhancement module. Global motion characteristics and motion texture features F c The texture feature enhancement module predicts a set of scale weights from the global motion features, and then uses these weights to achieve a global adjustment of the entire dense displacement, ultimately obtaining the enhanced texture features.

[0040]

[0041] Global motion characteristics and texture enhancement features The input is fed into the iterative refinement module, which first processes the motion features F. m Global motion characteristics and texture enhancement features The optical flow is stitched along the channel dimension and then fed into a gated loop unit to iteratively update the predicted optical flow. Subsequently, the corresponding optical flow is upsampled to obtain the original resolution optical flow map.

[0042] Example 1:

[0043] like Figure 3 As shown, input two consecutive frames of images I t and I t+1 The convolutional network for feature extraction consists of a 7×7 convolution with a stride of 2, a 1×1 convolution, and three residual modules, each composed of two 3×3 convolutions with a stride of 1. The feature extraction module then produces... Image features F t and F t+1 Where H and W represent the height and width of the input image, respectively, C represents the number of channels in the input image, and then the feature image F...t and F t+1 An image feature enhancement module consisting of N Mamba modules is used. Each Mamba module captures the remote dependencies between each pixel in the input image features. Its Mamba module operation is defined as follows:

[0044]

[0045]

[0046] Among them, F t k and This represents the input and output of the k-th Mamba module; F represents t k Regularized image features, LN represents the normalization layer; f linear DWConV represents a linear function; DWConV represents a depthwise separable convolutional layer; SiLU represents an activation function. SS2D represents a pixel-wise dot product operation; SS2D represents a 2D selective scan operation, defined as:

[0047] Δ=Δ(x)=w Δ x, B=B(x)=w B x, C = C(x) = w c x (3)

[0048]

[0049] y = [y1, y2, ..., y t ,…,y L (6)

[0050] Where x represents the input feature in the SS2D operation, x = [x1, x2, ..., x t ,…,x L ], its length is L; Δ represents the time scale parameter; B and C represent the linear projections of the input features, respectively; w Δ ,w B ,w C Representing three different parameter matrices; A represents the state matrix. and These are the discrete parameters of A and B, respectively; D represents the skip link, h t Let y be the intermediate state at time t, and y represent the output characteristics of the SS2D operation.

[0051] Enhanced image features output from the image feature enhancement network based on the CNN-Mamba hybrid module and The global correlation volume is obtained by performing pixel-by-pixel dot product operations, and then the motion features F are obtained through the motion feature encoding module. m Its motion feature encoding module consists of 4 CNN blocks. Each CNN block contains a convolution with a stride of 1 and a size of 3×3 kernel, as well as a ReLU activation function.

[0052] Image I t The input is fed into a texture feature extraction network based on a CNN-Mamba hybrid module to obtain texture features F. c Its texture feature extraction network consists of a 7×7 convolution with a stride of 2, a 1×1 convolution, three residual modules, and an image feature enhancement module composed of N Mamba modules. Each residual module consists of two 3×3 convolutions with a stride of 1. Figure 4 As shown, the Mamba feature cross-enhancement module consists of the Mamba motion enhancement module and the texture feature enhancement module, with motion feature F... m and texture features F c The long-range dependencies between pixels are explored using the Mamba motion enhancement module, and the learned weights are then aggregated into the motion feature F. m Then, global motion features are obtained through a skip link, and the operation is as follows:

[0053]

[0054] in, and represents the regularized motion features and texture features, respectively; Z represents the linear projection of the input features; GS represents global feature search; Represents global motion characteristics. For example... Figure 4 As shown, then global motion features and texture features F c The texture feature enhancement module predicts a set of scale weights from the global motion features, and then uses them to achieve a global adjustment of the entire dense displacement, ultimately obtaining the texture enhancement features. The operation is as follows:

[0055]

[0056] Where Sigmoid represents normalization, Conv represents convolution, ReLU represents activation, and AVG represents average pooling. This represents a texture enhancement feature.

[0057] This invention constructs a neural network with a CNN-Mamba hybrid recurrent regression framework. This network architecture includes an image feature enhancement network based on a CNN-Mamba hybrid module, a Mamba motion aggregation module, a relevant volumetric encoding module, and a recurrent iterative refinement module. The global motion features and texture enhancement features output from the Mamba feature cross-enhancement module are input into the recurrent iterative refinement module to calculate the corresponding optical flow. In the recurrent iterative refinement module, the motion features F are first processed... m Global motion characteristics and texture enhancement features The optical flow is stitched along the channel dimension and then fed into a gated recurrent unit (GRU) to iteratively update the predicted optical flow. The corresponding optical flow is then upsampled to obtain the original resolution optical flow map.

[0058] Two consecutive image frames are input to the network, and the network is trained under supervision using the overall network loss function. Ground Truth is used to monitor all traffic predictions. Two consecutive image frames are input to the trained CNN-Mamba hybrid recurrent regression framework neural network, and the output is the corresponding estimated optical flow.

[0059] The above description is merely a preferred embodiment of the present invention and is not intended to limit the invention. Various modifications and variations can be made to the present invention by those skilled in the art. Any modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of the present invention should be included within the scope of protection of the present invention.

Claims

1. An optical flow estimation method based on a state-space model, characterized in that: Two consecutive frames of images I t and I t+1 The input is fed into the pre-trained neural network of the CNN-Mamba hybrid recurrent regression framework; The neural network of the CNN-Mamba hybrid recurrent regression framework includes an image feature enhancement network module, a texture feature extraction network module, a motion feature encoding module, a Mamba feature cross enhancement module, and a recurrent iterative refinement module; the image feature enhancement network module includes a CNN-based convolutional network module and multiple Mamba modules; the Mamba feature cross enhancement module includes a Mamba motion enhancement module and a texture feature enhancement module; The CNN-based convolutional network module processes the input image I t and I t+1 Local features are extracted and image features are downsampled simultaneously to obtain image features F. t and F t+1 Image features F t and F t+1 Image feature enhancement is performed using the Mamba module to obtain the enhanced image features. and Each Mamba module captures the long-range dependencies between every pixel in the input image features, while maintaining linear complexity and reducing computational overhead; the operation of a Mamba module is represented as follows: Among them, F t k and This represents the input and output of the k-th Mamba module; F represents t k Regularized image features LN represents the normalization layer; f linear DWConV represents a linear function; DWConV represents a depthwise separable convolutional layer; SiLU represents an activation function. SS2D represents pixel-wise dot product operation; SS2D represents 2D selective scan operation. The texture feature extraction network module extracts the texture from the input image I. t Extracting texture features F c ; and The global correlation volume is obtained by performing pixel-by-pixel dot product operations, and then the motion features F are obtained through the motion feature encoding module. m Motion characteristics F m and texture features F c The long-range dependencies between pixels are explored using the Mamba motion enhancement module, and the learned weights are then aggregated into the motion feature F. m Then, global motion features are obtained through skip links. in, and represents the regularized motion features and texture features, respectively; Z represents the linear projection of the input features; GS represents global feature search; Global motion characteristics and motion texture features F c The texture feature enhancement module predicts a set of scale weights from the global motion features, and then uses these weights to achieve a global adjustment of the entire dense displacement, ultimately obtaining the enhanced texture features. Where Sigmoid represents normalization; Conv represents convolution; ReLU represents activation function; and AVG represents average pooling. Global motion characteristics and texture enhancement features The input is fed into the iterative refinement module, which first processes the motion features F... m Global motion characteristics and texture enhancement features The optical flow is stitched along the channel dimension and then fed into a gated loop unit to iteratively update the predicted optical flow. Subsequently, the corresponding optical flow is upsampled to obtain the original resolution optical flow map, thus obtaining the optical flow estimate.

2. The optical flow estimation method based on a state-space model according to claim 1, characterized in that: The SS2D selective scanning operation described above is specifically represented as follows: Δ=Δ(x)=w Δ x,B=B(x)=w B x,C=C(x)=w c x y t =Ch t +Dx t y=[y1,y2,…,y t ,…,y L ] Where x represents the input feature in the SS2D operation, x = [x1, x2, ..., x t ,…,x L ]; Δ represents the time scale parameter; B and C represent the linear projections of the input features, respectively; w Δ ,w B ,w C Representing three different parameter matrices; A represents the state matrix. and These are the discrete parameters of A and B, respectively; D represents the skip link, h t Let y be the intermediate state at time t, and y represent the output characteristics of the SS2D operation.

3. The optical flow estimation method based on a state-space model according to claim 1, characterized in that: The texture feature extraction network module is constructed in the same way as the image feature enhancement network module.

4. A computer device, comprising a memory, a processor, and a computer program stored in the memory, characterized in that: The processor executes the computer program to implement the steps of the method according to any one of claims 1 to 3.

5. A computer-readable storage medium having a computer program stored thereon, characterized in that: When executed by a processor, the computer program implements the steps of the method according to any one of claims 1 to 3.

6. A computer program product comprising computer instructions, characterized in that: When executed by a processor, the computer instructions implement the steps of the method according to any one of claims 1 to 3.

Images