A small sample semantic segmentation method based on query self-correlation reinforcement

Abstract

The application discloses a small sample semantic segmentation method based on query self-correlation reinforcement, and belongs to the technical field of image segmentation. In order to solve the problems existing in the prior art, the method comprises the following steps: step 1, constructing a network model: a backbone network sharing weights; prototype generation; a transformer decoder, input support prototype and query feature; a reinforcement prior module, using the interaction between high-level features to obtain prior mask information with high guidance; a feature enrichment module, using support prototypes and prior masks to enrich query features from different scales; using the output of the decoder to filter out the predicted mask from the query feature; step 2, preparing a data set: preprocessing the data set according to the small sample segmentation task; step 3, training the network model: inputting the data set prepared in step 2 into the network model constructed in step 1 for training; step 4, selecting a minimum loss function and an optimal evaluation index; step 5, fine-tuning and saving the model.

Description

TECHNICAL FIELD

[0001] The application relates to a small sample semantic segmentation method based on query autocorrelation reinforcement and belongs to the technical field of image segmentation. BACKGROUND

[0002] Image semantic segmentation is one of the key problems in the field of computer vision, which aims to classify image data at the pixel level. Specifically, the goal is to convert an image into a mask with highlighted areas, and each pixel in the image is assigned a class according to the highlighted area. However, in order to obtain a reliable segmentation model in the semantic segmentation task, it is necessary to use a large number of densely annotated images for supervised training, which is difficult to achieve in practical applications. In order to alleviate the problem of data scarcity in real scenarios, small sample semantic segmentation has emerged. Due to the small number of samples, there is a difficult problem in small sample segmentation that cannot be ignored, that is, the intra-class difference is large. When the target class appears as different appearances between the query image and the support image, it is difficult for the model to associate the target class information on the support image with the target class information on the query image, resulting in poor prediction results.

[0003] Chinese Patent Publication No. CN117726809A, entitled "Small Sample Semantic Segmentation Method Based on Information Interaction Enhancement", first extracts support features and query features through a backbone network, processes the features using a guided attention module and a spatial information interaction attention module, then refines the query features using a query prior mask, and finally inputs the refined query features into a decoder for processing to obtain a predicted segmentation result. However, this method has high computational complexity, blindly mines sample information, and cannot effectively alleviate the problem of appearance difference between support samples and query images. SUMMARY

[0004] The present application solves the problem of high computational complexity, blind mining of sample information, and inability to effectively alleviate the problem of appearance difference between support samples and query images in the prior art. A small sample semantic segmentation method based on query autocorrelation reinforcement is proposed. This method maintains a low parameter amount, uses the autocorrelation between pixels within the query image to drive the transfer of interaction information to the query data, enhances the utilization rate, and avoids overfitting caused by excessive mining of sample information. This method can achieve better segmentation results under data scarcity conditions and effectively improve the generalization ability of the model.

[0005] The technical problem solving scheme of the present application is:

[0006] A small sample semantic segmentation method based on query autocorrelation reinforcement, comprising the following steps:

[0007] Step 1, constructing a network model: the entire network includes six parts, the first part is a backbone network with shared weights, which is used to extract support features and query features; the second part is prototype generation, which is used to abstract features into prototype vectors with semantic information; the third part is a Transformer decoder, which inputs support prototypes and query features to obtain mask embeddings of target objects; the fourth part is a reinforcement prior module, which uses the interaction between high-level features to obtain prior mask information with high guidance; the fifth part is a feature enrichment module, which uses support prototypes and prior masks to enrich query features from different scales; the sixth part is to filter out the predicted mask from the query features using the output of the decoder;

[0008] Step 2, preparing the data set: the data set is preprocessed according to the small sample segmentation task;

[0009] Step 3, training the network model: input the data set prepared in step 2 into the network model constructed in step 1 for training;

[0010] Step 4, selecting the minimum loss function and the optimal evaluation index: by minimizing the loss function of the output image and the label, until the training times reach the set threshold or the value of the loss function reaches the set range, the model parameters can be considered to have been pre-trained, and the model parameters are saved; At the same time, select the optimal evaluation index to measure the accuracy of the algorithm and evaluate the performance of the system;

[0011] Step 5, fine-tuning and saving the model: adjust part of the parameters and train again until the network model reaches the optimal generalization ability and segmentation result under a certain parameter.

[0012] The backbone network in step 1 adopts an expanded version of Resnet, and the output is a dictionary of four-dimensional features. Each layer of features contains 1x1 and 3x3 convolution operations, ReLU operations, and prototypes use the mask average pooling method in traditional research.

[0013] The data set in step 2 is the small sample benchmark data set Pascal-5i. First, divide the data set into a training set and a test set, and then divide the training set and the test set into a support set and a query set. During training, select k support samples containing class c from the support set, and select one query sample containing class c from the query set, aiming to learn a model that can segment class c in the support sample and also segment class c in the query sample. Each time a support-query combination is taken as a training task.

[0014] Before the data set is input to step 1 in step 3, the pictures are preprocessed, that is, all picture sizes are adjusted to 473*473, and images that do not contain any valid categories and valid objects that are too small, that is, less than 2*32*32 pixels, are deleted, so as to ensure that all images input into the model contain the specified categories and the target object is large enough.

[0015] The loss function is selected as Dice loss in the training process in step 4, and Dice loss directly measures the proportion of overlapping areas, which can better guide the model to focus on the target area in the case of sample scarcity.

[0016] The mean intersection over union (mIOU) and foreground-background intersection over union (FB-IoU) are selected as evaluation indexes in the training process in step 4, which can effectively evaluate the precision and quality of the algorithm and measure the effect of the segmentation network.

[0017] The beneficial effects of the present application are as follows:

[0018] 1. The middle layer features are used to generate prototypes, because the middle layer features constitute the target object part of the new class, which is more helpful to improve the generalization ability of the model. The high layer features are used to generate reinforced priori, which can not only reduce the preference for the base class, but also realize the acquisition of high-quality meta-knowledge, and effectively improve the segmentation precision of the model.

[0019] 2. The designed reinforced priori module is used, which aims to obtain more guiding priori information for subsequent input into the FEM module to enrich the query features. The module can effectively improve the interaction between the support and the query, enhance the relevance of the two, and effectively alleviate the appearance difference problem between the same classes.

[0020] 3. The Dice loss function is selected, which has more advantages compared with other loss functions. For example, cross-entropy mainly measures the difference between the predicted probability distribution and the true distribution, and in the case of serious class imbalance, the model is prone to ignore the minority class. The Dice loss directly optimizes the overlap of the target area, so that the model has stronger robustness when predicting the minority class, and can better improve the segmentation quality of the minority sample class. BRIEF DESCRIPTION OF DRAWINGS

[0021] Figure 1 The network structure diagram of the query self-correlation reinforced small sample semantic segmentation method of the application.

[0022] Figure 2 The structure diagram of the backbone network.

[0023] Figure 3 The structure diagram of the Transformer decoder.

[0024] Figure 4 The specific operation steps of the reinforcement prior module of the present application. DETAILED DESCRIPTION

[0025] The present application will be further described in detail below with reference to the accompanying drawings.

[0026] As shown in Figure 1 , a small sample semantic segmentation method based on query self-correlation reinforcement, which specifically comprises the following steps:

[0027] Step 1, constructing a network model: the entire network includes six parts, the first part is a backbone network with shared weights, which is used to extract support features and query features; the second part is a prototype generation, which is used to abstract the features into a prototype vector with semantic information; the third part is a Transformer decoder, which inputs the support prototype and the query feature, and is used to obtain the mask embedding of the target object; the fourth part is a reinforcement prior module, which uses the interaction between high-level features to obtain prior mask information with high guidance; the fifth part is a feature enrichment module, which uses the support prototype and the prior mask to enrich the query feature from different scales; the sixth part is to filter out the predicted mask from the query feature using the output of the decoder.

[0028] Step 2, preparing the data set: the data set is preprocessed according to the small sample segmentation task;

[0029] Step 3, training the network model: input the data set prepared in step 2 into the network model constructed in step 1 for training;

[0030] Step 4, selecting the minimum loss function and the optimal evaluation index: by minimizing the loss function of the output image and the label, until the training times reach the set threshold or the value of the loss function reaches the set range, the model parameters can be considered to have been pre-trained, and the model parameters are saved; At the same time, select the optimal evaluation index to measure the accuracy of the algorithm and evaluate the performance of the system;

[0031] Step 5, fine-tuning and saving the model: adjust part of the parameters, and train again until the network model reaches the optimal generalization ability and segmentation result under a certain parameter.

[0032] Embodiment:

[0033] The first part of step 1 is a backbone network structure with shared weights, as shown in Figure 2 , the final output of the backbone network is a dictionary containing four layers of features, which are the outputs of four residual blocks conv2_x, conv3_x, conv4_x, and conv5_x, respectively, and the structure of each residual block is similar to Figure 2The conv3_x block shown on the right undergoes a 1×1 convolution, a 3×3 convolution, and a 1×1 convolution from input to output. The input is then expanded to the same dimension as the output using a 1×1 convolution and summed. The output dimensions of the four residual blocks are 256d, 512d, 1024d, and 2048d, respectively. conv3_x and conv4_x are defined as mid-level features, and conv5_x as high-level features. The mid-level features are concatenated along the channel dimension, followed by a 1×1 convolution to reduce the number of channels to 60. BatchNorm2d, ReLU, and Dropout are then used to normalize each channel to prevent exploding and vanishing gradients. Regularization is used to randomly set the output of some channels to 0 during training to reduce overfitting. The concatenated mid-level features are used to generate mid-level support prototypes, which are then fed into the Transformer decoder as conditional queries. The high-level features conv5_x are subsequently fed into the enhanced prior module.

[0034] The second part is prototype generation, which uses masked average pooling to aggregate support prototypes. First, the mask is applied to the feature map, then average pooling is applied to the feature map to obtain a single feature value, which is then multiplied by the total number of pixels in the feature map and finally divided by the area of ​​the masked region to obtain the prototype vector. The process is shown in the following formula:

[0035]

[0036] Among them, f pool For the pooling operation, is the feature map after the middle-layer support features are concatenated, ⊙ represents the Hadamard product, R e Representatives will support Mask M s Interpolation is Same size operation, That is, the generated initial prototype. The area of ​​the mask region is obtained by using average pooling to calculate the average value of the target class in the mask and then multiplying it by the total number of pixels in the feature map.

[0037] The third part is the Transformer decoder module, which has the following structure: Figure 3 First, the prototype The input is fed into the self-attention module in the form of residual connection, and after layer normalization, it is used as the query embedding of the cross attention module. As the key and value embedding of cross attention. After implementing cross attention, apply the feedforward layer and perform normalization to get the output This operation can enhance the separability of the query image foreground and background, enabling the extraction of more discriminative likelihood embeddings.

[0038] The specific steps of the fourth part of strengthening the prior module are as followsFigure 4 The support features, query features and mask of support image are input into the reinforced prior module, and the initial auxiliary prior information is generated by using the three. Specifically, first, the background pixels of the support features are filtered out by using the support mask, so that the pixels of the query features are not corresponding to the background on the support features, and then the cosine similarity between the pixels of the query features and the support features is calculated, and for each pixel of the query features, a support pixel point with the highest correspondence is taken, and the result is reshaped from h x w, 1 to h, w, 1, denoted as C q , and all values are normalized to the interval of 0-1 according to formula (3), and the result is Corr s , which represents the initial auxiliary prior mask, and ε is set to 10 -7 . Then, the flattened high-level query features are filtered out from each channel of the auxiliary prior to obtain high-confidence feature information, and the self-matching query vector is obtained by taking the average in the dimension. Since the similarity between the target class pixels (foreground pixels) is high and the noise tolerance is low, in order to ensure that the pixels in the high-confidence region exist and are sufficient, the threshold τ = 0.7 is selected. Then, in order to realize sufficient mining of the segmentation prior and enhance the high interaction between the support image and the query image, the support prototype under the high channel feature and the self-matching query vector FP pro are combined as a high-level class prototype, and the similarity spectrum C prior between the class prototype and the high-level query feature is used as the reinforced prior information. The pixel points with high median value in the similarity spectrum mean that they are strongly related to the pixels in the target object and are most likely to belong to the target class.

[0039] The fifth part feature enrichment module inputs the middle-level query features, support prototypes and reinforced prior information, respectively maps them to different scales (the sizes are 60 x 60, 30 x 30, 15 x 15 and 8 x 8 respectively), splices the three inputs into new feature information, independently interacts between different scales, and vertically uses the hierarchical relationship to enrich the rough feature map by using the basic information extracted from the more fine features through the top-down information path. After horizontal and vertical optimization, the spatial information on each scale is combined, up-sampling is performed, and the query feature with rich context information is obtained.

[0040] The sixth part in step 1 is to embed the possibility output by the Transformer decoder into the query feature F q according to the following formula to obtain the final predicted query mask

[0041] M q =upsam(ESC(F q ,p q ))

[0042] where ESC represents Einstein summation operation, which can be interpreted as the output of the Transformer generates weights for different scales, so that each pixel adaptively selects the scale information that is most powerful for its own classification. Specifically, F q with p q Pixel-wise multiplication and summation in the channel dimension, and the results are up-sampled to restore local details, thus obtaining the final prediction mask.

[0043] The data set in step 2 is a small sample benchmark data set Pascal-5i. The data set contains 12031 images. When obtaining support-query pairs from the data set, all image sizes are adjusted to 473x473, and images that do not contain any valid class and valid objects that are too small (less than 2x32x32 pixels) are deleted, ensuring that all images input into the model contain the specified class and the target object is large enough.

[0044] No other data augmentation strategy is used for training in step 3.

[0045] In step 4, the average value of the results of 5 experiments with different random seeds is taken as the evaluation result. The evaluation indicators are mean intersection over union (mIoU) and foreground-background intersection over union (FB-IoU). mIoU is the average value of the IoU of each class in the current fold, and IoU represents the ratio of the number of pixels in the intersection region between the predicted mask and the real mask of class C to the number of pixels in the union region, which is a global evaluation indicator. The implementation formula of IoU is as follows:

[0046]

[0047] where TP C , FP C , and FN C represent the number of true positives, false positives, and false negatives for class C in the segmentation prediction result, respectively.

[0048] The FB-IoU indicator ignores the difference between image classes, only needs to calculate the IoU of the foreground and background, only considers the separability between the two, has lower computational complexity, and is more sensitive to small sample evaluation.

[0049]

[0050] where IoU FG and IoU BG represent the foreground IoU and background IoU scores of the image, respectively.

[0051] The loss function selection selects a weighted Dice loss function, receives the prediction result of the model and the target segmentation mask as input. The function will expand the target mask by one channel dimension and convert it to a floating point type. Then, the prediction result and the target mask are flattened into one-dimensional vectors. Next, the function calculates the Dice coefficient by calculating the sum of squares of the prediction result and the target mask, and the intersection between them. By subtracting the Dice coefficient, the Dice loss is obtained. In order to further adjust the loss value, the function will weight the loss according to the set weight, and sum or average the loss value according to the configuration.

[0052] The experimental configuration environment is based on Windows-10 system, using PyTorch deep learning framework, on the GPU is NVIDIA GeForce RTX 4090D. During training, the Adam optimizer is used, the learning rate is 1e-3, and the epoch is set to 60.

[0053] The step 5 uses the image data set to fine-tune the whole network model parameters.

[0054] The application can complete more accurate binary segmentation under the condition of sample scarcity by constructing a small sample semantic segmentation network based on query autocorrelation reinforcement. The autocorrelation of the query image is effectively utilized, the pointing information more conducive to segmentation is mined from the query image itself, the high interactivity of the support image and the query image is enhanced, the multi-scale spatial information of the query feature is enriched, and the same class appearance difference problem between the support sample and the query sample is effectively alleviated. Under the same conditions, by calculating the relevant indicators of the images obtained by the existing method, the feasibility and superiority of the method are further verified. The relevant indicators of the existing technology and the method proposed in the application are compared as shown in Table 1:

[0055] Table 1 Comparison of relevant indicators of baseline and improved network of the application

[0056]

[0057] From the mIoU and FB-IoU results of the four folds in the table, it can be seen that the improved network proposed in the application effectively improves the segmentation effect of the model, and the segmentation effect of the new class is better than before.

Claims

1. A small sample semantic segmentation method based on query autocorrelation enhancement, characterized by: The method comprises the following steps: Step 1: Build the network model: The entire network consists of six parts. The first part is the backbone network with shared weights, which is used to extract support features and query features. The second part is prototype generation, which is used to abstract features into prototype vectors with semantic information. The third part is the Transformer decoder, which inputs the prototype and query features to obtain the mask embedding of the target object. The fourth part is the enhanced prior module, which uses the interaction between high-level features to obtain highly instructive prior mask information. The fifth part is the feature enrichment module, which uses support prototypes and prior masks to enrich query features at different scales. The sixth part is to use the decoder output to filter out the predicted mask from the query features; The backbone network uses an expanded version of Resnet, and the output is a dictionary of four-dimensional features. Each layer of features contains 1x1 and 3x3 convolution operations and ReLU operations. The prototype adopts the masked average pooling method used in traditional research. The fifth partial feature enrichment module takes in mid-level query features, supporting prototypes, and enhanced prior information, maps them to different scales, and concatenates the three inputs into new feature information. It independently interacts between different scales and vertically exploits hierarchical relationships, enriching the coarse feature map with basic information extracted from finer features through a top-down information path. After horizontal and vertical optimization, spatial information at each scale is merged and upsampled to obtain query features rich in contextual information; Step 2, prepare the dataset: preprocess the dataset according to the small sample segmentation task; Step 3: Train the network model: Input the data set prepared in step 2 into the network model built in step 1 for training; Step 4: Select the minimization loss function and optimal evaluation metric: By minimizing the loss function of the network output image and label until the number of training times reaches the set threshold or the value of the loss function falls within the set range, the model parameters are considered to have been pre-trained and saved. At the same time, the optimal evaluation metric is selected to measure the accuracy of the algorithm and evaluate the performance of the system. Step 5: Fine-tune and save the model: Adjust some parameters and train again until the network model achieves the optimal generalization ability and segmentation results under certain parameters.

2. The small sample semantic segmentation method based on query autocorrelation enhancement according to claim 1, characterized in that: The dataset described in step 2 is the small-sample benchmark dataset Pascal-5i. First, the dataset is divided into a training set and a test set, and then the training set and the test set are further divided into a support set and a query set, respectively. During training, k support samples containing category c are selected from the support set, and 1 query sample that also contains category c is selected from the query set. The goal is to learn a model that can segment category c in the query sample according to the paradigm of segmenting category c in the support sample. Each support-query combination is considered a training task.

3. The small sample semantic segmentation method based on query autocorrelation enhancement according to claim 1, characterized in that: Before inputting the dataset into step 1 as described in step 3, the images are preprocessed. This involves resizing all images to 473×473 pixels and removing images that do not contain any valid categories or that have valid objects that are too small, i.e., less than 2×32×32 pixels. This ensures that all images input to the model contain the specified categories and that the target objects are large enough.

4. The small sample semantic segmentation method based on query autocorrelation enhancement according to claim 1, characterized in that: As described in step 4, the Dice loss is selected as the loss function during the training process. The Dice loss directly measures the proportion of the overlapping area and can better guide the model to focus on the target area when samples are scarce.

5. The small sample semantic segmentation method based on query autocorrelation enhancement according to claim 1, characterized in that: As described in step 4, selecting mIOU and FB-IoU as evaluation indicators during the training process can effectively evaluate the accuracy and quality of the algorithm and measure the role of the segmentation network.

Images