Image rain removal method of space state sequence model based on learnable function

Abstract

The invention discloses an image rain removal method of a space state sequence model based on a learnable function, an image to be subjected to rain removal is input into a constructed rain removal network with a space state sequence module based on the learnable function to obtain an image after rain removal, and the method comprises the following steps: S1, constructing a space state sequence model facing a 2D image; s2, constructing a KAN layer based on a learnable function; s3, fusing a space state sequence module SKLayer based on a learnable function; s4, constructing a rain removal network structure; S5, inputting the data into the rain removal network structure for training iteration, and completing network training; and S6, after training iteration, inputting an image to be subjected to rain removal into the trained rain removal network structure to obtain a rain-removed image. According to the invention, a space state sequence module based on a learnable function is introduced, so that a global receptive field and linear calculation complexity are achieved, and a relatively good rain removal effect is achieved.

Description

Technical Field

[0001] The present invention relates to the technical field of image processing, and in particular to an image deraining method based on a spatial state sequence model of a learnable function. Background Art

[0002] Convolutional neural networks (CNNs) and vision-based transformers (ViTs) have achieved remarkable success in the field of computer vision. However, CNNs are usually not as effective as ViTs due to their limited receptive field. This is because ViTs have a global receptive field based on an attention mechanism that can better extract global features. However, since the computational complexity of the attention mechanism is quadratic with the image size, the computational overhead of using ViT in downstream computer vision tasks such as image deraining is very high. The present invention proposes a spatial state sequence model based on a learnable function to solve the problem of image deraining. The method has a global receptive field and linear computational complexity. Summary of the invention

[0003] The technical problem to be solved by the present invention is to provide an image deraining method based on a spatial state sequence model of a learnable function, introduce a spatial state sequence module based on a learnable function, so as to have a global receptive field and linear computational complexity, and achieve a better deraining effect.

[0004] In order to solve the above technical problems, an embodiment of the present invention provides an image deraining method based on a spatial state sequence model of a learnable function, comprising the following steps:

[0005] S1. Construct a spatial state sequence model ISSM for 2D images;

[0006] S2, construct a KAN layer based on learnable functions;

[0007] S3, integrating step S1 and step S2 to form a spatial state sequence module SKLayer based on a learnable function;

[0008] S4. Construct a deraining network structure based on the spatial state sequence module SKLayer of learnable functions.

[0009] Further, step S1 includes the following steps:

[0010] S1.1, divide the 2D image into M×N non-overlapping sub-image blocks, and scan the image horizontally and vertically to obtain image embedding sequences X0 and X2, and reverse X0 and X2 to obtain X1 and X3;

[0011] S1.2, define the space state model SSM, the formula is as follows:

[0012] h t=Ah t-1 +Bx t

[0013] y t =Ch t +Dx t ;

[0014] Among them, matrices A, B, C, and D are spatial parameters, and x t is the input sequence, h t-1 is the state parameter at time t-1, h t is the state parameter at time t, y t is the output sequence;

[0015] S1.3, the sequences X0, X1, X2, X3 obtained in step S1.1 are used as the input sequences of step S1.2 respectively, and the output sequences y0, y1, y2, y3 are obtained, and the 2D images Z0, Z1, Z2, Z3 are restored in the horizontal and vertical scanning order in step S1.1, and the images Z0, Z1, Z2, Z3 are combined to output Z:

[0016] Z=Z0+Z1+Z2+Z3.

[0017] Further, step S2 includes the following steps:

[0018] According to the Kolmogorov-Arnold theory, a multivariable continuous function f can be expressed as a combination of a finite number of single-variable continuous functions. The input is Z, and the output can be expressed as:

[0019]

[0020] Among them, Φ represents a learnable univariate function, represents matrix multiplication, n represents the number of functions, and KANLayer(U) represents the output.

[0021] Further, in step S3, the input is defined as input, including the following steps:

[0022] S3.1, use layer normalization to process the input input, obtain the output LM(input), and concatenate the output of the layer normalization as the input of the KAN layer based on the learnable function in step S2, and output KANLayer(LM(input));

[0023] S3.2, after splicing, it is used as the input of the space state sequence model ISSM in step S1, and the ISSM is obtained after being used n times continuously. n(KANLayer(LM(input))); The output is concatenated again as the input of the KAN layer based on the learnable function in step S2, and the output is KANLayer(ISSM n (KANLayer(LM(input))));

[0024] S3.3, the spliced ​​output and input are output together to obtain KANLayer (ISSM n (KANLayer(LM(input))))+input.

[0025] Further, in step S4, the deraining network structure consists of a 3*3 convolutional layer for pre-processing, several SKLayers for extracting deep features, and a 3*3 convolutional layer for post-processing, including the following steps:

[0026] S3.1, use layer normalization to process the input input, obtain the output LM(input), and concatenate the output of the layer normalization as the input of the KAN layer based on the learnable function in step S2, and output KANLayer(LM(input));

[0027] S3.2, after splicing, it is used as the input of the space state sequence model ISSM in step S1, and the ISSM is obtained after being used n times continuously. n (KANLayer(LM(input))); The output is concatenated again as the input of the KAN layer based on the learnable function in S2, and the output is KANLayer(ISSM n (KANLayer(LM(input))));

[0028] S3.3, the spliced ​​output and input are output together to obtain KANLayer (ISSM n (KANLayer(LM(input))))+input.

[0029] The deraining network structure consists of a 3*3 convolutional layer for pre-processing, several SKLayers for extracting deep features, and a 3*3 convolutional layer for post-processing, including the following steps:

[0030] S4.1, the input image is processed using a 3×3 convolutional layer to obtain conv(img);

[0031] S4.2. Process conv(img) obtained in step S4.1 through m SKLayer processes, and outputs SKLayer m (conv(img));

[0032] S4.3. Use a 3×3 convolutional layer to process the SKLayer obtained in step S4.2 m (conv(img)), get conv(SKLayer m (conv(img)));

[0033] S4.4. Output the restored image img+conv(SKLayer m (conv(img))).

[0034] Among them, step S4 also includes establishing a reconstruction loss function between the output restored image and its corresponding true value image, and using the loss function to constrain the deraining network structure, wherein the loss function is defined as:

[0035]

[0036] in, represents the pixel value of channel i of the true image, represents the derained image obtained after network processing, express The pixel value of channel i, N is the number of image pairs, F s is a smooth function, defined as:

[0037]

[0038] Among them, in step S5, using the gradient descent method to complete the training of the network includes the following steps:

[0039] S5.1, set the learning rate to 0.0001 and the batch size to 8;

[0040] S5.2. Adopt the adam algorithm as the optimization algorithm. After 200 iterations, calculate the SSIM and PSNR on the test set every 20 rounds. After 500 iterations, take the result with the highest SSIM and PSNR as the final training result.

[0041] The beneficial effects of the above technical solution of the present invention are as follows: The present invention proposes an image deraining method based on a spatial state sequence model of a learnable function, designs a spatial state sequence module based on a learnable function, and is used to extract global and local features of the image so as to have a global receptive field and linear computational complexity, thereby achieving a better deraining effect. BRIEF DESCRIPTION OF THE DRAWINGS

[0042] Figure 1 It is a schematic diagram of the scanning sequence in the space state sequence model ISSM;

[0043] Figure 2A schematic diagram of the overall network structure. DETAILED DESCRIPTION

[0044] In order to make the technical problems, technical solutions and advantages to be solved by the present invention more clear, a detailed description will be given below with reference to the accompanying drawings and specific embodiments.

[0045] An embodiment of the present invention provides an image deraining method based on a spatial state sequence model of a learnable function, comprising the following steps:

[0046] S1. Constructing a spatial state sequence model ISSM for 2D images, including the following steps:

[0047] S1.1, divide the 2D image into M×N non-overlapping sub-image blocks, and perform horizontal and vertical scanning on the image (the specific scanning sequence is as follows Figure 1 As shown), we get the image embedding sequences X0 and X2, and reverse X0 and X2 to get X1 and X3 (where 0, 1, 2, and 3 represent four different scanning directions);

[0048] S1.2, define the space state model SSM, the formula is as follows:

[0049] h t =Ah t-1 +Bx t

[0050] y t =Ch t +Dx t ;

[0051] Among them, matrices A, B, C, and D are spatial parameters that can be aware of the input context. t is the input sequence, h t-1 is the state parameter at time t-1, h t is the state parameter at time t, y t is the output sequence;

[0052] S1.3, the sequences X0, X1, X2, and X3 obtained in step S1.1 are used as the input sequences of step S1.2 respectively, and the output sequences y0, y1, y2, and y3 are obtained, and four independent sequences are obtained. The horizontal and vertical scanning strategies ensure that each element in the feature map integrates information from all other positions in different directions, thereby creating a global receptive field without increasing the linear computational complexity, and the horizontal and vertical scanning order in step S1.1 is used to restore the 2D images Z0, Z1, Z2, and Z3. The images Z0, Z1, Z2, and Z3 are combined to output Z:

[0053] Z=Z0+Z1+Z2+Z3.

[0054] S2. Construct a KAN layer based on a learnable function, including the following steps:

[0055] According to the Kolmogorov-Arnold theory, a multivariable continuous function f can be expressed as a combination of a finite number of single-variable continuous functions. The input is Z, and the output can be expressed as:

[0056]

[0057] Among them, Φ represents a learnable univariate function, represents matrix multiplication, n represents the number of functions, and KANLayer(U) represents the output.

[0058] S3, integrating step S1 and step S2 to form a spatial state sequence module SKLayer based on a learnable function, including the following steps:

[0059] S3.1, use layer normalization to process the input input, obtain the output LM(input), and concatenate the output of the layer normalization as the input of the KAN layer based on the learnable function in step S2, and output KANLayer(LM(input));

[0060] S3.2, after splicing, it is used as the input of the space state sequence model ISSM in step S1, and the ISSM is obtained after being used n times continuously. n (KANLayer(LM(input))); The output is concatenated again as the input of the KAN layer based on the learnable function in step S2, and the output is KANLayer(ISSM n (KANLayer(LM(input))));

[0061] S3.3, the spliced ​​output and input are output together to obtain KANLayer (ISSM n (KANLayer(LM(input))))+input.

[0062] S4, based on the spatial state sequence module SKLayer of the learnable function, build a rain removal network structure, such as Figure 2 As shown in the figure, the deraining network structure consists of a 3*3 convolutional layer for pre-processing, several SKLayers for extracting deep features, and a 3*3 convolutional layer for post-processing, including the following steps:

[0063] S4.1, the input image is processed using a 3×3 convolutional layer to obtain conv(img);

[0064] S4.2. Process conv(img) obtained in step S4.1 through m SKLayer processes, and outputs SKLayer m (conv(img));

[0065] S4.3. Use a 3×3 convolutional layer to process the SKLayer obtained in step S4.2 m (conv(img)), get conv(SKLayer m (conv(img)));

[0066] S4.4. Output the restored image img+conv(SKLayer m (conv(img)));

[0067] S4.5. Establish a reconstruction loss function between the output restored image and its corresponding true value image, and use the loss function to constrain the deraining network structure, wherein the loss function is defined as:

[0068]

[0069] in, represents the pixel value of channel i of the true image, represents the derained image obtained after network processing, express The pixel value of channel i, N is the number of image pairs, F s is a smooth function, defined as:

[0070]

[0071] S5. The restored image and the true value image are used as training sets and input into the deraining network structure for training iterations. For each iteration, a gradient descent method is used to complete the training of the network, wherein the gradient descent method includes the following steps:

[0072] S5.1, set the learning rate to 0.0001 and the batch size to 8;

[0073] S5.2. Adopt the adam algorithm as the optimization algorithm, with parameters β1=0.9 and β2=0.99. After 200 iterations, calculate the structural similarity (SSIM) and peak signal-to-noise ratio (PSNR) on the test set every 20 rounds. After 500 rounds of iteration, take the result with the highest SSIM and PSNR as the final training result.

[0074] S6. After training iterations, the image to be derained is input into the trained deraining network structure to obtain a derained image.

[0075] In order to verify the innovation of the present invention, tests were conducted. The test results are as follows: Peak signal-to-noise ratio (PSNR) and structural similarity (SSIM) were used for measurement. Experiments were conducted on multiple data sets including DDN, SPA, DID, RAIN200L and RAIN200H. The test results were better than the existing methods. The experimental results are shown in the following table.

[0076]

[0077] The above experiments show that the present invention has improved indicators compared with some transfomer-based methods, further illustrating the effectiveness of the patented method.

[0078] 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 should also be regarded as the scope of protection of the present invention.

Claims

1. An image deraining method based on a spatial state sequence model of a learnable function, characterized in that: The image to be derained is input into the deraining network having a spatial state sequence module based on a learnable function to obtain a derained image, including the following steps: S1. Construct a spatial state sequence model ISSM for 2D images; S2, construct a KAN layer based on learnable functions; S3, integrating step S1 and step S2 to form a spatial state sequence module SKLayer based on a learnable function; S4. Build a rain removal network structure based on the spatial state sequence module SKLayer of the learnable function; S5, taking the restored image and the true value image as training sets, inputting them into the deraining network structure for training iteration, and using the gradient descent method to complete the network training for each iteration; S6. After training iterations, the image to be derained is input into the trained deraining network structure to obtain a derained image.

2. The image deraining method based on a spatial state sequence model of a learnable function according to claim 1, characterized in that: Step S1 includes the following steps: S1.1, divide the 2D image into M×N non-overlapping sub-image blocks, and scan the image horizontally and vertically to obtain image embedding sequences X0 and X2, and reverse X0 and X2 to obtain X1 and X3; S1.2, define the space state model SSM, the formula is as follows: h t =Ah t-1 +Bx t y t =Ch t +Dx t ; Among them, matrices A, B, C, and D are spatial parameters, and x t is the input sequence, h t-1 is the state parameter at time t-1, h t is the state parameter at time t, y t is the output sequence; S1.3, the sequences X0, X1, X2, X3 obtained in step S1.1 are used as the input sequences of step S1.2 respectively, and the output sequences y0, y1, y2, y3 are obtained, and the 2D images Z0, Z1, Z2, Z3 are restored in the horizontal and vertical scanning order in step S1.1, and the images Z0, Z1, Z2, Z3 are combined to output Z: Z=Z0+Z1+Z2+Z3.

3. The image deraining method based on a spatial state sequence model of a learnable function according to claim 1, characterized in that: Step S2 includes the following steps: According to the Kolmogorov-Arnold theory, a multivariable continuous function f can be expressed as a combination of a finite number of single-variable continuous functions. The input is Z, and the output can be expressed as: Among them, Φ represents a learnable univariate function, ° represents matrix multiplication, n represents the number of functions, and KANLayer(U) represents the output.

4. The image deraining method based on a spatial state sequence model of a learnable function according to claim 1, characterized in that: In step S3, the input is defined as input, including the following steps: S3.1, use layer normalization to process the input input, obtain the output LM(input), and concatenate the output of the layer normalization as the input of the KAN layer based on the learnable function in step S2, and output KANLayer(LM(input)); S3.2, after splicing, it is used as the input of the space state sequence model ISSM in step S1, and the ISSM is obtained after being used n times continuously. n (KANLayer(LM(input))); The output is concatenated again as the input of the KAN layer based on the learnable function in step S2, and the output is KANLayer(ISSM n (KANLayer(LM(input)))); S3.3, the spliced ​​output and input are output together to obtain KANLayer (ISSM n (KANLayer(LM(input))))+input.

5. The image deraining method based on a spatial state sequence model of a learnable function according to claim 1, characterized in that: In step S4, the deraining network structure consists of a 3*3 convolutional layer for pre-processing, several SKLayers for extracting deep features, and a 3*3 convolutional layer for post-processing, including the following steps: S4.1, the input image is processed using a 3×3 convolutional layer to obtain conv(img); S4.

2. Process conv(img) obtained in step S4.1 through m SKLayer processes, and outputs SKLayer m (conv(img)); S4.

3. Use a 3×3 convolutional layer to process the SKLayer obtained in step S4.2 m (conv(img)), get conv(SKLayer m (conv(img))); S4.

4. Output the restored image img+conv(SKLayer m (conv(img))).

6. The image deraining method based on a spatial state sequence model of a learnable function according to claim 1, characterized in that: In step S4, a reconstruction loss function is established between the output restored image and its corresponding true value image, and the loss function is used to constrain the deraining network structure, wherein the loss function is defined as: in, represents the pixel value of channel i of the true image, represents the derained image obtained after network processing, express The pixel value of channel i, N is the number of image pairs, F s is a smooth function, defined as:

7. The image deraining method based on a spatial state sequence model of a learnable function according to claim 1, characterized in that: In step S5, training the network using the gradient descent method includes the following steps: S5.1, set the learning rate to 0.0001 and the batch size to 8; S5.

2. Adopt the adam algorithm as the optimization algorithm. After 200 iterations, calculate the SSIM and PSNR on the test set every 20 rounds. After 500 iterations, take the result with the highest SSIM and PSNR as the final training result.

Images