Optical flow estimation method based on state space model, program, equipment and storage medium

Abstract

The invention belongs to the technical field of computer vision, and particularly relates to an optical flow estimation method based on a state space model, a program, equipment and a storage medium. According to the invention, an image feature enhancement network based on a CNN-Mama hybrid module and a Mama feature cross enhancement module are designed, and a neural network of a CNN-Mama hybrid cyclic regression framework is designed. The global information of the input image is efficiently extracted by using the image feature enhancement network of the CNN-Mama hybrid module, so that not only can the remote dependency relationship between pixels of the input image be explored, but also the linear complexity can be maintained. According to the method, the Mama feature cross enhancement module is utilized to obtain global motion features, global texture information is aggregated in a motion boundary, information enhancement of the texture features is realized, and finally, the precision of optical flow estimation is improved.

Description

Technical Field

[0001] The present invention belongs to the field of computer vision technology, and in particular relates to an optical flow estimation method, program, device and storage medium based on a state space model. Background Art

[0002] The estimation of optical flow of consecutive frames of a video is a long-standing, fundamental and challenging problem in the field of computer vision. It is essential for building higher-level cognitive capabilities in scene understanding, such as object recognition, object tracking, action recognition, scene segmentation, etc. With the development of deep learning, using neural network models to solve the problem of optical flow estimation has gradually become a mainstream solution. PWC-Net and RAFT are recent representative methods based on CNN convolutional neural networks. PWC-Net constructs hierarchical local cost volumes with warp features and gradually estimates the flow of these local costs. RAFT proposes an iterative framework that constructs multi-scale four-dimensional correlation volumes for all pixel pairs and deploys gated recurrent units (GRUs) to iteratively update the predicted optical flow. A large number of iterative improvements are made, and different local cost volumes are convolved at different iteration stages, thereby gradually reaching a near-global search space and achieving excellent performance on standard benchmarks. However, these methods are limited by the size of the CNN convolution window and fail to achieve better accuracy.

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

[0004] The object of the present 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 state space model is used to estimate the optical flow of two consecutive frames of images I t and I t+1 Input into the pre-trained CNN-Mamba hybrid recurrent regression framework neural network;

[0006] The neural network of the CNN-Mamba hybrid cyclic 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 cyclic 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] Based on the CNN convolutional network module, the input image It and I t+1 Extract local features separately and downsample image features to obtain image features F t and F t+1 ; Image feature F t and F t+1 The image features are enhanced through the Mamba module to obtain the enhanced image features. and The texture feature extraction network module is used to extract the input image I t Extract texture features F c ; and Perform pixel-by-pixel dot product operations to obtain the global correlation volume, and then obtain the motion feature F through the motion feature encoding module. m , motion feature F m and texture feature F c Get global motion features through Mamba motion enhancement module Global motion features and motion texture feature F c The texture enhancement feature is obtained through the texture feature enhancement module Global motion features and texture enhancement features Input into the loop iterative refinement module to calculate the optical flow estimate.

[0008] Furthermore, the image feature enhancement network includes multiple groups of Mamba modules, each of which 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 expressed as:

[0009]

[0010] Among them, F t k and represents the input and output of the kth Mamba module; Indicates F t k The regularized image features are LN represents the normalization layer; f linear represents a linear function; DWConV represents a depth-separable convolutional layer; SiLU represents an activation function; represents a pixel-by-pixel dot product operation; SS2D represents a 2D selective scan operation.

[0011] Furthermore, the 2D selective scanning operation SS2D is specifically performed 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] Among them, 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 projection of the input features respectively; w Δ ,w B ,w C Represents three different parameter matrices; A represents the state matrix, and are the discrete parameters of A and B respectively; D represents the skip link, h t is the intermediate state at time t, and y represents the output feature of the SS2D operation.

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

[0017]

[0018]

[0019] in, and They represent 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 feature and motion texture feature F c Through the texture feature enhancement module, the texture feature enhancement module predicts a set of scale weights from the global motion features, and then uses them to achieve global adjustment of the entire dense displacement, and finally obtains the texture enhancement feature The specific operation is expressed as:

[0021]

[0022] Among them, Sigmoid represents normalization operation; Conv represents convolution operation; ReLU represents activation function; AVG represents average pooling.

[0023] Furthermore, the iterative refinement module first transforms the motion feature F m , global motion features and texture enhancement features It is spliced ​​along the channel dimension and then input into the gated recurrent unit to iteratively update the predicted optical flow, and then the corresponding optical flow is upsampled to obtain the original resolution optical flow map.

[0024] Furthermore, the structure of the texture feature extraction network module is the same as that of the image feature enhancement network module.

[0025] A computer device / equipment / system comprises 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 optical flow estimation method based on the state space model.

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

[0027] A computer program product comprises a computer program / instruction, which implements the steps of the above-mentioned optical flow estimation method based on the state space model when executed by a processor.

[0028] The beneficial effects of the present invention are:

[0029] The present invention designs an image feature enhancement network and a Mamba feature cross-enhancement module based on a CNN-Mamba hybrid module, and designs a neural network of a CNN-Mamba hybrid cyclic regression framework. The present invention utilizes the image feature enhancement network of the CNN-Mamba hybrid module to efficiently extract global information of an input image, which can not only explore long-range dependencies between pixels of the input image, but also maintain linear complexity. The present invention utilizes the Mamba feature cross-enhancement module to obtain global motion features, aggregate global texture information within a motion boundary, and enhance information on texture features, thereby ultimately improving the accuracy of optical flow estimation. BRIEF DESCRIPTION OF THE DRAWINGS

[0030] Figure 1 It is a flow chart of the overall concept of the present invention.

[0031] Figure 2 It is the overall structural diagram of the neural network of the CNN-Mamba hybrid recurrent regression framework in the present invention.

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

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

[0034] Figure 5 This is a structural diagram of the 2D selection search (SS2D) in the present invention. DETAILED DESCRIPTION

[0035] The present invention is further described below in conjunction with the accompanying drawings.

[0036] The present invention proposes an optical flow estimation method based on a state space model, and the optical flow estimation method adopts a neural network of a CNN-Mamba hybrid recurrent regression framework. The present invention constructs an image feature enhancement network based on a CNN-Mamba hybrid module, which is composed of a feature extraction convolutional network and an image feature enhancement module composed of Mamba modules. It can not only efficiently extract the global information of the input image, but also maintain linear complexity and reduce computational overhead. In order to better aggregate texture information within the motion boundary, the present invention constructs a Mamba feature cross enhancement module. The module is composed of a Mamba motion enhancement module and a texture feature enhancement module. The Mamba motion enhancement module uses a global feature search mechanism to explore the long-term dependency 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] The present invention constructs a neural network of a CNN-Mamba hybrid cyclic regression framework, comprising 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 cyclic iterative refinement module, wherein the Mamba feature cross enhancement module comprises a Mamba motion enhancement module and a texture feature enhancement module.

[0038] Two consecutive frames of image I t and I t+1 The image feature enhancement network based on the CNN-Mamba hybrid module first extracts local features from the input image by the feature extraction convolution network and downsamples the image features to obtain the image feature F t and F t+1 ; Then F t and F t+1Input image feature enhancement module, the image feature enhancement module contains N Mamba modules, each Mamba module captures the long-range dependency 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 Perform pixel-by-pixel dot product operations to obtain the global correlation volume, and then obtain the motion feature F through the motion feature encoding module. m , the texture feature extraction network module based on the CNN-Mamba hybrid module is used to extract the input image I t Extract texture features F c ; Motion characteristics F m and texture feature F c Get global motion features through Mamba motion enhancement module Global motion features and motion texture feature F c Through the texture feature enhancement module, the texture feature enhancement module predicts a set of scale weights from the global motion features, and then uses them to achieve global adjustment of the entire dense displacement, and finally obtains the texture enhancement feature

[0040]

[0041] Global motion features and texture enhancement features The input is sent to the iterative refinement module, which first transforms the motion feature F m , global motion features and texture enhancement features It is spliced ​​along the channel dimension and then input into the gated recurrent unit to iteratively update the predicted optical flow, and then the corresponding optical flow is upsampled to obtain the original resolution optical flow map.

[0042] Embodiment 1:

[0043] like Figure 3 As shown, two consecutive frames of image I are input 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 residual module consists of two 3×3 convolutions with a stride of 1. The image feature F t and F t+1 . Where H, W represent the height and width of the input image respectively, C represents the number of channels of the input image, and then the feature image Ft and F t+1 Through the image feature enhancement module composed of N Mamba modules, each Mamba module captures the long-range dependencies between each pixel in the input image features. The Mamba module operation is defined as:

[0044]

[0045]

[0046] Among them, F t k and represents the input and output of the kth Mamba module; Indicates F t k Regularized image features, LN represents the standardization layer; f linear represents a linear function; DWConV represents a depth-separable convolutional layer; SiLU represents an activation function; represents the pixel-by-pixel dot product operation; SS2D represents the 2D selective scan operation, which is 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] Among them, x represents the input feature in the SS2D operation, x = [x1, x2, ..., x t ,…,x L ], whose length is L; Δ represents the time scale parameter; B and C represent the linear projection of the input features respectively; w Δ ,w B ,w C Represents three different parameter matrices; A represents the state matrix, and are the discrete parameters of A and B respectively; D represents the skip link, h t is the intermediate state at time t, and y represents the output feature of the SS2D operation.

[0051] Enhance the image features output from the network based on the image feature enhancement network of the CNN-Mamba hybrid module and Perform pixel-by-pixel dot product operations to obtain the global correlation volume, and then obtain the motion feature F through the motion feature encoding module. m , its motion feature encoding module consists of 4 CNN blocks, each of which contains a convolution with a stride of 1, a convolution kernel of size 3×3, and a ReLU activation function.

[0052] Image I t Input to the texture feature extraction network based on the CNN-Mamba hybrid module to obtain the texture feature F c , its texture feature extraction network consists of a stride of 2, size 7×7 convolution, a 1×1 convolution, three residual modules, and an image feature enhancement module composed of N Mamba modules. Each residual module consists of two stride of 1, size 3×3 convolutions. Figure 4 As shown in the figure, the Mamba feature cross enhancement module consists of the Mamba motion enhancement module and the texture feature enhancement module. m and texture feature F c The Mamba motion enhancement module explores the long-range dependencies between pixels and then aggregates the learned weights into the motion features F. m Then, a skip link is used to obtain the global motion feature, which is done as follows:

[0053]

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

[0055]

[0056] Among them, Sigmoid represents the normalization operation, Conv represents the convolution operation, ReLU represents the activation function, and AVG represents the average pooling. Represents texture enhancement feature.

[0057] The present invention constructs a neural network of a CNN-Mamba hybrid cyclic regression framework, and the network architecture includes an image feature enhancement network based on a CNN-Mamba hybrid module, a Mamba motion aggregation module, a related volume encoding module, and a cyclic iterative refinement module. The global motion features and texture enhancement features output by the Mamba feature cross enhancement module are input into the cyclic iterative refinement module to calculate the corresponding optical flow. In the cyclic iterative refinement module, the motion feature F is firstly converted into m , global motion features and texture enhancement features It is concatenated along the channel dimension and then input into the gated recurrent unit (GRU) to iteratively update the predicted optical flow, and then the corresponding optical flow is upsampled to obtain the original resolution optical flow map.

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

[0059] The above description is only a preferred embodiment of the present invention and is not intended to limit the present invention. For those skilled in the art, the present invention may have various modifications and variations. Any modification, equivalent replacement, improvement, etc. made within the spirit and principle of the present invention shall be included in the protection scope of the present invention.

Claims

1. An optical flow estimation method based on a state space model, characterized in that: Two consecutive frames of image I t and I t+1 Input into the pre-trained CNN-Mamba hybrid recurrent regression framework neural network; The neural network of the CNN-Mamba hybrid cyclic 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 cyclic 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; Based on the CNN convolutional network module, the input image I t and I t+1 Extract local features separately and downsample image features to obtain image features F t and F t+1 ; Image feature F t and F t+1 The image features are enhanced through the Mamba module to obtain the enhanced image features. and The texture feature extraction network module is used to extract the input image I t Extract texture features F c ; and Perform pixel-by-pixel dot product operations to obtain the global correlation volume, and then obtain the motion feature F through the motion feature encoding module. m , motion feature F m and texture feature F c Get global motion features through Mamba motion enhancement module Global motion features and motion texture feature F c The texture enhancement feature is obtained through the texture feature enhancement module Global motion features and texture enhancement features Input into the loop iterative refinement module to calculate the optical flow estimate.

2. The optical flow estimation method based on the state space model according to claim 1, characterized in that: The image feature enhancement network includes multiple groups of Mamba modules, each of which 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 expressed as: Among them, F t k and represents the input and output of the kth Mamba module; Indicates F t k The regularized image features are LN represents the normalization layer; f linear represents a linear function; DWConV represents a depth-separable convolutional layer; SiLU represents an activation function; represents a pixel-by-pixel dot product operation; SS2D represents a 2D selective scan operation.

3. The optical flow estimation method based on the state space model according to claim 2, characterized in that: The 2D selective scanning operation SS2D is shown, and the specific operation is expressed 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 ] Among them, 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 projection of the input features respectively; w Δ ,w B ,w C Represents three different parameter matrices; A represents the state matrix, and are the discrete parameters of A and B respectively; D represents the skip link, h t is the intermediate state at time t, and y represents the output feature of the SS2D operation.

4. The optical flow estimation method based on the state space model according to claim 2, characterized in that: The motion feature F m and texture feature F c The Mamba motion enhancement module explores the long-range dependencies between pixels and then aggregates the learned weights into the motion features F. m Then, the global motion feature is obtained through skip links. The specific operation is expressed as: in, and They represent the regularized motion features and texture features respectively; Z represents the linear projection of the input features; GS represents global feature search.

5. The optical flow estimation method based on the state space model according to claim 2, characterized in that: The global motion feature and motion texture feature F c Through the texture feature enhancement module, the texture feature enhancement module predicts a set of scale weights from the global motion features, and then uses them to achieve global adjustment of the entire dense displacement, and finally obtains the texture enhancement feature The specific operation is expressed as: Among them, Sigmoid represents normalization operation; Conv represents convolution operation; ReLU represents activation function; AVG represents average pooling.

6. The optical flow estimation method based on the state space model according to claim 1, characterized in that: The iterative refinement module first transforms the motion feature F m , global motion features and texture enhancement features It is concatenated along the channel dimension and then input into the gated recurrent unit to iteratively update the predicted optical flow, and then the corresponding optical flow is upsampled to obtain the original resolution optical flow map.

7. The optical flow estimation method based on the state space model according to claim 1, characterized in that: The structure of the texture feature extraction network module is the same as that of the image feature enhancement network module.

8. A computer device / equipment / system 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 7.

9. A computer-readable storage medium having a computer program / instruction stored thereon, characterized in that: When the computer program / instructions are executed by a processor, the steps of the method according to any one of claims 1 to 7 are implemented.

10. A computer program product comprising a computer program / instructions, characterized in that: When the computer program / instructions are executed by a processor, the steps of the method according to any one of claims 1 to 7 are implemented.

Images