Multi-task image processing method based on wavelet transform and cross attention mechanism

Abstract

The invention discloses a multi-task image processing method based on wavelet transform and a cross attention mechanism, and relates to the technical field of image processing, and the method comprises the steps: respectively extracting global features and local texture details of an image through an RGB feature extraction network and a wavelet high and low frequency component feature extraction network; by calculating cross attention, high-frequency information, low-frequency information and color space information are effectively fused; in the aspect of multiple task heads, sub-task heads are separated through a neck branch network, adaptive weights are added to shared features through an attention module, meanwhile, loss weights and masks are adjusted according to actual needs, the efficiency and performance of associated sub-tasks are improved, and mutual interference between irrelevant sub-tasks is reduced.

Description

Technical Field

[0001] The present invention relates to the field of image processing technology, and in particular to a multi-task image processing method based on wavelet transform and cross-attention mechanism. Background Art

[0002] Currently, deep neural networks used in image processing typically use a backbone network to extract image features, a neck network to fuse features, and a head network to output predictions. When multiple subtasks are required, a shared convolutional layer is often used to extract features, with each subtask using a specific head network. This allows for joint multi-task training and feature sharing.

[0003] However, existing image processing techniques still have shortcomings in multi-task processing: relying on a single input image to extract multi-scale features makes it difficult to effectively capture the image's global structure and local details; methods such as feature pyramids struggle to rationally and fully combine global features with texture information in an image, failing to achieve optimal feature expression; and when multi-task heads are weakly correlated and independent of each other, joint training can interfere with each other, leading to a decline in overall performance. Therefore, a multi-task image processing neural network based on wavelet transform and cross-attention mechanism is urgently needed to address these issues. Summary of the Invention

[0004] The purpose of the present invention is to provide a multi-task image processing method based on wavelet transform and cross-attention mechanism to solve the problems raised in the prior art.

[0005] To achieve the above objectives, the present invention provides the following technical solution: a multi-task image processing method based on wavelet transform and cross-attention mechanism, the processing method comprising:

[0006] Step S1: Read the target image, extract the feature information in the target image through the RGB feature extraction network and the wavelet high and low frequency component feature extraction network, and obtain a feature map;

[0007] Step S2: calculating the cross-attention feature map between the feature map obtained by the RGB feature extraction network and the feature map obtained by the wavelet high- and low-frequency component feature extraction network, and realizing feature fusion using residual connection, upsampling and downsampling;

[0008] Step S3: The high-resolution feature map collected by upsampling and the low-resolution feature map collected by downsampling are passed through the attention module and the neck branch network in sequence, and the high-resolution subtask and the low-resolution subtask are performed respectively.

[0009] Furthermore, step S1 includes:

[0010] The RGB feature extraction network uses a pre-trained lightweight backbone network and removes the classification layer; when the RGB feature extraction network extracts feature information from the target image, it generates several feature maps of different resolutions.

[0011] Furthermore, step S1 includes:

[0012] The wavelet high- and low-frequency component feature extraction network includes four branches, and the feature extraction components corresponding to the four branches are low-frequency component LL, high-frequency component LH, high-frequency component HL, and high-frequency component HH, respectively; wherein each branch is composed of several groups of convolutional pooling blocks; each group of convolutional pooling blocks is composed of a convolutional layer, a batch normalization layer, and an activation function layer through a residual connection;

[0013] In the above steps, the low-frequency component LL, high-frequency component LH, high-frequency component HL, and high-frequency component HH are extracted as follows:

[0014]

[0015] Among them, I is the input grayscale image; i and j are the pixel coordinates in the output image respectively. After extraction, the output resolution is reduced to half of the original feature map.

[0016] Furthermore, step S2 includes:

[0017] The calculation process of the cross-attention feature map includes: flattening the last two dimensions of the shape matrices of the feature map obtained by the RGB feature extraction network and the feature map obtained by the wavelet high and low frequency component feature extraction network into one-dimensional matrices respectively to obtain a query matrix, a key matrix, and a value matrix; inputting the query matrix, the key matrix, and the value matrix into the attention matrix calculation formula to obtain an attention matrix; restoring the last dimension of the attention matrix to the height-by-width dimension of the feature map obtained by the RGB feature extraction network, so that the attention matrix of the cross-attention feature map is the same as the shape matrix of the feature map obtained by the RGB feature extraction network;

[0018] The residual connection includes adding the feature map obtained by the RGB feature extraction network to the cross attention feature map to output a feature map with a resolution of 128*128 and a feature map with a resolution of 32*32;

[0019] Upsampling includes upsampling the 128*128 resolution feature map to a 256*256 resolution feature map through a deconvolution network; adding the high-frequency feature map and the low-frequency feature map in the 256*256 resolution feature map to obtain a 256*256 resolution comprehensive feature map; and upsampling the 256*256 resolution comprehensive feature map to a 512*512 resolution feature map through deconvolution.

[0020] Downsampling includes downsampling the 32*32 resolution feature map to a 16*16 resolution feature map through point convolution and pooling layers; adding the high-frequency feature map and the low-frequency feature map in the 16*16 resolution feature map to obtain a comprehensive feature map with a resolution of 16*16; downsampling the comprehensive feature map with a resolution of 16*16 to a 8*8 resolution feature map through point convolution and pooling layers;

[0021] The cross attention calculation formula for the above steps is as follows:

[0022]

[0023] Among them, x1 and x2 are RGB feature map and wavelet component feature map respectively; the shape matrix of the feature map is (B, C, H, W); among them, B, C, H, and W represent the batch size, the number of feature channels, the feature map height, and the feature map width respectively;

[0024] The above steps flatten the last two dimensions H and W of the shape matrix of the feature map into a one-dimensional matrix to obtain three matrices for calculating attention. These three matrices are the query matrix Q, the key matrix K, and the value matrix V. The three matrices are all (B, C, N); where N = H × W;

[0025] The attention matrix calculation formula is as follows:

[0026]

[0027] Where C is the number of channels of the feature map; the attention matrix A is calculated by the above formula, and the shape of the attention matrix A is (B, C, N); the last dimension N of the attention matrix A is restored to H×W, that is, the output is transformed into A′(B, C, H, W), which is the same shape matrix as the input x1.

[0028] Furthermore, step S3 includes:

[0029] The attention module includes channel attention and spatial attention.

[0030] Furthermore, channel attention generates a channel weight matrix by performing average pooling, maximum pooling, two-layer point convolution and one-layer activation function operations on the input feature map, and multiplies the channel weight matrix with the input feature map;

[0031] The input of the channel attention in the above steps is the feature map x (the shape matrix of the feature map x is (B, C, H, W)), and the formula is as follows:

[0032] avg_out=Conv2d(ReLU(Conv2d(AvgPool(x)))));

[0033] max_out=Conv2d(ReLU(Conv2d(MaxPool(x)))));

[0034] out=Sigmoid(avg out +max out )×x;

[0035] Among them, Conv2d is 1×1 point convolution; ReLU is the activation function; AvgPool is average pooling; MaxPool is maximum pooling; the above steps average pool and max pool the last two dimensions (H×W) of the shape matrix of the input feature map x to 1 point respectively, and obtain a feature map x with a shape matrix of (B, C, 1, 1); then the feature map x is subjected to 2 layers of point convolution and 1 layer of activation function to generate the weights of the C channels of the input feature map; the two channel weights are added to obtain the channel weight matrix (B, C, 1, 1), which is normalized by the Sigmoid function and multiplied with the input feature map x to complete the channel attention weighting of the input features.

[0036] Furthermore, spatial attention calculates the average and maximum values of each channel of the input feature map, concatenates the weight matrix corresponding to the average value with the weight matrix corresponding to the maximum value, and sequentially passes the concatenated weight matrix through one layer of point convolution and one layer of Sigmoid function to generate a spatial weight matrix; the spatial weight matrix is multiplied by the input feature map;

[0037] In the above steps, the input of spatial attention is also the feature map x, and the formula is as follows:

[0038]

[0039] out=Sigmoid(Conv2d(concat(avg_out, max_out)))×x;

[0040] The input feature map x is calculated along each channel, and the average and maximum values in the feature map x are calculated. The shape matrix at this time is (B, 1, H, W). After splicing, the spatial weight matrix (B, 2, H, W) is obtained. Then, it is passed through one layer of point convolution and one layer of Sigmoid function to generate the weights in the H×W space of the input feature map. The spatial attention weight matrix is (B, 1, H, W). The spatial attention weight matrix is multiplied by the input feature map x to complete the spatial attention weighting of the input features.

[0041] Furthermore, step S3 includes:

[0042] The neck branch network consists of several 1×1 point convolutions, batch normalization, and activation functions. The detection task head of the high-resolution subtask uses a heat map to predict the center point position and target box size, and the segmentation task head uses a heat map to predict the probability of belonging to the target pixel. The classification task head of the low-resolution subtask is connected to a fully connected layer, and the sigmoid function or softmax function is used for normalization and then multiplied by the corresponding total score.

[0043] Compared with the existing technology, the beneficial effects of the present invention are: the present invention combines the RGB feature extraction network with the wavelet transform feature extraction network to simultaneously extract the global information and local details of the image; combines cross attention to fully and effectively integrate high-frequency information, low-frequency information and color space information; by introducing the neck branch network and attention mechanism, the efficiency and performance of related subtasks are improved, and the negative images between irrelevant subtasks are reduced. The model has a simple structure and strong scalability. BRIEF DESCRIPTION OF THE DRAWINGS

[0044] Figure 1 Schematic diagram of the RGB feature extraction network flow of the multi-task image processing method based on wavelet transform and cross attention mechanism of the present invention;

[0045] Figure 2 Schematic diagram of the network flow for extracting high and low frequency component features of the multi-task image processing method based on wavelet transform and cross attention mechanism of the present invention;

[0046] Figure 3 Schematic diagram of the feature fusion network structure of the multi-task image processing method based on wavelet transform and cross attention mechanism of the present invention;

[0047] Figure 4 Schematic diagram of the high-resolution overall task flow of the multi-task image processing method based on wavelet transform and cross-attention mechanism of the present invention;

[0048] Figure 5 Schematic diagram of the low-resolution subtask process of the multi-task image processing method based on wavelet transform and cross-attention mechanism of the present invention;

[0049] Figure 6 Schematic diagram of the model training process of the multi-task image processing method based on wavelet transform and cross attention mechanism of the present invention;

[0050] Figure 7 Schematic diagram of the model reasoning process of the multi-task image processing method based on wavelet transform and cross-attention mechanism of the present invention. DETAILED DESCRIPTION

[0051] Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without making any creative work shall fall within the scope of protection of the present invention.

[0052] Example: Figure 1-Figure 7 As shown, the present invention provides a technical solution, a multi-task image processing method based on wavelet transform and cross attention mechanism, the processing method includes:

[0053] Step S1: Read the target image, extract the feature information in the target image through the RGB feature extraction network and the wavelet high and low frequency component feature extraction network, and obtain a feature map;

[0054] Wherein, step S1 includes:

[0055] The RGB feature extraction network uses a pre-trained lightweight backbone network and removes the classification layer; when the RGB feature extraction network extracts feature information from the target image, it generates several feature maps of different resolutions;

[0056] See Figure 1 , the input image resolution is 512×512, and after feature extraction, the output feature maps with resolutions of 128×128 and 32×32 respectively;

[0057] The wavelet high- and low-frequency component feature extraction network includes four branches, and the feature extraction components corresponding to the four branches are low-frequency component LL, high-frequency component LH, high-frequency component HL, and high-frequency component HH respectively; wherein each branch is composed of several groups of convolutional pooling blocks; each group of convolutional pooling blocks is composed of a convolutional layer, a batch normalization layer, and an activation function layer through a residual connection;

[0058] See Figure 2 ,The input image resolution is 512×512, and after wavelet transform the resolution is 256×256. After passing through three sets of feature extraction modules, the ,output feature maps of 128×128 and 32×32 respectively;

[0059] Step S2: calculating the cross-attention feature map between the feature map obtained by the RGB feature extraction network and the feature map obtained by the wavelet high- and low-frequency component feature extraction network, and realizing feature fusion using residual connection, upsampling and downsampling;

[0060] See Figure 3 , wherein step S2 includes:

[0061] The calculation process of the cross-attention feature map includes: flattening the last two dimensions of the shape matrices of the feature map obtained by the RGB feature extraction network and the feature map obtained by the wavelet high and low frequency component feature extraction network into one-dimensional matrices respectively to obtain a query matrix, a key matrix, and a value matrix; inputting the query matrix, the key matrix, and the value matrix into the attention matrix calculation formula to obtain an attention matrix; restoring the last dimension of the attention matrix to the height-by-width dimension of the feature map obtained by the RGB feature extraction network, so that the attention matrix of the cross-attention feature map is the same as the shape matrix of the feature map obtained by the RGB feature extraction network;

[0062] The residual connection includes adding the feature map obtained by the RGB feature extraction network to the cross attention feature map to output a feature map with a resolution of 128*128 and a feature map with a resolution of 32*32;

[0063] The upsampling includes upsampling the 128*128 resolution feature map to a 256*256 resolution feature map through a deconvolution network; adding the high-frequency feature map and the low-frequency feature map in the 256*256 resolution feature map to obtain a 256*256 resolution comprehensive feature map; and upsampling the 256*256 resolution comprehensive feature map to a 512*512 resolution feature map through deconvolution;

[0064] The downsampling includes downsampling a 32*32 resolution feature map to a 16*16 resolution feature map through point convolution and a pooling layer; adding a high-frequency feature map and a low-frequency feature map in the 16*16 resolution feature map to obtain a comprehensive feature map with a 16*16 resolution; and downsampling the 16*16 resolution comprehensive feature map to an 8*8 resolution feature map through point convolution and a pooling layer;

[0065] Step S3: The high-resolution feature map collected by upsampling and the low-resolution feature map collected by downsampling are sequentially passed through the attention module and the neck branch network, and the high-resolution subtask and the low-resolution subtask are performed respectively;

[0066] Wherein, step S3 includes:

[0067] The attention module includes channel attention and spatial attention;

[0068] Among them, the channel attention generates a channel weight matrix by performing average pooling, maximum pooling, 2-layer point convolution and 1-layer activation function operations on the input feature map in sequence, and multiplies the channel weight matrix with the input feature map;

[0069] Among them, spatial attention calculates the average and maximum values of each channel of the input feature map, and concatenates the weight matrix corresponding to the average value with the weight matrix corresponding to the maximum value. The concatenated weight matrix is sequentially subjected to one layer of point convolution operation and one layer of Sigmoid function to generate a spatial weight matrix; the spatial weight matrix is multiplied by the input feature map;

[0070] Wherein, step S3 further includes:

[0071] The neck branch network consists of several 1×1 point convolutions, batch normalization, and activation functions; see Figure 4 , where the detection task head of the high-resolution subtask uses the heat map to predict the center point position and target box size, and the segmentation task head uses the heat map to predict the attribution probability of the target pixel; see Figure 5 , where the classification task head of the low-resolution subtask is connected to the fully connected layer, and the sigmoid function or softmax function is used for normalization and then multiplied by the corresponding total score;

[0072] In an embodiment of the present invention, the above-mentioned model is used in the multi-index evaluation of colposcopy images, including: when using a colposcope for vaginal examination, it is necessary to simultaneously evaluate the vaginal wall, secretions, cervix and other areas, and the evaluation indicators include the elasticity, fullness, wall size, degree of wrinkles, and mucosal color of the vaginal wall; the amount, texture, and color of secretions; and the color, erosion, and hypertrophy of the cervix. The model evaluates the scores of these indicators;

[0073] The above model structure includes:

[0074] The high-resolution branch includes target segmentation, which includes the cervical os and vaginal wall. The loss function is used to calculate the loss during model training and is deleted during inference.

[0075] The cervical branch network at low resolution includes three branches: vaginal wall branch, secretion branch, and cervical os branch;

[0076] The vaginal wall branch contains a classification head to distinguish whether the image contains the vaginal wall branch. Several regression heads correspond to the evaluation results of indicators such as elasticity, fullness, wall size, wrinkles, and mucosal color.

[0077] The secretion branch is similar to the vaginal wall, and contains a classification head corresponding to whether the secretion branch is included, and several regression heads corresponding to the amount, texture, color, etc. of the secretion;

[0078] Among them, the cervical os branch is similar to the vaginal wall, including a classification head corresponding to whether the cervical os branch is included, and several regression heads corresponding to the color, erosion, and hypertrophy degree of the cervix.

[0079] The above model includes a segmentation, multiple classification, and multiple regression tasks in practical applications. Among them, the Dice loss function is used to optimize the prediction segmentation task, the binary cross entropy loss function is used to optimize the classification results, and the Smooth1Loss is used to optimize the regression score.

[0080] Among them, the Dice indicator measures the degree of overlap between the predicted segmentation area and the true segmentation area, and the Dice loss function formula is changed as follows:

[0081]

[0082] Wherein, K represents the number of regions to be segmented, which in the embodiment of the present invention is the vaginal wall and the cervical os; i and j are pixel coordinates; Represents the prediction result of region k at pixel (i, j); y k，i，j represents the true label of the k-region at pixel (i, j); ε is a constant 10 -6 , to prevent division by 0;

[0083] Among them, the binary cross entropy loss function is used for the two-classification model, and the formula is as follows:

[0084]

[0085] Wherein, N is the number of samples, which in the embodiment of the present invention is three branches: vaginal wall, cervix, and secretions; y i represents the true label; represents the predicted label;

[0086] Among them, SmoothL1Loss is a piecewise function that combines absolute error and mean square error to describe the gap between the predicted value and the true value. The formula is as follows:

[0087]

[0088] Among them, y i represents the true value; Represents the predicted value; when |x|≥1, it is the absolute error loss, which is used for fast convergence; when |x|<1, it is the mean square error loss, which is used for stable convergence;

[0089] Among them, the total loss function is:

[0090] Loss total =λ1·Loss Dice +λ2·Weight·Loss bce +λ3·Mask·Loss smooth_l1 ;

[0091] Among them, λ1, λ2, and λ3 are weight coefficients, all of which are 1 here;

[0092] Among them, Weight is a weight matrix generated based on the actual value of "whether the region exists" in the current batch. The proportion of non-existent regions in the real data is small, so setting a higher weight when the region does not exist allows the model to better utilize the labeled data. Here, the weights are set to 0.1 and 1 respectively;

[0093] Among them, Mask is a mask generated according to the true value of "whether the region exists" in the current batch. For example, if the cervix region is not visible in the image, then all predictions related to cervix assessment will not be calculated.

[0094] See Figure 6 , the training process of the above model is:

[0095] (1) Model loading;

[0096] (2) Read in the colposcopic image, the target segmentation area, and the score numerical labels of each indicator, including the classification label of "whether the image contains the area corresponding to the indicator", and perform forward inference of the model;

[0097] (3) Obtain the multi-task prediction results such as the predicted target segmentation results, the probability value of whether each region exists, and the score of each indicator, respectively. Use the above loss functions to calculate the segmentation loss, classification loss, and regression loss, and use the weights and masks to calculate the total loss;

[0098] (4) Back propagation, update weights, and save the model;

[0099] See Figure 7 , the reasoning process of the above model is:

[0100] (1) Model loading;

[0101] (2) Read the colposcopy image and perform forward model reasoning; delete the target segmentation related branches to improve the reasoning speed;

[0102] (3) Obtain the classification results of whether each region exists and the score regression results of each indicator respectively, and use the classification results to assist in giving the final comprehensive evaluation score.

[0103] Finally, it should be noted that the above descriptions are merely preferred embodiments of the present invention and are not intended to limit the present invention. Although the present invention has been described in detail with reference to the aforementioned embodiments, those skilled in the art will be able to modify the technical solutions described in the aforementioned embodiments or substitute equivalents for some of the technical features. Any modifications, equivalent substitutions, and improvements made within the spirit and principles of the present invention shall be included within the scope of protection of the present invention.

Claims

1. A multi-task image processing method based on wavelet transform and cross-attention mechanism, characterized by: The processing method comprises: Step S1: Read the target image, extract the feature information in the target image through the RGB feature extraction network and the wavelet high and low frequency component feature extraction network, and obtain a feature map; Step S2: calculating the cross-attention feature map between the feature map obtained by the RGB feature extraction network and the feature map obtained by the wavelet high- and low-frequency component feature extraction network, and realizing feature fusion using residual connection, upsampling and downsampling; Step S3: The high-resolution feature map collected by upsampling and the low-resolution feature map collected by downsampling are passed through the attention module and the neck branch network in sequence, and the high-resolution subtask and the low-resolution subtask are performed respectively.

2. The multi-task image processing method based on wavelet transform and cross attention mechanism according to claim 1, characterized in that: Step S1 includes: the RGB feature extraction network uses a pre-trained lightweight backbone network and removes the classification layer; when the RGB feature extraction network extracts feature information from the target image, it generates several feature maps with different resolutions.

3. The multi-task image processing method based on wavelet transform and cross attention mechanism according to claim 1, characterized in that: Step S1 includes: the wavelet high and low frequency component feature extraction network includes four branches, and the feature extraction components corresponding to the four branches are low frequency component LL, high frequency component LH, high frequency component HL, and high frequency component HH, respectively; wherein each of the branches is composed of several groups of convolutional pooling blocks; each group of convolutional pooling blocks is composed of a convolutional layer, a batch normalization layer, and an activation function layer through a residual connection.

4. The multi-task image processing method based on wavelet transform and cross attention mechanism according to claim 1, characterized in that: Step S2 includes: The calculation process of the cross-attention feature map includes: flattening the last two dimensions of the shape matrices of the feature map obtained by the RGB feature extraction network and the feature map obtained by the wavelet high- and low-frequency component feature extraction network into one-dimensional matrices respectively to obtain a query matrix, a key matrix, and a value matrix; inputting the query matrix, the key matrix, and the value matrix into the attention matrix calculation formula to obtain an attention matrix; restoring the last dimension of the attention matrix to the height-by-width dimension of the feature map obtained by the RGB feature extraction network, so that the attention matrix of the cross-attention feature map is the same as the shape matrix of the feature map obtained by the RGB feature extraction network; The residual connection includes adding the feature map obtained by the RGB feature extraction network to the cross attention feature map to output a feature map with a resolution of 128*128 and a feature map with a resolution of 32*32; The upsampling includes upsampling the 128*128 resolution feature map to a 256*256 resolution feature map through a deconvolution network; adding the high-frequency feature map and the low-frequency feature map in the 256*256 resolution feature map to obtain a 256*256 resolution comprehensive feature map; and upsampling the 256*256 resolution comprehensive feature map to a 512*512 resolution feature map through deconvolution; The downsampling includes downsampling a 32*32 resolution feature map to a 16*16 resolution feature map through point convolution and a pooling layer; adding a high-frequency feature map and a low-frequency feature map in the 16*16 resolution feature map to obtain a comprehensive feature map with a 16*16 resolution; and downsampling the 16*16 resolution comprehensive feature map to an 8*8 resolution feature map through point convolution and a pooling layer.

5. The multi-task image processing method based on wavelet transform and cross attention mechanism according to claim 1, characterized in that: Step S3 includes: the attention module includes channel attention and spatial attention.

6. The multi-task image processing method based on wavelet transform and cross-attention mechanism according to claim 5, characterized in that: The channel attention generates a channel weight matrix by performing average pooling, maximum pooling, 2-layer point convolution and 1-layer activation function operations on the input feature map in sequence, and multiplies the channel weight matrix with the input feature map.

7. The multi-task image processing method based on wavelet transform and cross-attention mechanism according to claim 5, characterized in that: The spatial attention calculates the average and maximum values of the input feature map along each channel, and concatenates the weight matrix corresponding to the average value with the weight matrix corresponding to the maximum value. The concatenated weight matrix is sequentially subjected to one layer of point convolution operation and one layer of Sigmoid function to generate a spatial weight matrix; the spatial weight matrix is multiplied by the input feature map.

8. The multi-task image processing method based on wavelet transform and cross attention mechanism according to claim 1, characterized in that: Step S3 includes: the neck branch network is composed of several 1×1 point convolutions, batch normalization, and activation functions; the detection task head of the high-resolution subtask uses the heat map to predict the center point position and target box size, and the segmentation task head uses the heat map to predict the attribution probability of the target pixel; the classification task head of the low-resolution subtask is connected to the fully connected layer, and the sigmoid function or softmax function is used for normalization and then multiplied by the corresponding total score.

Images