Image rain removal method based on learnable prior distillation and frequency domain cross attention
By adopting an image deraining method based on learnable prior distillation and frequency domain cross-attention, the problems of incomplete rain pattern removal and insufficient model generalization ability in single image deraining techniques are solved, and efficient image deraining effect is achieved in complex rain pattern scenes.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- JIANGSU HAOHAN INFORMATION TECH
- Filing Date
- 2026-02-28
- Publication Date
- 2026-05-29
Smart Images
Figure CN122115247A_ABST
Abstract
Description
Technical Field
[0001] This application relates to the field of image processing, and in particular to an image deraining method based on learnable prior distillation and frequency domain cross-attention. Background Technology
[0002] Single-image rain removal, a fundamental task in computer vision, plays a crucial supporting role in industrial scenarios such as autonomous driving, security monitoring, and remote sensing imaging. Its processing performance directly impacts the accuracy and reliability of subsequent core tasks like object detection and image recognition. Current mainstream technologies for addressing this problem primarily rely on deep learning methods, including feature modeling based on the Transformer architecture, traditional image restoration methods with manually designed priors, and conventional knowledge distillation-driven model optimization methods. However, existing methods suffer from several drawbacks. These include incomplete rain removal, limited model deployment, weak generalization ability, and incomplete feature representation. These drawbacks stem from issues such as reliance on manually designed priors, difficulty adapting to complex rain pattern distributions, high computational complexity of Transformer spatial self-attention, lack of semantic-level dynamic interaction mechanisms in knowledge distillation, insufficient separation of amplitude-phase information in frequency domain feature extraction, and unclear collaboration with spatial convolution.
[0003] Currently, single-image deraining techniques suffer from poor rain pattern removal and limited model usability. Summary of the Invention
[0004] This application provides an image deraining method based on learnable prior distillation and frequency domain cross-attention. It utilizes a prior representation network to perform degradation analysis on rainy images, extracting rain streak degradation patterns containing prior knowledge. Simultaneously, the rainy image is input into a pre-trained teacher network to acquire teacher knowledge. Using prior knowledge as the query and teacher knowledge as the key and value, a student model driven by a cross-attention mechanism performs prior distillation to generate prior features. In parallel, a frequency domain feature extraction module performs spatiotemporal feature fusion on the rainy image to obtain fused features. Combining the prior features and fused features, image reconstruction is completed, outputting a clear image. These techniques address the technical problems of poor rain streak removal and limited model practicality in existing single-image deraining techniques, achieving accurate rain streak separation and enhanced model generalization ability.
[0005] This application provides an image deraining method based on learnable prior distillation and frequency domain cross-attention, comprising: based on rainy image data, performing degradation analysis by combining the rainy image data with a prior representation network to extract rain streak degradation patterns, wherein the rain streak degradation patterns contain prior knowledge; inputting the rainy image data into a pre-trained teacher network to extract teacher knowledge; constructing a student model based on a cross-attention mechanism, using the prior knowledge as a query and the teacher knowledge as a key and value to perform prior distillation to generate prior features; introducing a frequency domain feature extraction module to perform spatiotemporal feature fusion on the rainy image data to generate fused features; and performing image reconstruction and restoration based on the prior features and the fused features to generate clear image data.
[0006] In a possible implementation, based on rainy image data, degradation analysis is performed using a prior representation network combined with the rainy image data to extract rain streak degradation patterns. These rain streak degradation patterns contain prior knowledge. The following processing is then performed: a lightweight network is constructed using multi-layer depthwise separable convolutions; rain streak learning is performed based on the rainy image data to generate rain streak structured parameters; adaptive training is performed based on the lightweight network and the rain streak structured parameters to construct a prior representation network; degradation priors are performed on the rainy image data using the prior representation network to generate prior knowledge; and the prior knowledge is added to the rain streak degradation patterns.
[0007] In a possible implementation, the rainy image data is input into a pre-trained teacher network to extract teacher knowledge, and the following processing is performed: an image deraining task is constructed based on the rainy image data; a Restormer network is trained according to the image deraining task, and the Restormer network is used as the teacher network; the rainy image data is input into the teacher network for cross-attention analysis to determine a first feature map; the rainy image data is input into the teacher network for gating analysis to determine a second feature map; the first feature map and the second feature map are dimensionally aligned to generate the teacher knowledge.
[0008] In a possible implementation, a student model is constructed based on a cross-attention mechanism. Prior knowledge is used as the query, and teacher knowledge is used as the key and value to perform prior distillation, generating prior features. The following processes are performed: the prior knowledge is projected into a query matrix through a linear transformation layer; the teacher knowledge is projected into a key matrix and a value matrix through a linear transformation layer; the similarity between the query matrix and the key matrix is calculated to generate matrix similarity; the value matrix is weighted and summed based on the matrix similarity to generate attention features; the attention features are added to the prior knowledge, and convolutional analysis is performed based on the sum to generate a gating graph; the attention features are element-wise weighted using the gating graph, and the teacher knowledge is distilled and transferred to the student model based on the weighting result to generate the prior features.
[0009] In a possible implementation, a frequency domain feature extraction module is introduced to perform spatiotemporal feature fusion on the rainy image data to generate fused features, and the following processing is performed: the frequency domain feature extraction module performs a discrete Fourier transform on the rainy image data to obtain complex frequency domain features; a set of complex convolution kernels and complex biases are used to perform complex multiplication and addition operations on the complex frequency domain features to generate modulation frequency domain features; an inverse discrete Fourier transform is performed based on the modulation frequency domain features to generate spatial domain features; features are extracted from the rainy image data based on the spatial domain features to generate local detail features; the spatial domain features and the local detail features are summed, and convolution fusion is performed based on the summation result to generate fused features.
[0010] In a possible implementation, image reconstruction and restoration are performed based on the prior features and the fused features to generate clear image data. The following processing is performed: the prior features and the fused features are combined to generate aggregate features; an image reconstruction network is constructed based on an encoder-decoder, and the image reconstruction network is optimized end-to-end using a loss function to generate an image reconstruction optimization network; the aggregate features are input into the image reconstruction optimization network, and the quality of the clear image data is constrained in multiple dimensions through the loss function to generate the clear image data.
[0011] In a possible implementation, the construction process of the loss function involves the following steps: acquiring real image data; mapping the clear image data and the real image data to pixel space for calculation to obtain a first L2 norm distance; mapping the clear image data and the real image data to the frequency domain for discrete Fourier transform to obtain a first L1 norm distance; mapping the rainy image data and the clear image data to the frequency domain for discrete Fourier transform to obtain a second L1 norm distance; normalizing the first L1 norm distance by dividing it by the second L1 norm distance to construct a frequency domain consistency loss parameter; calculating the distance based on the prior knowledge and the teacher's knowledge to obtain a second L2 norm distance; and integrating the first L2 norm distance, the second L2 norm distance, and the frequency domain consistency loss parameter to construct the loss function.
[0012] In a possible implementation, the following process is performed: the method is publicly trained and tested on Rain200L, Rain200H, and SPA-Dat to verify the robust removal capability of rain streaks of different densities and scenes.
[0013] The proposed image deraining method, based on learnable prior distillation and frequency domain cross-attention, firstly extracts rain streak degradation patterns from rainy image data through a prior representation network. These patterns incorporate prior knowledge. Next, the rainy image data is input into a pre-trained teacher network to extract teacher knowledge. Then, a student model is constructed based on a cross-attention mechanism, using the prior knowledge as a query and the teacher knowledge as a key and value to perform prior distillation, generating prior features. A frequency domain feature extraction module is then introduced to perform spatiotemporal feature fusion on the rainy image data, generating fused features. Finally, image reconstruction and restoration are performed based on the prior features and the fused features, generating clear image data. Through this process, the proposed method achieves accurate rain streak separation and enhances the model's generalization ability. Attached Figure Description
[0014] To more clearly illustrate the technical solutions of the embodiments of the present invention, the accompanying drawings of the embodiments of the present invention will be briefly described below. Flowcharts are used in this application to illustrate the operations performed by the system according to the embodiments of the present application. It should be understood that the preceding or following operations are not necessarily performed precisely in sequence. Instead, various steps can be processed in reverse order or simultaneously as needed. Furthermore, other operations can be added to these processes, or one or more steps can be removed from these processes.
[0015] Figure 1 This is a flowchart illustrating the image deraining method based on learnable prior distillation and frequency domain cross-attention provided in an embodiment of this application.
[0016] Figure 2 This is a schematic diagram illustrating the process of generating fused features in an image deraining method based on learnable prior distillation and frequency domain cross-attention, as provided in an embodiment of this application. Detailed Implementation
[0017] To further illustrate the technical means and effects of the present invention in achieving its intended purpose, the following detailed description of the specific implementation methods, structures, features, and effects of the present invention, in conjunction with the accompanying drawings and preferred embodiments, is provided below.
[0018] This application provides an image deraining method based on learnable prior distillation and frequency domain cross-attention, such as... Figure 1 As shown, the method includes: Step S100: Based on the rainy image data, degradation analysis is performed by combining the rainy image data with a prior representation network to extract rain streak degradation patterns, wherein the rain streak degradation patterns contain prior knowledge.
[0019] Specifically, a lightweight prior representation network is constructed, which is used to adaptively learn the pixel distribution and rain pattern texture features of rain images to identify the morphology (such as thin lines and dense dots), distribution density, direction and other degradation features of rain patterns. Prior knowledge containing these features is then extracted to form rain pattern degradation patterns.
[0020] In one possible implementation, based on rainy image data, a prior representation network is used to perform degradation analysis on the rainy image data to extract rain streak degradation patterns. These rain streak degradation patterns contain prior knowledge. Step S100 further includes step S110, which employs a multi-layer depthwise separable convolutional network to construct a lightweight network. Specifically, based on the network's lightweight requirements and the accuracy of rain streak feature extraction, the number of depthwise separable convolutional layers, the kernel parameters of each layer, and the variation rules of the number of channels are determined. Channel-by-channel convolution reduces computational load while preserving cross-channel feature interaction capabilities.
[0021] For example, a lightweight network can be constructed using three depthwise separable convolutional layers. The first layer is a depthwise convolutional layer with a 3×3 kernel, a stride of 1, and padding of 1. It has 3 input channels (corresponding to an RGB image) and 3 output channels. After each channel is independently convolved, a point convolutional layer (1×1 convolution) is used to increase the number of output channels to 32. The second layer has a 3×3 kernel, a stride of 1, and padding of 1. It has 32 input channels and 32 output channels, and the point convolutional layer increases the number of output channels to 64. The third layer has a 3×3 kernel, a stride of 1, and padding of 1. It has 64 input channels and 64 output channels, and the point convolutional layer adjusts the number of output channels to 32. Each depthwise convolutional layer is normalized using BatchNorm2d, and the point convolutional layer is activated using the ReLU activation function. The computational cost of the entire network is only 1 / 8 of that of a traditional convolutional network, making it efficient for mobile devices.
[0022] Step S120: Based on the rainy image data, rain pattern learning is performed to generate structured parameters for the rain patterns. Specifically, the rainy image data is preprocessed, such as normalization and cropping, and then input into a lightweight network. Through supervised learning, using clean images as labels, the network learns the difference features between rain patterns and the background, and then extracts the structured parameters of the rain patterns, including quantifiable parameters such as the length, width, grayscale range, and distribution probability of the rain patterns.
[0023] For example, rainy images from the Rain200L dataset are cropped into 256×256 sub-images, normalized, and their pixel values are divided by 255 before being input into the lightweight network constructed in step S110. Using the corresponding clean image as the label, the MSE loss between the network's output rain pattern feature map and the real rain pattern feature map is calculated, and the network is trained using the SGD optimizer. After training, the network can generate structured parameters for rain patterns, such as a width range of 1-3 pixels, a length range of 5-20 pixels, a grayscale value range of 0.3-0.7 (after normalization), and probability distribution matrices of rain patterns along the x and y axes in the image (256×256 dimensions, each element representing the probability of a rain pattern existing at the corresponding location).
[0024] Step S130: Based on the lightweight network and the rain ripple structure parameters, adaptive training is performed to construct a prior representation network. Specifically, the rain ripple structure parameters are incorporated as constraints into the training process of the lightweight network. By dynamically adjusting the network weights, the network can adapt to the structure parameters of different rain ripple patterns, forming a prior representation network specifically used to extract rain ripple degradation priors.
[0025] For example, based on the lightweight network from step S110, a constraint term for rain ripple structured parameters is added to the loss function. This constraint term calculates the L1 distance between the network-predicted rain ripple structured parameters and the true structured parameters. Adaptive training is performed using the Rain200H dataset. During training, if the network's predicted rain ripple width error is less than 0.5 pixels and its predicted length error is less than 2 pixels, the network is considered to have adapted to the current structured parameters. The final constructed prior representation network, upon inputting a rainy image, can output feature labels corresponding to the rain ripple structured parameters, used for prior knowledge generation.
[0026] Step S140: The prior representation network is used to perform degradation prior on the rainy image data to generate prior knowledge. Specifically, the preprocessed rainy image is input into the prior representation network. The network extracts high-order features related to rain stripe degradation in the image through multi-layer convolution and activation operations. These features are the prior knowledge containing information such as rain stripe distribution patterns and degradation degree. The output form is a feature map tensor.
[0027] For example, a rainy image with a resolution of 512×512×3 is normalized, with pixel values mapped to [-1,1], and input into the prior representation network constructed in step S130. The network first extracts low-level rain pattern texture features through a first-layer depthwise separable convolution, and then gradually generates higher-order features through subsequent convolutional layers and activation functions, finally outputting a feature map tensor with a dimension of 512×512×32 as prior knowledge. Each channel in this feature map corresponds to a rain pattern degradation feature, such as channel 1 corresponding to rain pattern density features (the larger the pixel value, the denser the rain pattern), channel 2 corresponding to rain pattern direction features (encoding different directions through pixel values), and channel 3 corresponding to rain pattern grayscale features, etc.
[0028] Step S150: Add the prior knowledge to the rain pattern degradation pattern. Specifically, feature splicing or weighted fusion is used to fuse the feature map corresponding to the prior knowledge with the basic features of the rain pattern degradation pattern, such as rain pattern morphology and distribution area, so that the rain pattern degradation pattern contains richer prior information, providing support for distillation and feature fusion.
[0029] For example, the basic feature map of the rain streak degradation pattern is a 512×512×16 feature map, containing information on the rain streak morphology and distribution area. The prior knowledge is a 512×512×32 feature map. A weighted fusion method is used, setting the weight of the basic feature map to 0.3 and the weight of the prior knowledge map to 0.7. The fusion is performed by multiplying element-wise and then adding the results, i.e., the pixel value of the fused feature map = pixel value of the basic feature map × 0.3 + pixel value of the prior knowledge map × 0.7. The resulting 512×512×32 rain streak degradation pattern feature map retains the basic rain streak distribution information while incorporating higher-order prior knowledge, thus improving the representational ability of the degradation pattern.
[0030] Step S200: Input the rainy image data into a pre-trained teacher network to extract teacher knowledge.
[0031] Specifically, a teacher network with strong image deraining capabilities is constructed and trained. This network can extract high-precision rain stripe removal-related features from rainy images, i.e., teacher knowledge, including background features and features that distinguish rain stripes from the background.
[0032] In one possible implementation, the rainy image data is input into a pre-trained teacher network to extract teacher knowledge. Step S200 further includes step S210, constructing an image deraining task based on the rainy image data, training a Restormer network according to the image deraining task, and using the Restormer network as the teacher network. Specifically, the input of the image deraining task is defined as a rainy image, the output as a clean image, and the evaluation metrics as PSNR and SSIM. The structural parameters of the Restormer network are configured, supervised training is performed using a rainy-clean image pairing dataset, and the network weights are optimized through backpropagation to enable the network to have high-quality image deraining capabilities. Finally, the trained network is used as the teacher network.
[0033] For example, the image deraining task takes an RGB image with rain as input (512×512 resolution) and outputs a clean RGB image. The Restormer network is configured as follows: the input layer maps the image to a 128-dimensional feature space, containing 8 stages. Each stage consists of 6 multi-head attention Transformer blocks and 2 Feed-Forward network blocks, with 16 attention heads. The hidden layer dimension of the Feed-Forward network is 512. The training set is composed of Rain200L, Rain200H, and SPA-Data, with a batch size of 4. The AdamW optimizer is used, with an initial learning rate of 0.0002, decaying every 100 epochs using a cosine annealing strategy. During training, the L1 loss function is used as the primary loss, combined with perceptual loss. After 500 epochs, the network achieves a PSNR of 40.99 and an SSIM of 0.9890 on the test set, meeting the performance requirements of the teacher network. Therefore, it is selected as the final teacher network.
[0034] Step S220: The rainy image data is input into the teacher network for cross-attention analysis to determine the first feature map. Specifically, the rainy image is input into the Transformer attention module of the teacher network. By calculating the correlation between features at different spatial locations in the image, the distribution of rain streaks and background structure features in the global range are captured to generate a first feature map containing global dependency information.
[0035] For example, a 512×512×3 rainy image is input into a pre-trained teacher network, and cross-attention analysis is performed in the multi-head attention module of the network's third stage. This module contains 16 attention heads, each with a feature dimension of 8 and a scaling factor of 8. For each location in the image feature map, its feature similarity with all other locations is calculated, and attention weights are obtained by normalization using the Softmax function. Then, the features of all locations are weighted and summed based on the weights. Each attention head outputs a 512×512×8 feature map. The outputs of the 16 attention heads are concatenated, and the number of channels is adjusted to 128 through a 1×1 convolution to obtain a 512×512×128 first feature map, which contains global rain texture and background correlation information of the image.
[0036] Step S230: The rainy image data is input into the teacher network for gating analysis to determine the second feature map. Specifically, a gating mechanism module is introduced into the teacher network. By learning adaptive weights, the feature maps of the rainy image are filtered and enhanced to highlight background features and suppress rain streak features, generating a second feature map that focuses on effective background information.
[0037] For example, a gated convolutional module is added after the fourth stage of the teacher network. This module contains two parallel 1×1 convolutional layers, used to generate gating weights and process features, respectively. The input is a first feature map (512×512×128). The first 1×1 convolutional layer outputs a 512×512×128 gated weight map, which maps the weight values to the [0,1] interval using a Sigmoid activation function. The second 1×1 convolutional layer performs a linear transformation on the input feature map, outputting a 512×512×128 feature map. The two outputs are multiplied element-wise to obtain a 512×512×128 second feature map, where regions with weight values close to 1 correspond to enhanced background features, and regions with weight values close to 0 correspond to suppressed rain pattern features.
[0038] Step S240: Align the first feature map with the second feature map in terms of dimensions to generate the teacher knowledge. Specifically, the first and second feature maps have the same number of channels and spatial resolution. A stitching and fusion strategy is used to integrate the two feature maps, and then the number of channels is adjusted through a convolutional layer. After processing by an activation function, teacher knowledge containing global correlation information and effective background information is formed.
[0039] For example, if both the first and second feature maps are 512×512×128 and their dimensions are aligned, a concatenation and fusion strategy is used to concatenate the two feature maps along the channel dimension, resulting in a 512×512×256 feature map. Then, a 3×3 convolutional layer is used to adjust the number of channels to 128, and after passing through the ReLU activation function, a 512×512×128 teacher knowledge feature map is obtained. This feature map contains both the global correlation information of the first feature map and the effective background information of the second feature map, providing high-quality teacher knowledge for subsequent prior distillation.
[0040] Step S300: Construct a student model based on the cross-attention mechanism, use the prior knowledge as the query, and use the teacher knowledge as the key and value to perform prior distillation to generate prior features.
[0041] Specifically, a lightweight student model is constructed, embedding a cross-attention mechanism module. Prior knowledge is mapped to query vectors, and teacher knowledge is mapped to key and value vectors. Attention weights are obtained by calculating the similarity between the query and the key. Then, the value vectors are weighted and summed based on the weights to realize the transfer of teacher knowledge to the student model and generate prior features that integrate prior knowledge and teacher knowledge.
[0042] In one possible implementation, a student model is constructed based on a cross-attention mechanism. Prior knowledge is used as the query, and teacher knowledge is used as the key and value to perform prior distillation, generating prior features. Step S300 further includes step S310, projecting the prior knowledge into a query matrix through a linear transformation layer. Specifically, a 1×1 convolutional linear transformation layer is designed to convert the feature map tensor of the prior knowledge into a query matrix format that meets the requirements of the cross-attention mechanism, ensuring that the feature dimensions match the key-value matrix.
[0043] For example, given a 512×512×32 feature map as prior knowledge, a 1×1 convolutional layer is used for the linear transformation layer, with 32 input channels, 64 output channels, 64 convolutional kernels, and a bias term initialized to 0. Through convolution, the 512×512×32 prior knowledge feature map is transformed into a 512×512×64 feature map, which is then converted into a (512×512)×64 two-dimensional matrix, serving as the query matrix for the cross-attention mechanism. Each row corresponds to a query vector representing a spatial location in the image.
[0044] Step S320: Project the teacher knowledge into a key matrix and a value matrix through a linear transformation layer. Specifically, two independent 1×1 convolutional linear transformation layers are set to perform dimensional transformation on the teacher knowledge, generating a key matrix and a value matrix respectively, ensuring that the key matrix has the same dimension as the query matrix and that the value matrix matches the spatial dimension of the key matrix.
[0045] For example, given a 512×512×128 feature map of teacher knowledge, the first linear transformation layer uses a 1×1 convolution with 128 input channels, 64 output channels, 64 kernels, and a bias term initialized to 0. This transforms the teacher knowledge into a 512×512×64 feature map, resulting in a (512×512)×64 key matrix. The second linear transformation layer uses a 1×1 convolution with 128 input channels, 64 output channels, 64 kernels, and a bias term initialized to 0. This transforms the teacher knowledge into a 512×512×64 feature map, resulting in a (512×512)×64 value matrix.
[0046] Step S330: Calculate the similarity between the query matrix and the key matrix to generate a matrix similarity. Specifically, a dot product attention calculation method is used. The transposes of the query matrix and the key matrix are multiplied, and then the result is scaled by the square root of the key matrix dimension to avoid excessively large values that could saturate the Softmax function. The final result is a matrix similarity representing the degree of association between the query and each key, where a larger element value indicates a higher degree of association between the corresponding query and the key.
[0047] Step S340: The value matrix is weighted and summed based on the matrix similarity to generate attention features. Specifically, the matrix similarity is Softmax normalized to obtain an attention weight matrix. Then, the attention weight matrix is multiplied by the value matrix to achieve weighted fusion of value features based on the degree of association, thereby generating attention features.
[0048] For example, if the matrix similarity dimension is (262144, 262144), performing Softmax normalization on each row yields the attention weight matrix, where each row represents the attention weight assigned to all key positions by the corresponding query position. The value matrix has a dimension of (262144, 64). Calculating the matrix multiplication between the attention weight matrix and the value matrix results in an attention feature matrix of dimension (262144, 64), which is then converted into a 512×512×64 feature map, i.e., the attention features.
[0049] Step S350: The attention features are added to the prior knowledge, and convolutional analysis is performed based on the addition result to generate a gated graph. Specifically, the attention features and prior knowledge are first dimensionally aligned, then element-wise addition and fusion are performed. The fused features are input into a convolutional layer to extract high-order features, and finally a gated graph is generated using a sigmoid activation function for subsequent adaptive selection of attention features.
[0050] For example, if the attention feature is 512×512×64 and the prior knowledge is 512×512×32, the number of channels in the prior knowledge is first adjusted to 64 using a 1×1 convolutional layer, and then added element-wise with the attention feature to obtain a 512×512×64 fused feature. This fused feature is then input into a 3×3 convolutional layer to extract higher-order fused features. Finally, a sigmoid activation function maps the feature values to the [0,1] interval, generating a 512×512×64 gating graph. Each element in the gating graph represents the retention weight of the feature at the corresponding location.
[0051] Step S360: The attention features are weighted element-wise using the gating graph. Based on the weighting result, the teacher knowledge is distilled and transferred to the student model to generate the prior features. Specifically, the gating graph and attention features are multiplied element-wise to filter and enhance effective features and suppress redundant features. The weighted features are the core features for the teacher knowledge distillation and transfer to the student model. The feature representation is then optimized through convolutional blocks in the student model to finally generate the prior features.
[0052] For example, if the gating graph is 512×512×64 and the attention feature is also 512×512×64, multiplying them element-wise yields a weighted feature of 512×512×64, enhancing effective features and suppressing redundant features. This weighted feature is then input into a two-layer convolutional block of the student model. The first layer is a 3×3 convolution with a stride of 1 and padding of 1, producing 64 output channels; the second layer is a 1×1 convolution with 32 output channels. After ReLU activation and BatchNorm2d normalization, a 512×512×32 prior feature is generated. This feature integrates the rain ripple degradation pattern of prior knowledge and the high-precision features of teacher knowledge, completing the distillation transfer of teacher knowledge to the student model.
[0053] Step S400: Introduce the frequency domain feature extraction module to perform spatiotemporal feature fusion on the rainy image data to generate fused features.
[0054] Specifically, the frequency domain feature extraction module first converts the spatial features of the rainy image to the frequency domain, modulates the frequency domain features through complex convolution, and then converts them back to the spatial domain. At the same time, it extracts local spatial detail features through a spatial feature extraction network composed of convolutions. Finally, it fuses the frequency domain processed spatial features with the local spatial detail features to generate a fused feature that combines global frequency domain information and local spatial details.
[0055] like Figure 2As shown, in one possible implementation, a frequency domain feature extraction module is introduced to perform spatiotemporal feature fusion on the rainy image data to generate fused features. Step S400 further includes step S410, where the frequency domain feature extraction module performs a discrete Fourier transform on the rainy image data to obtain complex frequency domain features. Specifically, a 2D discrete Fourier transform is performed on the spatial features of each channel of the rainy image to convert the spatial distribution information in the spatial domain into the frequency distribution information in the frequency domain, resulting in a complex frequency domain feature tensor containing real and imaginary parts.
[0056] For example, the spatial domain features of rainy image data are 512×512×3. A 2D Discrete Fourier Transform is performed on the 512×512 pixel matrix of each channel. Through calculation, each channel obtains a 512×512 complex matrix, and the entire image obtains a 512×512×3 complex frequency domain feature, where each element contains real and imaginary part information.
[0057] Step S420 involves using a set of complex convolution kernels and complex biases to perform complex multiplication and addition operations on the complex frequency domain features to generate modulation frequency domain features. Specifically, complex convolution kernels and complex bias terms are designed to match the number of frequency domain feature channels. The convolution kernel for each channel is a complex number. Complex multiplication and addition are performed on the complex frequency domain features of each channel, i.e., the convolution kernel is multiplied by the frequency domain features and the bias term is added to achieve the separation of rain ripple frequency components from background frequency components.
[0058] For example, the complex frequency domain feature is 512×512×3. The complex convolution kernel dimension is set to 3×1×1, with each channel having a complex kernel and three complex bias terms. Complex multiplication and addition operations are performed on all channels and all frequency domain coordinates to generate a 512×512×3 modulation frequency domain feature.
[0059] Step S430: Perform an inverse discrete Fourier transform based on the modulation frequency domain features to generate spatial domain features. Specifically, perform a 2D inverse discrete Fourier transform on the modulation frequency domain features of each channel to convert the modulated frequency distribution information in the frequency domain back to the pixel distribution information in the spatial domain, obtaining spatial domain features that have undergone global frequency domain processing. These features have separated some rain ripple frequency components.
[0060] For example, if the modulation frequency domain feature is 512×512×3, perform a 2D inverse discrete Fourier transform on the 512×512 complex matrix of each channel. Through calculation, each channel obtains a 512×512 spatial pixel matrix, and the entire image obtains a 512×512×3 spatial domain feature, which is the spatial domain feature after global frequency domain processing.
[0061] Step S440: Based on the spatial features, feature extraction is performed on the rainy image data to generate local detail features. Specifically, a spatial feature extraction network is constructed, and convolutional layers are used to perform convolution operations on the original spatial features of the rainy image to extract local texture, edge, and other detail features of the image. These features can supplement the local information of the spatial features after frequency domain processing.
[0062] For example, the spatial feature extraction network consists of three convolutional layers, with the input being the original spatial features of a rainy image (512×512×3). The first convolutional layer has a kernel size of 3×3, a stride of 1, padding of 1, and 16 output channels, using the ReLU activation function. The second convolutional layer has a kernel size of 3×3, a stride of 1, padding of 1, and 32 output channels, also using the ReLU activation function. The third convolutional layer has a kernel size of 1×1, a stride of 1, no padding, and 3 output channels. Through the convolutional operations of this network, detailed features such as local edges (e.g., object contour edges) and textures (e.g., ground texture) of the rainy image are extracted, generating 512×512×3 local detail features.
[0063] Step S450: The spatial features and the local detail features are summed, and convolutional fusion is performed based on the summation result to generate fused features. Specifically, the spatial features and local detail features have the same dimension, and element-wise summation and fusion are performed. The summed features are then input into a 1×1 convolutional layer for feature integration, and finally added to the original rainy image features to generate fused features that combine global frequency domain information and local spatial details.
[0064] For example, the spatial features after frequency domain processing are 512×512×3, and the local detail features are also 512×512×3. The two are summed element-wise to obtain a 512×512×3 summed feature. This summed feature is then input into a 1×1 convolutional layer (3 output channels) for convolutional fusion. The integrated feature is extracted and then added element-wise to the original rainy image features (512×512×3) to generate a 512×512×3 fused feature. This feature contains both the global rain streak suppression information from frequency domain processing and retains the local spatial detail information.
[0065] Step S500: Based on the prior features and the fused features, perform image reconstruction and restoration to generate clear image data.
[0066] Specifically, prior features and fused features are fused and input into an image reconstruction network based on an encoder-decoder structure. The network is optimized through end-to-end training, and a multi-dimensional loss function is used to constrain the difference between the reconstructed image and the real clean image, ultimately outputting clear image data.
[0067] In one possible implementation, image reconstruction and restoration are performed based on the prior features and the fused features to generate clear image data. Step S500 further includes step S510, which combines the prior features and the fused features to generate aggregated features. Specifically, since the prior features and the fused features have different number of channels, the number of channels in the fused features is adjusted to match that of the prior features using a 1×1 convolutional layer, and then combined using channel concatenation to integrate the advantages of both features and generate aggregated features.
[0068] For example, the prior feature is 512×512×32, and the fused feature is 512×512×3. A combination of channel adjustment and concatenation is used. First, a 1×1 convolutional layer is used to adjust the number of channels in the fused feature to 32, aligning it with the number of channels in the prior feature. Then, the adjusted fused feature is concatenated with the prior feature along the channel dimension, resulting in a 512×512×64 feature map. This feature map is the aggregated feature, which includes both the prior knowledge of rain ripple degradation from the prior feature and the global frequency domain information and local spatial details from the fused feature.
[0069] Step S520: Construct an image reconstruction network based on an encoder-decoder architecture, and perform end-to-end optimization of the image reconstruction network using a loss function to generate an optimized image reconstruction network. Specifically, an encoder-decoder structure image reconstruction network is constructed, where the encoder is responsible for extracting high-order features of the aggregated features, and the decoder is responsible for mapping the high-order features back to the image pixel space. A multi-dimensional loss function is designed, and the loss function is minimized through backpropagation to optimize the network weights, generating an optimized image reconstruction network with high-quality image reconstruction capabilities.
[0070] For example, the encoder of an image reconstruction network contains four convolutional layers: the first is a 3×3 convolution with 64 input channels and 128 output channels; the second is a 3×3 convolution with 256 output channels; the third is a 3×3 convolution with 512 output channels; and the fourth is a 3×3 convolution with 1024 output channels. Each convolutional layer is followed by a Leaky ReLU activation function. The decoder contains four deconvolutional layers: the first is a 3×3 deconvolution with 1024 input channels and 512 output channels; the second is a 3×3 deconvolution with 256 output channels; the third is a 3×3 deconvolution with 128 output channels; and the fourth is a 3×3 deconvolution with 3 output channels. Each deconvolutional layer is followed by a ReLU activation function. The encoder and decoder fuse features through skip connections.
[0071] Step S530: The aggregated features are input into the image reconstruction optimization network. The loss function is used to impose multi-dimensional constraints on the quality of the clear image data, generating the clear image data. Specifically, the aggregated features are input into the image reconstruction optimization network. The decoder maps the high-order features to the image pixel space output through deconvolution and skip connections. The loss function constrains the quality of the candidate clear image from three dimensions: pixel space, frequency domain space, and prior knowledge matching degree, ensuring consistency with the real clean image in multiple dimensions. Finally, clear image data that meets the constraints is output. For example, a 512×512×64 aggregated feature map is input into the image reconstruction optimization network. After the encoder extracts high-order features, the decoder gradually restores the feature map resolution to 512×512 through deconvolution, outputting a 3-channel candidate sharp image. The normalized values of the pixel L2 distance and frequency domain L1 distance between the candidate sharp image and the real clean image, as well as the L2 distance between prior knowledge and teacher knowledge, are calculated using a loss function. The network output is adjusted based on the loss values, ultimately generating 512×512×3 sharp image data. This data is highly consistent with the real clean image in dimensions such as pixels and frequency domain, achieving a PSNR of 41.88 and an SSIM of 0.9908 on the Rain200L test set.
[0072] In one possible implementation, the loss function construction process, step S520, further includes step S521: acquiring real image data, mapping the clear image data and the real image data to pixel space for calculation, and obtaining the first L2 norm distance. Specifically, acquiring real clean image data corresponding to rainy images with different scenes and different rain streak densities, such as clean images from the Rain200L, Rain200H, and SPA-Data datasets, converting the clear image data output by the image reconstruction network and the real clean image data into pixel value vectors of the same dimension, and calculating the L2 norm distance between the two, i.e., the Euclidean distance, which reflects the degree of difference in pixel space.
[0073] For example, 100 pairs of rainy-clean images are collected from the Rain200L dataset. The real clean image data is 512×512×3, and the clear image data is the 512×512×3 output by the network. The pixel values of each clear image are expanded row-wise into a vector of 1×(512×512×3)=1×786432, and the pixel values of the real clean image are expanded into a vector of 1×786432. The first L2 norm distance is calculated based on the two vectors.
[0074] Step S522: Map the clear image data and the real image data to the frequency domain and perform a Discrete Fourier Transform to obtain the first L1 norm distance. Specifically, perform a 2D Discrete Fourier Transform on the clear image data and the real clean image data respectively, mapping them from the pixel space to the frequency domain to obtain the corresponding complex frequency domain features. Calculate the L1 norm distance between the two frequency domain features, which is the sum of their absolute values. This distance reflects the degree of difference in the frequency domain distribution.
[0075] For example, the sharp image data is 512×512×3, and the real clean image data is also 512×512×3. A 2D Discrete Fourier Transform is performed on each channel of both images to obtain their respective complex frequency domain features (512×512×3). For the frequency domain feature of each channel, the complex difference at the corresponding frequency domain coordinates is calculated, i.e., the frequency domain value of the sharp image minus the frequency domain value of the real image. The absolute value of this complex difference is taken, and then the magnitudes of all channels and all frequency domain coordinates are summed to obtain the first L1 norm distance.
[0076] Step S523: Map the rainy image data and the clear image data to the frequency domain and perform a Discrete Fourier Transform to obtain the second L1 norm distance. Specifically, perform a 2D Discrete Fourier Transform on the rainy image data and the clear image data respectively, map them to the frequency domain, obtain the corresponding complex frequency domain features, and calculate the L1 norm distance between the two frequency domain features. This distance reflects the difference between the rainy image and the clear image in the frequency domain, that is, the difference in the frequency domain components corresponding to the rain streaks.
[0077] For example, the rainy image data is 512×512×3, and the clear image data is also 512×512×3. A 2D Discrete Fourier Transform is performed on each channel of both images to obtain the complex frequency domain features (512×512×3). For each channel's frequency domain feature, the complex difference at the corresponding frequency domain coordinates is calculated, i.e., the frequency domain value of the rainy image minus the frequency domain value of the clear image. The modulus of this complex difference is then taken, and the modulus values of all channels and all frequency domain coordinates are summed to obtain the second L1 norm distance.
[0078] Step S524: Normalize the first L1 norm distance by dividing the second L1 norm distance to construct the frequency domain consistency loss parameter. Specifically, divide the first L1 norm distance by the second L1 norm distance, adding a small constant to the denominator to avoid it being zero, to obtain the normalized ratio. This ratio is the frequency domain consistency loss parameter, used to constrain the consistency of the frequency domain distribution between the sharp image and the real image. The closer this parameter is to 0, the more consistent the frequency domain distribution between the sharp image and the real image.
[0079] Step S525: Calculate the distance between the prior knowledge and the teacher knowledge to obtain the second L2 norm distance. Specifically, convert the prior knowledge feature map and the teacher knowledge feature map into vectors of the same dimension, and calculate their L2 norm distance. This distance reflects the similarity between the prior knowledge and the teacher knowledge and is used to constrain the learning direction of the prior representation network.
[0080] For example, the prior knowledge is a feature map of 512×512×32, and the teacher's knowledge is a feature map of 512×512×128. First, a 1×1 convolutional layer is used to adjust the number of channels in the teacher's knowledge to 32, consistent with the prior knowledge. Then, both are expanded row-wise into vectors of 1×(512×512×32)=1×8388608. The second L2 norm distance is calculated based on the two vectors. The smaller this distance, the more similar the prior knowledge and the teacher's knowledge are, and the more correct the learning direction of the prior representation network.
[0081] Step S526: Integrate the first L2 norm distance, the second L2 norm distance, and the frequency domain consistency loss parameter to construct the loss function. Specifically, the first L2 norm distance, the second L2 norm distance, and the frequency domain consistency loss parameter are integrated by weighted summation to construct a multi-dimensional loss function, which is used to comprehensively constrain the training of the image reconstruction network. The weight coefficients are determined through experimental verification.
[0082] In one possible implementation, the method further includes step S600, wherein the method is trained and tested on Rain200L, Rain200H, and SPA-Dat to verify the robust removal capability of rain streaks of different densities and scenes.
[0083] Specifically, publicly available deraining datasets containing different rain streak densities and different scenes were selected. Rain200L is for low density, Rain200H is for high density, and SPA-Data covers various scenes, including outdoor streets, forests, buildings, and skies. Each dataset was divided into training and test sets in an 8:2 ratio. The image deraining model was trained using the mixed training set, and the model's deraining performance under different rain streak densities and scenes was evaluated using the test set to verify the model's robustness.
[0084] For example, the Rain200L dataset contains 200 pairs of rain-clean images, Rain200H contains 200 pairs, and SPA-Data contains 500 pairs. Each dataset is divided into training and testing sets in an 8:2 ratio. A mixed training set (merging the training sets of the three datasets) is used to train the model with a total of 720 image pairs. After training, the model is tested on the testing sets of the three datasets, and the PSNR and SSIM of each image after rain removal are calculated. The test results show that the PSNR reaches 41.88 and the SSIM reaches 0.9908 on the Rain200L test set; the PSNR reaches 32.51 and the SSIM reaches 0.9361 on the Rain200H test set; and the PSNR reaches 49.29 and the SSIM reaches 0.9936 on the SPA-Data test set. All of these results outperform existing mainstream methods such as Uformer, Restormer, and DRSformer, validating the model's robust ability to remove rain streaks of different densities and scenes.
[0085] This application utilizes a priori representation network to perform degradation analysis on rainy images, extracting rain streak degradation patterns containing prior knowledge. Simultaneously, the rainy image is input into a pre-trained teacher network to obtain teacher knowledge. Using prior knowledge as the query and teacher knowledge as the key and value, a student model driven by a cross-attention mechanism performs prior distillation to generate prior features. In parallel, a frequency domain feature extraction module is used to perform spatiotemporal feature fusion on the rainy image to obtain fused features. Combining the prior features and fused features, image reconstruction is completed, outputting a clear image. These techniques solve the technical problems of poor rain streak removal effect and limited model practicality in existing single-image rain removal techniques, achieving the technical effects of accurate rain streak separation and enhanced model generalization ability.
[0086] The above description is merely a preferred embodiment of the present invention and is not intended to limit the present invention in any way. Although the present invention has been disclosed above with reference to preferred embodiments, it is not intended to limit the present invention. Any person skilled in the art can make some modifications or alterations to the above-disclosed technical content to create equivalent embodiments without departing from the scope of the present invention. Any modifications, equivalent changes, and alterations made to the above embodiments based on the technical essence of the present invention without departing from the scope of the present invention shall still fall within the scope of the present invention.
Claims
1. An image rain removal method based on learnable prior distillation and frequency domain cross-attention, characterized in that, The method includes: Based on rainy image data, degradation analysis is performed by combining the rainy image data with a prior representation network to extract rain streak degradation patterns, which contain prior knowledge. The rainy image data is input into a pre-trained teacher network to extract teacher knowledge; A student model is constructed based on the cross-attention mechanism. Prior knowledge is used as the query and teacher knowledge is used as the key and value to perform prior distillation and generate prior features. A frequency domain feature extraction module is introduced to perform spatiotemporal feature fusion on the rainy image data to generate fused features; Image reconstruction and restoration are performed based on the prior features and the fused features to generate clear image data.
2. The image deraining method based on learnable prior distillation and frequency domain cross-attention as described in claim 1, characterized in that, The method involves using rainy image data as the basis, and performing degradation analysis through a prior representation network to extract rain streak degradation patterns. These rain streak degradation patterns incorporate prior knowledge. The method includes: A lightweight network is constructed using multiple depthwise separable convolutions. Rain pattern learning is performed based on the rainy image data to generate rain pattern structure parameters; Based on the lightweight network and the rain ripple structure parameters, an adaptive training is performed to construct a prior representation network; The prior representation network is used to perform degradation priors on rainy image data to generate prior knowledge. The prior knowledge is added to the rain pattern degradation mode.
3. The image deraining method based on learnable prior distillation and frequency domain cross-attention as described in claim 1, characterized in that, The rainy image data is input into a pre-trained teacher network to extract teacher knowledge. The method includes: Based on the rainy image data, an image deraining task is constructed, and a Restormer network is trained according to the image deraining task. The Restormer network is used as the teacher network. The rainy image data is input into the teacher network for cross-attention analysis to determine the first feature map; The rainy image data is input into the teacher network for gating analysis to determine the second feature map; The first feature map and the second feature map are dimensionally aligned to generate the teacher knowledge.
4. The image deraining method based on learnable prior distillation and frequency domain cross-attention as described in claim 1, characterized in that, A student model is constructed based on a cross-attention mechanism. Prior knowledge is used as the query, and teacher knowledge is used as the key and value to perform prior distillation to generate prior features. The method includes: The prior knowledge is projected into a query matrix through a linear transformation layer; The teacher's knowledge is projected into a key matrix and a value matrix through a linear transformation layer, respectively; Calculate the similarity between the query matrix and the key matrix to generate matrix similarity. Based on the matrix similarity, the value matrix is weighted and summed to generate attention features; The attention features are added to the prior knowledge, and convolutional analysis is performed based on the addition result to generate a gated graph; The attention features are weighted element-wise using the gating graph, and the teacher knowledge is distilled and transferred to the student model based on the weighting result to generate the prior features.
5. The image deraining method based on learnable prior distillation and frequency domain cross-attention as described in claim 1, characterized in that, A frequency domain feature extraction module is introduced to perform spatiotemporal feature fusion on the rainy image data to generate fused features. The method includes: A frequency domain feature extraction module is introduced to perform a discrete Fourier transform on the rainy image data to obtain complex frequency domain features; A set of complex convolution kernels and complex biases are used to perform complex multiplication and addition operations on the complex frequency domain features to generate modulation frequency domain features; Based on the modulation frequency domain features, perform an inverse discrete Fourier transform to generate spatial domain features; Based on the aforementioned spatial features, feature extraction is performed on rainy image data to generate local detail features; The spatial features and the local detail features are summed, and the summation result is used for convolution fusion to generate fused features.
6. The image deraining method based on learnable prior distillation and frequency domain cross-attention as described in claim 1, characterized in that, Image reconstruction and restoration based on the prior features and the fused features to generate clear image data includes the following methods: The prior features are combined with the fused features to generate aggregated features; An image reconstruction network is constructed based on an encoder-decoder, and the image reconstruction network is optimized end-to-end using a loss function to generate an optimized image reconstruction network. The aggregated features are input into the image reconstruction optimization network, and the loss function is used to impose multidimensional constraints on the quality of the clear image data to generate the clear image data.
7. The image deraining method based on learnable prior distillation and frequency domain cross-attention as described in claim 6, characterized in that, The process of constructing the loss function includes the following methods: Collect real image data, map the clear image data and the real image data to the pixel space for calculation, and obtain the first L2 norm distance; The clear image data and the real image data are mapped to the frequency domain and then subjected to discrete Fourier transform to obtain the first L1 norm distance; The rainy image data and the clear image data are mapped to the frequency domain and subjected to discrete Fourier transform to obtain the second L1 norm distance; Normalize by dividing the first L1 norm distance by the second L1 norm distance to construct the frequency domain consistency loss parameter; The distance between the prior knowledge and the teacher's knowledge is calculated to obtain the second L2 norm distance. The first L2 norm distance, the second L2 norm distance, and the frequency domain consistency loss parameter are integrated to construct the loss function.
8. The image deraining method based on learnable prior distillation and frequency domain cross-attention as described in claim 1, characterized in that, The method was publicly trained and tested on Rain200L, Rain200H, and SPA-Dat to verify its robust removal capability for rain streaks of different densities and scenes.