Video rain removal method and device based on RWKV model, equipment and medium

Abstract

The invention discloses a video rain removal method, device and equipment based on an RWKV model and a medium, and belongs to the field of video rain removal, and the method comprises the steps: carrying out the feature extraction of a video frame sequence containing rain, and obtaining a coding feature sequence; based on the coding feature sequence, global low-frequency degradation features in a single frame are modeled through a spatial RWKV model, and a spatial feature sequence is obtained; wherein the space RWKV model comprises a wavelet transform shift layer, a space mixing layer and a channel mixing layer; based on the spatial feature sequence, modeling high-frequency time sequence dependence features among multiple frames through a time RWKV model to obtain a spatial-temporal feature sequence; wherein the time RWKV model comprises a tubular embedding layer, a space mixing layer, a channel mixing layer and an inverse tubular embedding layer; and decoding the spatio-temporal feature sequence and adding the decoded spatio-temporal feature sequence with the video frame sequence to obtain a rain-removed video frame sequence. Therefore, by implementing the method and the device, the video rain removal effect can be improved.

Description

Technical Field

[0001] The present invention relates to the field of video deraining, and in particular to a video deraining method, device, equipment and medium based on a RWKV model. Background Art

[0002] Outdoor video recordings using mobile sensors in surveillance equipment or autonomous vehicles are often adversely affected by raindrops and rain streaks. The presence of rain not only negatively impacts the visual quality of the video but also significantly degrades the performance of subsequent video processing algorithms, such as object detection, semantic segmentation, and autonomous driving. Video deraining technology aims to remove raindrops and streaks from videos recorded under rainy conditions, restoring obscured background details to enhance the reliability of systems such as autonomous driving and video surveillance. Therefore, video deraining technology is an important preprocessing step to improve the reliability and effectiveness of outdoor intelligent systems.

[0003] Traditional video deraining methods fall into two main categories: physical modeling-based methods optimize using hand-crafted priors (such as the physical and photometric properties of raindrops), but these methods rely on complex prior assumptions, are difficult to optimize, and have poor generalization capabilities. Deep learning-based methods primarily employ convolutional neural networks (CNNs), recurrent neural networks (RNNs), or Transformer architectures. CNNs and RNNs exploit temporal information through optical flow estimation or deformable convolutions, but their spatiotemporal receptive fields are limited and they cannot model global degradation. The Transformer architecture captures long-term dependencies through multi-head self-attention (MSA), but its computational complexity is quadratic, making it difficult to process long video sequences. In recent years, the Receptance Weighted Key Value (RWKV) model has shown promise in modeling long sequences due to its linear computational complexity. However, its original design for one-dimensional sequences makes it difficult to effectively capture the spatiotemporal dependencies of three-dimensional video data and suffers from limitations in preserving low-frequency features and recovering high-frequency details. Summary of the Invention

[0004] The present invention provides a video deraining method, apparatus, device and medium based on a RWKV model, which can improve the effect of video deraining.

[0005] An embodiment of the present invention provides a video deraining method based on the RWKV model, comprising:

[0006] Extract features from the video frame sequence containing rain to obtain a coding feature sequence;

[0007] Based on the coding feature sequence, a global low-frequency degradation feature in a single frame is modeled by a spatial RWKV model to obtain a spatial feature sequence; wherein the spatial RWKV model includes a wavelet transform shift layer, a spatial mixing layer, and a channel mixing layer;

[0008] Based on the spatial feature sequence, a temporal RWKV model is used to model the high-frequency temporal dependency features between multiple frames to obtain a spatiotemporal feature sequence; wherein the temporal RWKV model includes a tubular embedding layer, a spatial mixing layer, a channel mixing layer, and an inverse tubular embedding layer;

[0009] The spatiotemporal feature sequence is decoded and added to the video frame sequence to obtain a rain-removed video frame sequence.

[0010] The embodiments of the present invention extract features from a sequence of rain-containing video frames, providing a data foundation for subsequent low-frequency and high-frequency feature processing. The spatial RWKV model is used to model global low-frequency degradation features within a single frame, improving the ability to retain low-frequency information. The temporal RWKV model is used to model high-frequency temporal dependency features between multiple frames, improving the ability to recover high-frequency details. By decoding the feature sequence, a video frame sequence after rain removal is obtained. Compared to the existing techniques that suffer from poor video rain removal performance, the present application improves video rain removal performance by retaining low-frequency information and recovering high-frequency details.

[0011] Furthermore, the feature extraction of the rain-containing video frame sequence includes:

[0012] Extracting shallow features of each frame of the video frame sequence; wherein the shallow features include multiple scales;

[0013] Based on the shallow features, scale fusion is performed through the lightweight head to obtain a coding feature sequence.

[0014] The embodiment of the present invention can provide a data basis for subsequent low-frequency and high-frequency feature processing by extracting features from a video frame sequence containing rain.

[0015] Furthermore, the modeling of global low-frequency degradation features within a single frame using a spatial RWKV model includes:

[0016] Flattening each frame of the coding feature sequence into a one-dimensional coding feature;

[0017] Based on each of the one-dimensional coding features, frequency division processing is performed through a wavelet transform shift layer to obtain multiple first one-dimensional coding features;

[0018] Based on each of the first one-dimensional coding features, performing global enhancement through a spatial mixing layer to obtain a plurality of second one-dimensional coding features;

[0019] Based on each of the second one-dimensional coding features and the coding feature sequence, performing channel enhancement through a channel mixing layer to obtain a plurality of third one-dimensional coding features;

[0020] Each of the second one-dimensional coding features and the third one-dimensional coding features is correspondingly added and inversely flattened to obtain a spatial feature sequence.

[0021] The embodiment of the present invention uses a spatial RWKV model to model the global low-frequency degradation characteristics within a single frame, thereby improving the ability to retain low-frequency information.

[0022] Furthermore, the frequency division processing by wavelet transform shift layer includes:

[0023] Performing a two-dimensional discrete wavelet transform on the one-dimensional coding feature to obtain a low-frequency component and a high-frequency component;

[0024] Performing downsampling and convolution operations on the low-frequency components to obtain low-frequency features;

[0025] Performing a small kernel convolution operation on the high-frequency component to obtain high-frequency features;

[0026] Perform inverse wavelet transform on the low-frequency features and the high-frequency features to obtain a first one-dimensional coding feature.

[0027] The embodiment of the present invention performs frequency division processing through wavelet transform shift layers, which can enhance the ability to retain low-frequency features and avoid the problem of blurred details caused by full-band convolution.

[0028] Furthermore, the global enhancement by the spatial mixing layer includes:

[0029] Performing linear projection on the first one-dimensional coding feature to obtain a first key, a first value, and a first gating signal;

[0030] Global attention weighting is performed on the first key and the first value, and the output intensity is controlled by the first gating signal to obtain a second one-dimensional encoding feature.

[0031] The embodiment of the present invention performs global enhancement through a spatial mixing layer, and can repair the global low-frequency degradation characteristics of a single-frame image.

[0032] Furthermore, channel enhancement is performed through a channel mixing layer, including:

[0033] Obtaining a second key and a second gating signal by linear projection based on the second one-dimensional coding feature and the coding feature sequence;

[0034] The second key is subjected to nonlinear transformation and the output intensity is controlled by the second gating signal to obtain a third one-dimensional coding feature.

[0035] The embodiment of the present invention performs channel enhancement through a channel mixing layer, and can optimize the local channel feature expression of a single-frame image.

[0036] Furthermore, the method of modeling high-frequency temporal dependency features between multiple frames through the temporal RWKV model includes:

[0037] Based on the spatial feature sequence, performing spatiotemporal window division through a tubular embedding layer to obtain a label sequence;

[0038] Based on the marker sequence, performing global enhancement through a spatial mixing layer to obtain a first marker sequence;

[0039] Based on the first marker sequence and the marker sequence, channel enhancement is performed through a channel mixing layer to obtain a second marker sequence;

[0040] The first label sequence and the second label sequence are added together and passed through an inverse tubular embedding layer to obtain a spatiotemporal feature sequence.

[0041] The embodiment of the present invention uses a temporal RWKV model to model high-frequency temporal dependency characteristics between multiple frames, thereby improving the ability to restore high-frequency details.

[0042] Another embodiment of the present invention further provides a video deraining device based on the RWKV model, comprising: a feature encoding module, a spatial feature module, a spatiotemporal feature module, and a feature decoding module;

[0043] The feature encoding module is used to extract features from a video frame sequence containing rain to obtain a coded feature sequence;

[0044] The spatial feature module is used to model the global low-frequency degradation features in a single frame through a spatial RWKV model based on the coding feature sequence to obtain a spatial feature sequence; wherein the spatial RWKV model includes a wavelet transform shift layer, a spatial mixing layer and a channel mixing layer;

[0045] The spatiotemporal feature module is configured to model the high-frequency temporal dependency features between multiple frames using a temporal RWKV model based on the spatial feature sequence to obtain a spatiotemporal feature sequence; wherein the temporal RWKV model comprises a tubular embedding layer, a spatial mixing layer, a channel mixing layer, and an inverse tubular embedding layer;

[0046] The feature decoding module is used to decode the spatiotemporal feature sequence and add it to the video frame sequence to obtain a rain-free video frame sequence.

[0047] Another embodiment of the present invention further provides a terminal device, comprising: a processor, a memory, and a computer program stored in the memory and configured to be executed by the processor. When the processor executes the computer program, it implements the steps of a video deraining method based on the RWKV model of the present invention.

[0048] Another embodiment of the present invention further provides a computer-readable storage medium item, comprising: a stored computer program, which, when the computer program is executed, controls the device where the computer-readable storage medium is located to execute the steps of a video deraining method based on the RWKV model of the present invention. BRIEF DESCRIPTION OF THE DRAWINGS

[0049] Figure 1 A schematic diagram of a flow chart of an embodiment of a video deraining method based on the RWKV model provided by the present invention;

[0050] Figure 2 A schematic diagram of the structure of an embodiment of a video rain removal neural network based on the RWKV model provided by the present invention;

[0051] Figure 3 A comparative experimental result diagram of an embodiment of a video rain removal neural network based on the RWKV model provided by the present invention;

[0052] Figure 4 This is a comparative experimental effect diagram of an embodiment of a video deraining neural network based on the RWKV model provided by the present invention;

[0053] Figure 5 This is a diagram showing the efficiency comparison results of an embodiment of a video rain removal neural network based on the RWKV model provided by the present invention;

[0054] Figure 6 This is a structural diagram of an embodiment of a video rain removal device based on the RWKV model provided by the present invention. DETAILED DESCRIPTION

[0055] To make the objectives, technical solutions, and advantages of this application more clear, the technical solutions in this application will be clearly and completely described below in conjunction with the accompanying drawings in the embodiments of this application. Obviously, the described embodiments are only part of the embodiments of this application, not all of them. Based on the embodiments in this application, all other embodiments obtained by ordinary technicians in this field without making any creative efforts are within the scope of protection of this application.

[0056] Unless otherwise defined, all technical and scientific terms used herein have the same meaning as commonly understood by those skilled in the art to which this application belongs; the terms used herein are only for the purpose of describing specific embodiments and are not intended to limit this application; the terms "including" and "having" and any variations thereof in the specification and claims of this application and the above-mentioned figure descriptions are intended to cover non-exclusive inclusions.

[0057] In the description of the embodiments of this application, the technical terms "first" and "second" are used only to distinguish different objects and should not be understood to indicate or imply relative importance or implicitly specify the quantity, specific order, or primary and secondary relationship of the indicated technical features. In the description of the embodiments of this application, the meaning of "plurality" is more than two, unless otherwise clearly and specifically defined.

[0058] References herein to "embodiments" mean that a particular feature, structure, or characteristic described in connection with the embodiments may be included in at least one embodiment of the present application. The appearance of this phrase in various places in the specification does not necessarily refer to the same embodiment, nor does it constitute an independent or alternative embodiment that is mutually exclusive of other embodiments. It is understood, both explicitly and implicitly, by those skilled in the art that the embodiments described herein may be combined with other embodiments.

[0059] In the description of the embodiments of this application, the term "and / or" is simply a description of the association relationship between associated objects, indicating that three relationships can exist. For example, A and / or B can represent the following three situations: A exists alone, A and B exist simultaneously, and B exists alone. In addition, the character " / " in this document generally indicates that the associated objects are in an "or" relationship.

[0060] In the description of the embodiments of the present application, the term "multiple" refers to more than two (including two). Similarly, "multiple groups" refers to more than two groups (including two groups), and "multiple pieces" refers to more than two pieces (including two pieces).

[0061] See also Figure 1 To solve the problem of poor video deraining effect in the prior art, an embodiment of the present invention provides a video deraining method based on the RWKV model, including steps S101 to S104:

[0062] Step S101 : extracting features from a video frame sequence containing rain to obtain a coding feature sequence.

[0063] Specifically, a pre-trained ConvNeXt model can be used as the backbone network to extract multi-scale shallow features of each frame in the video frame sequence, where the number of channels can be adjusted according to the network layer; a lightweight head (LightweightHead) is used to fuse shallow features of different scales to generate a coded feature sequence.

[0064] Step S102 : Based on the coding feature sequence, a spatial RWKV model is used to model the global low-frequency degradation features in a single frame to obtain a spatial feature sequence; wherein the spatial RWKV model includes a wavelet transform shift layer, a spatial mixing layer, and a channel mixing layer.

[0065] Specifically, the spatial RWKV model can be used to model the global low-frequency degradation features within a single frame, such as the low-frequency structure covered by rain streaks. The spatial RWKV model includes a wavelet transform shift layer, a spatial mixing layer, and a channel mixing layer. The wavelet transform shift layer is used to perform frequency division processing on the coding features of each frame, which can enhance the low-frequency feature retention capability of a single-frame image. The spatial mixing layer is used to perform global enhancement on the coding features of each frame, which can repair the global low-frequency degradation features of a single-frame image. The channel mixing layer is used to perform channel enhancement on the coding features of each frame, which can optimize the local channel feature expression of a single-frame image.

[0066] Step S103: Based on the spatial feature sequence, a temporal RWKV model is used to model the high-frequency temporal dependency features between multiple frames to obtain a spatiotemporal feature sequence; wherein the temporal RWKV model includes a tubular embedding layer, a spatial mixing layer, a channel mixing layer, and an inverse tubular embedding layer.

[0067] Specifically, the temporal RWKV model can be used to model high-frequency temporal dependency features between multiple frames, such as the local motion of dynamic raindrops. The temporal RWKV model includes a tubular embedding layer, a spatial mixing layer, a channel mixing layer, and an inverse tubular embedding layer. The tubular embedding layer is used to divide the spatial feature sequence into spatiotemporal windows, which can improve the high-frequency detail recovery capability of the video frame sequence. The spatial mixing layer is used to globally enhance the spatial feature sequence and repair the global low-frequency degradation features of the video frame sequence. The channel mixing layer is used to perform channel enhancement on the spatial feature sequence and optimize the local channel feature expression of the video frame sequence. The inverse tubular embedding layer is used to reorganize the output sequence features back to the original volume features.

[0068] Step S104 : Decode the spatiotemporal feature sequence and add it to the video frame sequence to obtain a rain-removed video frame sequence.

[0069] Specifically, a three-dimensional convolution layer and an upsampling layer may be used to convert the spatiotemporal feature sequence into a residual image sequence, and then the residual image sequence is added to a rain-containing video frame sequence to obtain a rain-free video frame sequence.

[0070] The embodiments of the present invention extract features from a sequence of rain-containing video frames, providing a data foundation for subsequent low-frequency and high-frequency feature processing. The spatial RWKV model is used to model global low-frequency degradation features within a single frame, improving the ability to retain low-frequency information. The temporal RWKV model is used to model high-frequency temporal dependency features between multiple frames, improving the ability to recover high-frequency details. By decoding the feature sequence, a video frame sequence after rain removal is obtained. Compared to the existing techniques that suffer from poor video rain removal performance, the present application improves video rain removal performance by retaining low-frequency information and recovering high-frequency details.

[0071] Optionally, in this embodiment of the present invention, the feature extraction of the rain-containing video frame sequence includes:

[0072] Extracting shallow features of each frame of the video frame sequence; wherein the shallow features include multiple scales;

[0073] Based on the shallow features, scale fusion is performed through the lightweight head to obtain a coding feature sequence.

[0074] The embodiment of the present invention can provide a data basis for subsequent low-frequency and high-frequency feature processing by extracting features from a video frame sequence containing rain.

[0075] Optionally, in this embodiment of the present invention, modeling the global low-frequency degradation features in a single frame by using a spatial RWKV model includes:

[0076] Flattening each frame of the coding feature sequence into a one-dimensional coding feature;

[0077] Based on each of the one-dimensional coding features, frequency division processing is performed through a wavelet transform shift layer to obtain multiple first one-dimensional coding features;

[0078] Based on each of the first one-dimensional coding features, performing global enhancement through a spatial mixing layer to obtain a plurality of second one-dimensional coding features;

[0079] Based on each of the second one-dimensional coding features and the coding feature sequence, performing channel enhancement through a channel mixing layer to obtain a plurality of third one-dimensional coding features;

[0080] Each of the second one-dimensional coding features and the third one-dimensional coding features is correspondingly added and inversely flattened to obtain a spatial feature sequence.

[0081] The embodiment of the present invention uses a spatial RWKV model to model the global low-frequency degradation characteristics within a single frame, thereby improving the ability to retain low-frequency information.

[0082] Optionally, in the embodiment of the present invention, the frequency division processing by using the wavelet transform shift layer includes:

[0083] Performing a two-dimensional discrete wavelet transform on the one-dimensional coding feature to obtain a low-frequency component and a high-frequency component;

[0084] Performing downsampling and convolution operations on the low-frequency components to obtain low-frequency features;

[0085] Performing a small kernel convolution operation on the high-frequency component to obtain high-frequency features;

[0086] Perform inverse wavelet transform on the low-frequency features and the high-frequency features to obtain a first one-dimensional coding feature.

[0087] Specifically, the process of frequency division processing through wavelet transform shift layer is as follows:

[0088] WT-Shift(X)=IWT(Conv(K,WT(X)));

[0089] Among them, WT is a two-dimensional discrete wavelet transform; IWT is an inverse wavelet transform; and K is a depth convolution kernel weight. First, a two-dimensional discrete wavelet transform is performed on the one-dimensional coding feature to obtain a low-frequency component (LL) and three high-frequency components (LH / HL / HH); wherein the low-frequency component contains the global background structure, and the high-frequency component contains detail information such as edges and textures. Then, after downsampling the low-frequency component, a lightweight convolution (such as a 3×3 depth convolution) is applied to reduce the amount of calculation; a small kernel convolution (such as a 1×1 depth convolution or a 3×3 depth convolution) is applied to the high-frequency component to extract local details. Finally, the features obtained after the frequency division processing are merged into the output feature through an inverse wavelet transform.

[0090] The embodiment of the present invention performs frequency division processing through wavelet transform shift layers, which can enhance the ability to retain low-frequency features and avoid the problem of blurred details caused by full-band convolution.

[0091] Optionally, in this embodiment of the present invention, performing global enhancement through a spatial mixing layer includes:

[0092] Performing linear projection on the first one-dimensional coding feature to obtain a first key, a first value, and a first gating signal;

[0093] Global attention weighting is performed on the first key and the first value, and the output intensity is controlled by the first gating signal to obtain a second one-dimensional encoding feature.

[0094] Specifically, the input one-dimensional coding features are first linearly projected into keys, values, and gating signals (Receptance); then the Bi-WKV bidirectional linear attention mechanism can be used to perform global attention weighting on the keys and values; wherein, the Bi-WKV bidirectional linear attention mechanism models global spatial dependencies through relative position bias and learnable parameters; finally, the reception ratio of the attention results is controlled by the gating signal after Sigmoid activation, and the globally enhanced one-dimensional coding features are output.

[0095] The embodiment of the present invention performs global enhancement through a spatial mixing layer, and can repair the global low-frequency degradation characteristics of a single-frame image.

[0096] Optionally, in an embodiment of the present invention, performing channel enhancement through a channel mixing layer includes:

[0097] Obtaining a second key and a second gating signal by linear projection based on the second one-dimensional coding feature and the coding feature sequence;

[0098] The second key is subjected to nonlinear transformation and the output intensity is controlled by the second gating signal to obtain a third one-dimensional coding feature.

[0099] Specifically, the input one-dimensional coding features are first linearly projected into keys and gating signals; then the keys are nonlinearly transformed (squared ReLU and projection) to obtain values; wherein the nonlinear transformation is used to enhance the nonlinear interaction between channels; finally, the intensity is controlled by the gating signal after Sigmoid activation, and the one-dimensional coding features after channel enhancement are output.

[0100] The embodiment of the present invention performs channel enhancement through a channel mixing layer, and can optimize the local channel feature expression of a single-frame image.

[0101] Optionally, in this embodiment of the present invention, modeling high-frequency temporal dependency features between multiple frames using a temporal RWKV model includes:

[0102] Based on the spatial feature sequence, performing spatiotemporal window division through a tubular embedding layer to obtain a label sequence;

[0103] Based on the marker sequence, performing global enhancement through a spatial mixing layer to obtain a first marker sequence;

[0104] Based on the first marker sequence and the marker sequence, channel enhancement is performed through a channel mixing layer to obtain a second marker sequence;

[0105] The first label sequence and the second label sequence are added together and passed through an inverse tubular embedding layer to obtain a spatiotemporal feature sequence.

[0106] Specifically, the input feature sequence is first reduced in dimensionality through downsampling convolution; then the tubular embedding mechanism (Tubelet Embedding) is used to divide the reduced input feature sequence into spatiotemporal windows to improve the ability to recover high-frequency details; the token sequence reorganized by the tubular embedding mechanism is enhanced in turn through the spatial mixing layer and the channel mixing layer to retain the correlation between adjacent frames and spatial regions in the Token sequence, avoiding the spatial dispersion caused by one-dimensional serialization; finally, the output feature sequence volume is reorganized back to the input feature sequence volume through the inverse tubular embedding mechanism.

[0107] In one embodiment, the shape of the input feature sequence of the tubular embedding mechanism is N∈R T×H×W×C ; Where T is the time dimension, H and W are the height dimension and width dimension respectively, and C is the channel dimension; the shape of the output feature sequence obtained by reorganizing through the tubular embedding mechanism is in, and are the number of tags in the height dimension and width dimension respectively; h and w are the height dimension and width dimension of the spatiotemporal window respectively; d is the tag dimension.

[0108] The embodiment of the present invention uses a temporal RWKV model to model high-frequency temporal dependency characteristics between multiple frames, thereby improving the ability to restore high-frequency details.

[0109] like Figure 2 As shown, based on the above method embodiment, an embodiment of a neural network is provided; the neural network includes: an encoding module (Encoder), a spatial RWKV module (Spatial RWKV Block), a temporal RWKV module (Temporal RWKV Block) and a decoding module (Decoder);

[0110] The encoding module is used to extract features from the rain-containing video frame sequence to obtain an encoding feature sequence, which is equivalent to executing step S101;

[0111] The spatial RWKV module is used to model the global low-frequency degradation characteristics in a single frame based on the coding feature sequence to obtain a spatial feature sequence, which is equivalent to executing step S102;

[0112] The temporal RWKV module is used to model high-frequency temporal dependency features between multiple frames based on the spatial feature sequence to obtain a spatiotemporal feature sequence, which is equivalent to executing step S103;

[0113] The decoding module is used to decode the spatiotemporal feature sequence, which is equivalent to executing step S104.

[0114] The embodiment of the present invention uses an encoding module to extract features from a rain-containing video frame sequence, which can provide a data basis for subsequent low-frequency and high-frequency feature processing; uses a spatial RWKV module to model the global low-frequency degradation features within a single frame, which can improve the ability to retain low-frequency information; uses a temporal RWKV module to model the high-frequency temporal dependency features between multiple frames, which can improve the ability to restore high-frequency details; and uses a decoding module to decode the feature sequence to obtain a video frame sequence after rain removal.

[0115] Based on the above neural network embodiment, a comparative experiment embodiment is provided. The neural network is trained based on the Charbonnier loss function and the perceptual loss function. The video deraining algorithm based on the trained neural network is quantitatively compared with 12 image or video deraining algorithms on the VRDS dataset (including 72 sets of training videos and 30 sets of test videos). The specific quantitative comparison results are shown in Figure 2. Figure 3As shown in the figure, “Ours” is a video deraining algorithm based on the neural network. The video deraining algorithm based on the neural network outperforms the other 12 video deraining algorithms in terms of PSNR (the larger the value, the better the effect), SSIM (the larger the value, the better the effect), and LPIPS (the smaller the value, the better the effect). Figure 4 As shown in the figure, a visual comparison of the rain removal effects of multiple video rain removal algorithms is shown. Compared with other video rain removal algorithms that produce dark pixel artifacts (such as Restormer or BasicVSR++) or edge contour distortion problems (such as DRSformer or ViMP-Net) when removing rain marks, the video rain removal algorithm based on the neural network achieves dual optimization by effectively modeling temporal dependencies, completely eliminating rain interference while fully preserving the scene's geometric structure characteristics. Figure 5 Figure 2 shows a comparison of the performance and efficiency of multiple video deraining algorithms. Each algorithm was run on an NVIDIA RTX4090 GPU, inferring a five-frame video clip with a frame resolution of 256×256 pixels. The neural network-based video deraining algorithm outperformed six other video deraining algorithms in both GFLOPs (lower values, better performance) and Runtime (lower values, better performance), demonstrating that the neural network-based video deraining algorithm not only excels in deraining performance but also maintains a minimal computational overhead.

[0116] like Figure 6 As shown, based on the above method embodiment, a corresponding device embodiment is provided;

[0117] An embodiment of the present invention provides a video deraining device based on the RWKV model, comprising: a feature encoding module 601, a spatial feature module 602, a spatiotemporal feature module 603, and a feature decoding module 604;

[0118] The feature encoding module 601 is used to extract features from a video frame sequence containing rain to obtain a coded feature sequence;

[0119] The spatial feature module 602 is configured to model the global low-frequency degradation features within a single frame using a spatial RWKV model based on the coding feature sequence to obtain a spatial feature sequence; wherein the spatial RWKV model includes a wavelet transform shift layer, a spatial mixing layer, and a channel mixing layer;

[0120] The spatiotemporal feature module 603 is configured to model the high-frequency temporal dependency features between multiple frames using a temporal RWKV model based on the spatial feature sequence to obtain a spatiotemporal feature sequence; wherein the temporal RWKV model includes a tubular embedding layer, a spatial mixing layer, a channel mixing layer, and an inverse tubular embedding layer;

[0121] The feature decoding module 604 is configured to decode the spatiotemporal feature sequence and add the decoded feature sequence to the video frame sequence to obtain a rain-free video frame sequence.

[0122] Optionally, in the embodiment of the present invention, the feature encoding module 601 includes: a feature extraction submodule and a feature fusion submodule;

[0123] The feature extraction submodule is used to extract shallow features of each frame in the video frame sequence; wherein the shallow features include multiple scales;

[0124] The feature fusion submodule is used to perform scale fusion based on each of the shallow features through a lightweight head to obtain a coding feature sequence.

[0125] The embodiment of the present invention can provide a data basis for subsequent low-frequency and high-frequency feature processing by extracting features from a video frame sequence containing rain.

[0126] Optionally, in an embodiment of the present invention, the spatial feature module 602 includes: a feature flattening submodule, a wavelet transform submodule, a first spatial mixing submodule, a first channel mixing submodule and a feature inverse flattening submodule;

[0127] The feature flattening submodule is used to flatten each frame of the coding feature sequence into a one-dimensional coding feature;

[0128] The wavelet transform submodule is configured to perform frequency division processing through a wavelet transform shift layer based on each of the one-dimensional coding features to obtain a plurality of first one-dimensional coding features;

[0129] The first spatial mixing submodule is configured to perform global enhancement based on each of the first one-dimensional coding features through a spatial mixing layer to obtain a plurality of second one-dimensional coding features;

[0130] The first channel mixing submodule is configured to perform channel enhancement through a channel mixing layer based on each of the second one-dimensional coding features and the coding feature sequence to obtain a plurality of third one-dimensional coding features;

[0131] The feature inverse flattening submodule is used to add and inverse flatten the second one-dimensional coding features and the third one-dimensional coding features respectively to obtain a spatial feature sequence.

[0132] The embodiment of the present invention uses a spatial RWKV model to model the global low-frequency degradation characteristics within a single frame, thereby improving the ability to retain low-frequency information.

[0133] Optionally, in an embodiment of the present invention, the wavelet transform submodule includes: a wavelet transform unit, a low-frequency component unit, a high-frequency component unit and an inverse wavelet transform unit;

[0134] The wavelet transform unit is used to perform a two-dimensional discrete wavelet transform on the one-dimensional coding feature to obtain a low-frequency component and a high-frequency component;

[0135] The low-frequency component unit is used to perform downsampling and convolution operations on the low-frequency component to obtain low-frequency features;

[0136] The high-frequency component unit is used to perform a small-kernel convolution operation on the high-frequency component to obtain high-frequency features;

[0137] The inverse wavelet transform unit is used to perform inverse wavelet transform on the low-frequency features and the high-frequency features to obtain a first one-dimensional coding feature.

[0138] The embodiment of the present invention performs frequency division processing through wavelet transform shift layers, which can enhance the ability to retain low-frequency features and avoid the problem of blurred details caused by full-band convolution.

[0139] Optionally, in an embodiment of the present invention, the first spatial mixing submodule includes: a first linear projection unit and a first gating unit;

[0140] The first linear projection unit is configured to perform linear projection on the first one-dimensional coding feature to obtain a first key, a first value, and a first gating signal;

[0141] The first gating unit is used to perform global attention weighting on the first key and the first value, and control the output intensity through the first gating signal to obtain a second one-dimensional coding feature.

[0142] The embodiment of the present invention performs global enhancement through a spatial mixing layer, and can repair the global low-frequency degradation characteristics of a single-frame image.

[0143] Optionally, in an embodiment of the present invention, the first channel mixing submodule includes: a second linear projection unit and a second gating unit;

[0144] The second linear projection unit is configured to obtain a second key and a second gating signal by linear projection based on the second one-dimensional coding feature and the coding feature sequence;

[0145] The second gating unit is used to perform nonlinear transformation on the second key and control the output intensity through the second gating signal to obtain a third one-dimensional coding feature.

[0146] The embodiment of the present invention performs channel enhancement through a channel mixing layer, and can optimize the local channel feature expression of a single-frame image.

[0147] Optionally, in an embodiment of the present invention, the spatiotemporal feature module 603 includes: a tubular embedding submodule, a second spatial mixing submodule, a second channel mixing submodule, and an inverse tubular embedding submodule;

[0148] The tubular embedding submodule is configured to perform spatiotemporal window division based on the spatial feature sequence through the tubular embedding layer to obtain a label sequence;

[0149] The second spatial mixing submodule is configured to perform global enhancement through a spatial mixing layer based on the marker sequence to obtain a first marker sequence;

[0150] The second channel mixing submodule is configured to perform channel enhancement through a channel mixing layer based on the first label sequence and the label sequence to obtain a second label sequence;

[0151] The inverse tubular embedding submodule is configured to add the first label sequence and the second label sequence, and obtain a spatiotemporal feature sequence through an inverse tubular embedding layer.

[0152] The embodiment of the present invention uses a temporal RWKV model to model high-frequency temporal dependency characteristics between multiple frames, thereby improving the ability to restore high-frequency details.

[0153] It can be understood that the above-mentioned device embodiment corresponds to the method embodiment of the present invention, and can implement a video deraining method based on the RWKV model provided by any of the above-mentioned method embodiments of the present invention.

[0154] The embodiment of the present invention uses a feature encoding module to extract features from a rain-containing video frame sequence, providing a data foundation for subsequent low-frequency and high-frequency feature processing. The spatial feature module models global low-frequency degradation features within a single frame, improving the ability to retain low-frequency information. The spatiotemporal feature module models high-frequency temporal dependency features between multiple frames, improving the ability to recover high-frequency details. The feature decoding module decodes the feature sequence to obtain a video frame sequence after rain removal. Compared to the existing technology that suffers from poor video rain removal, the present application improves the effectiveness of video rain removal by retaining low-frequency information and recovering high-frequency details.

[0155] It should be noted that the device embodiments described above are merely illustrative, and some or all of the modules may be selected according to actual needs to achieve the purpose of the present embodiment. Furthermore, in the drawings of the device embodiments provided by the present invention, the connection relationship between modules indicates that they have a communication connection, which may be implemented as one or more communication buses or signal lines. Those skilled in the art can understand and implement the present invention without inventive effort.

[0156] Based on the above-mentioned embodiment of a video deraining method based on the RWKV model, another embodiment of the present invention provides a terminal device, which includes a processor, a memory, and a computer program stored in the memory and configured to be executed by the processor. When the processor executes the computer program, a video deraining method based on the RWKV model according to any embodiment of the present invention is implemented.

[0157] For example, in this embodiment, the computer program may be divided into one or more modules, which are stored in the memory and executed by the processor to implement the present invention. The one or more module elements may be a series of computer program instruction segments capable of performing specific functions, and the instruction segments are used to describe the execution process of the computer program in the terminal device.

[0158] The terminal device may be a computing device such as a desktop computer, a notebook computer, a PDA, a cloud server, etc. The terminal device may include, but is not limited to, a processor and a memory.

[0159] The processor may be a central processing unit (CPU), other general-purpose processors, digital signal processors (DSP), application-specific integrated circuits (ASIC), field-programmable gate arrays (FPGA) or other programmable logic devices, discrete gate or transistor logic devices, discrete hardware components, etc. The general-purpose processor may be a microprocessor or any conventional processor, etc. The processor is the control center of the terminal device, connecting various parts of the entire terminal device using various interfaces and lines.

[0160] Based on the above method embodiments, another embodiment of the present invention provides a computer-readable storage medium, including a stored computer program, wherein, when the computer program is running, the device where the computer-readable storage medium is located is controlled to execute the video deraining method based on the RWKV model described in any one of the above method embodiments of the present invention.

[0161] Wherein, the module / unit integrated in the device / terminal equipment, if implemented in the form of a software functional unit and sold or used as an independent product, can be stored in a computer-readable storage medium. Based on this understanding, the present invention implements all or part of the process in the above-mentioned embodiment method, and can also be completed by instructing the relevant hardware through a computer program. The computer program can be stored in a computer-readable storage medium, and when the computer program is executed by the processor, it can implement the steps of the above-mentioned various method embodiments. Wherein, the computer program includes computer program code, and the computer program code can be in source code form, object code form, executable file or some intermediate form. The computer-readable medium may include: any entity or device that can carry the computer program code, recording medium, USB flash drive, mobile hard disk, magnetic disk, optical disk, computer memory, read-only memory (ROM), random access memory (RAM), electric carrier signal, telecommunication signal and software distribution medium, etc.

[0162] The above is a preferred embodiment of the present invention. It should be pointed out that for ordinary technicians in this technical field, several improvements and modifications can be made without departing from the principles of the present invention. These improvements and modifications are also considered to be within the scope of protection of the present invention.

Claims

1. A video deraining method based on the RWKV model, characterized in that: include: Extract features from the video frame sequence containing rain to obtain a coding feature sequence; Based on the coding feature sequence, a global low-frequency degradation feature in a single frame is modeled by a spatial RWKV model to obtain a spatial feature sequence; wherein the spatial RWKV model includes a wavelet transform shift layer, a spatial mixing layer, and a channel mixing layer; Based on the spatial feature sequence, a temporal RWKV model is used to model the high-frequency temporal dependency features between multiple frames to obtain a spatiotemporal feature sequence; wherein the temporal RWKV model includes a tubular embedding layer, a spatial mixing layer, a channel mixing layer, and an inverse tubular embedding layer; The spatiotemporal feature sequence is decoded and added to the video frame sequence to obtain a rain-removed video frame sequence.

2. The video deraining method based on the RWKV model according to claim 1, characterized in that: The feature extraction of the rain-containing video frame sequence includes: Extracting shallow features of each frame of the video frame sequence; wherein the shallow features include multiple scales; Based on the shallow features, scale fusion is performed through the lightweight head to obtain a coding feature sequence.

3. The video deraining method based on the RWKV model according to claim 1, characterized in that: The global low-frequency degradation feature within a single frame is modeled using a spatial RWKV model, including: Flattening each frame of the coding feature sequence into a one-dimensional coding feature; Based on each of the one-dimensional coding features, frequency division processing is performed through a wavelet transform shift layer to obtain multiple first one-dimensional coding features; Based on each of the first one-dimensional coding features, performing global enhancement through a spatial mixing layer to obtain a plurality of second one-dimensional coding features; Based on each of the second one-dimensional coding features and the coding feature sequence, performing channel enhancement through a channel mixing layer to obtain a plurality of third one-dimensional coding features; Each of the second one-dimensional coding features and the third one-dimensional coding features is correspondingly added and inversely flattened to obtain a spatial feature sequence.

4. The video deraining method based on the RWKV model according to claim 3, characterized in that: The frequency division processing by wavelet transform shift layer includes: Performing a two-dimensional discrete wavelet transform on the one-dimensional coding feature to obtain a low-frequency component and a high-frequency component; Performing downsampling and convolution operations on the low-frequency components to obtain low-frequency features; Performing a small kernel convolution operation on the high-frequency component to obtain high-frequency features; Perform inverse wavelet transform on the low-frequency features and the high-frequency features to obtain a first one-dimensional coding feature.

5. The video deraining method based on the RWKV model according to claim 3, characterized in that: The global enhancement is performed by a spatial mixing layer, comprising: Performing linear projection on the first one-dimensional coding feature to obtain a first key, a first value, and a first gating signal; Global attention weighting is performed on the first key and the first value, and the output intensity is controlled by the first gating signal to obtain a second one-dimensional encoding feature.

6. The video deraining method based on the RWKV model according to claim 3, characterized in that: The channel enhancement is performed by the channel mixing layer, comprising: Obtaining a second key and a second gating signal by linear projection based on the second one-dimensional coding feature and the coding feature sequence; The second key is subjected to nonlinear transformation and the output intensity is controlled by the second gating signal to obtain a third one-dimensional coding feature.

7. The video deraining method based on the RWKV model according to claim 1, characterized in that: The high-frequency temporal dependency features between multiple frames are modeled by the temporal RWKV model, including: Based on the spatial feature sequence, performing spatiotemporal window division through a tubular embedding layer to obtain a label sequence; Based on the marker sequence, performing global enhancement through a spatial mixing layer to obtain a first marker sequence; Based on the first marker sequence and the marker sequence, channel enhancement is performed through a channel mixing layer to obtain a second marker sequence; The first label sequence and the second label sequence are added together and passed through an inverse tubular embedding layer to obtain a spatiotemporal feature sequence.

8. A video deraining device based on the RWKV model, characterized in that: include: Feature encoding module, spatial feature module, spatiotemporal feature module and feature decoding module; The feature encoding module is used to extract features from a video frame sequence containing rain to obtain a coded feature sequence; The spatial feature module is used to model the global low-frequency degradation features in a single frame through a spatial RWKV model based on the coding feature sequence to obtain a spatial feature sequence; wherein the spatial RWKV model includes a wavelet transform shift layer, a spatial mixing layer and a channel mixing layer; The spatiotemporal feature module is configured to model the high-frequency temporal dependency features between multiple frames using a temporal RWKV model based on the spatial feature sequence to obtain a spatiotemporal feature sequence; wherein the temporal RWKV model comprises a tubular embedding layer, a spatial mixing layer, a channel mixing layer, and an inverse tubular embedding layer; The feature decoding module is used to decode the spatiotemporal feature sequence and add it to the video frame sequence to obtain a rain-free video frame sequence.

9. A terminal device, characterized in that: The method comprises a processor, a memory, and a computer program stored in the memory and configured to be executed by the processor. When the processor executes the computer program, the method for video deraining based on the RWKV model according to any one of claims 1 to 7 is implemented.

10. A computer-readable storage medium, characterized in that include: A stored computer program, wherein when the computer program is executed, the device where the computer-readable storage medium is located is controlled to execute the video deraining method based on the RWKV model according to any one of claims 1 to 7.

Images