Small sample semantic segmentation method based on query autocorrelation enhancement

Abstract

A small sample semantic segmentation method based on query autocorrelation enhancement belongs to the technical field of image segmentation, aims to solve the problems in the prior art, and comprises the following steps: step 1, constructing a network model: sharing a backbone network of weight; generating a prototype; the Transform decoder is used for inputting a support prototype and a query feature; the prior enhancement module is used for obtaining prior mask information with high guidance by utilizing interaction between high-level features; the feature enrichment module enriches query features from different scales by using a support prototype and a prior mask; filtering out a predicted mask from the query feature with an output of the decoder; 2, a data set is prepared, wherein the data set is preprocessed according to a small sample segmentation task; 3, training a network model: inputting the data set prepared in the step 2 into the network model constructed in the step 1 for training; 4, selecting a minimum loss function and an optimal evaluation index; and 5, finely adjusting and storing the model.

Description

Technical Field

[0001] The invention relates to a small sample semantic segmentation method based on query autocorrelation enhancement, belonging to the technical field of image segmentation. Background Art

[0002] Image semantic segmentation is one of the key issues in the field of computer vision, which aims to classify image data at the pixel level. Specifically, its goal is to convert an image into a mask with a highlighted area, and each pixel in the image is assigned a category according to the highlighted area. However, if you want to get a reliable segmentation model in the semantic segmentation task, you can only 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-world scenarios, small sample semantic segmentation came into being. Due to the small number of samples, there are difficult problems in small sample segmentation that cannot be ignored, that is, the intra-class differences are large. When the target category has different appearances between the query image and the support image, it is difficult for the model to associate the target category information on the support image with the target category information on the query image, resulting in poor prediction results.

[0003] The Chinese patent publication number is "CN117726809A", and the name is "A small sample semantic segmentation method based on information interaction enhancement". This method first extracts support features and query features through the backbone network, uses the guided attention module and the spatial information interaction attention module to process the features, and then uses the query prior mask to refine the query features. Finally, the refined query features are input into the decoder for processing to obtain the predicted segmentation results. However, this method has the problems of high computational complexity, blindly mining sample information, and being unable to effectively alleviate the appearance differences between the support samples and the query images. Summary of the invention

[0004] In order to solve the problems of high computational complexity, blind mining of sample information, and inability to effectively alleviate the appearance differences between support samples and query images in the prior art, the present invention proposes a small sample semantic segmentation method based on query autocorrelation enhancement. While keeping the number of parameters low, the method uses the autocorrelation between pixels within the query image to drive the migration of interactive information to the query data. While enhancing the utilization rate, it avoids the overfitting phenomenon caused by excessive mining of sample information. This method can achieve better segmentation effects under conditions of data scarcity and effectively improve the generalization ability of the model.

[0005] The solution to the technical problem of the present invention is:

[0006] A small sample semantic segmentation method based on query autocorrelation enhancement includes the following steps:

[0007] Step 1, build the network model: the whole 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 support prototypes 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 from different scales; the sixth part is to use the output of the decoder to filter out the predicted mask from the query feature;

[0008] Step 2, prepare the data set: preprocess the data set 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 built in step 1 for training;

[0010] Step 4, select the minimized loss function and the optimal evaluation index: By minimizing the loss function of the network output image and the label, until the number of training times reaches the set threshold or the value of the loss function reaches the set range, the model parameters are considered to have been pre-trained and saved; at the same time, the optimal evaluation index is selected to measure the accuracy of the algorithm and evaluate the performance of the system;

[0011] 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.

[0012] The backbone network described in step 1 uses an expanded version of Resnet, and the output is a dictionary of features in four dimensions. Each layer of features contains 1x1, 3x3 convolution operations, and ReLU operations. The prototype uses the masked average pooling method used in traditional research.

[0013] 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 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 a query sample containing category c is selected from the query set. The purpose 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.

[0014] Before inputting the dataset into step 1 as described in step 3, the images are preprocessed by resizing all images to 473×473 and removing images that do not contain any valid categories and that have valid objects that are too small, i.e., less than 2×32×32 pixels, to ensure that all images input to the model contain the specified categories and the target objects are large enough.

[0015] As described in step 4, the loss function during the training process selects Dice loss. Dice loss directly measures the proportion of overlapping areas and can better guide the model to focus on the target area when samples are scarce.

[0016] As described in step 4, the mean intersection over union (mIOU) and foreground-background intersection over union (FB-IoU) are selected as evaluation indicators during the training process, which can effectively evaluate the accuracy and quality of the algorithm and measure the role of the segmentation network.

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

[0018] 1. Use the middle-level features to generate prototypes, because the middle-level features constitute the target object part of the new class, which is more helpful to improve the generalization ability of the model. Using high-level features to generate enhanced priors can not only reduce the preference for base classes, but also achieve the acquisition of high-quality meta-knowledge, effectively improving the segmentation accuracy of the model.

[0019] 2. The purpose of using the designed enhanced prior module is to obtain more instructive prior information, which is then input into the FEM module to enrich the query features. This module can effectively improve the interaction between support and query, enhance the relevance between the two, and effectively alleviate the problem of appearance differences between the two in the same category.

[0020] 3. Choosing to use the Dice loss function is more advantageous than other loss functions. For example, cross entropy mainly measures the difference between the predicted probability distribution and the true distribution. In the case of severe class imbalance, it is easy to cause the model to ignore the minority class. The Dice loss directly optimizes the overlap of the target area, making the model more robust when predicting the minority class, and can better improve the segmentation quality of the few-sample class. BRIEF DESCRIPTION OF THE DRAWINGS

[0021] Figure 1 This is a network structure diagram of a small sample semantic segmentation method based on query autocorrelation enhancement in the present invention.

[0022] Figure 2 This is a structural diagram of the backbone network described in the present invention.

[0023] Figure 3 This is a structural diagram of the Transformer decoder described in the present invention.

[0024] Figure 4 These are the specific operation steps of the enhanced priori module described in the present invention. DETAILED DESCRIPTION

[0025] The present invention is further described in detail below with reference to the accompanying drawings.

[0026] like Figure 1 As shown, a small sample semantic segmentation method based on query autocorrelation enhancement includes the following steps:

[0027] Step 1, build the network model: the whole network consists of six parts. The first part is the backbone network with shared weights, which is used to extract supporting 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 supporting prototypes 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 supporting prototypes and prior masks to enrich query features from different scales; the sixth part uses the output of the decoder to filter out the predicted mask from the query features.

[0028] Step 2, prepare the data set: preprocess the data set 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 built in step 1 for training;

[0030] Step 4, select the minimized loss function and the optimal evaluation index: By minimizing the loss function of the network output image and the label, until the number of training times reaches the set threshold or the value of the loss function reaches the set range, the model parameters are considered to have been pre-trained and saved; at the same time, the optimal evaluation index is selected to measure the accuracy of the algorithm and evaluate the performance of the system;

[0031] 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.

[0032] Example:

[0033] The first part of step 1 is the backbone network structure with shared weights, such as Figure 2 As shown, the final output of the backbone network is a dictionary containing four layers of features, which are the outputs of the four residual blocks conv2_x, conv3_x, conv4_x, and conv5_x. The structure of each residual block is similar Figure 2The conv3_x shown on the right undergoes a 1×1 convolution, a 3×3 convolution, and a 1×1 convolution from input to output, and then the input is expanded to the same dimension as the output through a 1×1 convolution and added. 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 is defined as high-level features. We concatenate the mid-level features in the channel dimension, and then use a 1×1 convolution to reduce the number of channels to 60. Then, we use BatchNorm2d, ReLU, and Dropout operations to normalize each channel to prevent gradient explosion and gradient disappearance. Regularization is used to randomly output 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 input into the Transformer decoder as conditional queries. The high-level features conv5_x will be input into the enhanced prior module in the subsequent step.

[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 initial prototype generated. 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 residual connection is input to the self-attention module, 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 then 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 follows Figure 4 As shown. The high-level support features, query features, and support image masks are input into the enhanced prior module, and the three are used to generate initial auxiliary prior information. Specifically, the support mask is first used to filter out the background pixels of the support features so that the pixels of the query features have no corresponding relationship with the background on the support features. Then, the cosine similarity between the query features and each pixel of the support features is calculated. For each pixel of the query feature, a support pixel with the highest correspondence is selected, and the result is reshaped from h×w,1 to h,w,1, denoted as C q , according to formula (3), use the maximum-minimum normalization to make all values ​​in the range of 0-1, and the result Corr s represents the initial auxiliary prior mask, ε is set to 10 -7 . Then, the flattened high-level query features are used to filter out high-confidence feature information from each channel in the auxiliary prior, and the self-matching query vector is obtained by taking the average in this dimension. Since the target category pixels (foreground pixels) show high similarity and low noise tolerance, in order to ensure that there are enough pixels in the high-confidence area, the threshold τ = 0.7 is selected. After that, in order to fully exploit the segmentation prior and enhance the high interactivity between the support image and the query image, the support prototypes under the high channel features are taken respectively. Matching query vector FP pro Combined into a high-level class prototype, the similarity spectrum C between the class prototype and the high-level query feature prior As enhanced prior information, pixels with higher values ​​in the similarity spectrum mean that they are strongly correlated with pixels in the target object and are very likely to belong to the target category.

[0039] The fifth part, the feature enrichment module, inputs the mid-level query features, supporting prototypes, and enhanced prior information, maps them to different scales (sizes are 60×60, 30×30, 15×15, 8×8), concatenates the three inputs into new feature information, interacts independently between different scales, and vertically utilizes hierarchical relationships to enrich the coarse feature map with basic information extracted from finer features through a top-down information path. After horizontal and vertical optimization, the spatial information at each scale is merged and upsampled to obtain query features rich in contextual information.

[0040] The sixth part of step 1 is to embed the probability output of the Transformer decoder into the query feature F according to the following formula: q On the top, we get the final predicted query mask

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

[0042] Among them, ESC stands for Einstein summation operation. This process can be explained as the output of Transformer generates weights for features of 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 The final prediction mask is obtained by multiplying and summing pixel by pixel in the channel dimension and upsampling the result to restore local details.

[0043] The dataset in step 2 is the small sample benchmark dataset Pascal-5i. The dataset contains 12031 images. When obtaining support-query pairs from the dataset, all images are resized to 473×473, and images that do not contain any valid categories and have valid objects that are too small (less than 2×32×32 pixels) are deleted to ensure that all images input to the model contain the specified category and the target object is large enough.

[0044] In step 3, no other data enhancement strategies are used for training.

[0045] In step 4, the average of the results of 5 experiments with different random seeds is finally 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 of the IoU of each class in the current fold. IoU represents the ratio of the number of pixels in the intersection area between the predicted mask and the true mask of category C to the number of pixels in the union area. It is a global evaluation indicator. The implementation formula of IoU is as follows:

[0046]

[0047] Among them, TP C , FP C , FN C They represent the number of true positive, false positive, and false negative pixels for category C in the segmentation prediction results.

[0048] The FB-IoU indicator ignores the differences in image categories and only needs to calculate the IoU of the foreground and background, considering only the separation between the two. It has low computational complexity and is more sensitive to small sample evaluation.

[0049]

[0050] Among them, IoU FG and IoU BG Represent the foreground IoU and background IoU scores of the image respectively.

[0051] The loss function selects the weighted Dice loss function, which receives the model's prediction results and the target segmentation mask as input. The function expands the target mask by one channel dimension and converts it to floating point type, and then flattens the prediction results and target masks into one-dimensional vectors. Next, the function calculates the Dice coefficient by calculating the sum of the squares of the prediction results and the target mask, as well as the intersection between them. By subtracting the Dice coefficient, the Dice loss is obtained. To further adjust the loss value, the function weights the loss according to the set weights, and sums or averages the loss values ​​according to the configuration.

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

[0053] In step 5, the image dataset is used to fine-tune the parameters of the entire network model.

[0054] The present invention constructs a small sample semantic segmentation network based on query autocorrelation enhancement, which can complete more accurate binary segmentation under the condition of scarce samples, effectively utilizes the autocorrelation of the query image, mines out the pointing information that is more conducive to segmentation from the query image itself, enhances the high interactivity between the support image and the query image, provides favorable guidance for enriching the multi-scale spatial information of the query features, and effectively alleviates the problem of similar appearance differences between the support samples and the query samples. Under the same conditions, the feasibility and superiority of this method are further verified by calculating the relevant indicators of the image obtained by the existing method. The comparison of the relevant indicators of the prior art and the method proposed in the present invention is shown in Table 1:

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

[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 present invention 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 in that: 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 supporting prototypes and prior masks to enrich query features from different scales. The sixth part is to use the decoder output to filter out the predicted mask from the query features; Step 2, prepare the data set: preprocess the data set according to the small sample segmentation task; Step 3, training 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 minimized loss function and the optimal evaluation index: By minimizing the loss function of the network output image and the label, until the number of training times reaches the set threshold or the value of the loss function reaches the set range, the model parameters are considered to have been pre-trained and saved; at the same time, the optimal evaluation index 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. A small sample semantic segmentation method based on query autocorrelation enhancement, characterized in that: The backbone network described in step 1 uses an expanded version of Resnet, and the output is a dictionary of features in four dimensions. Each layer of features contains 1x1, 3x3 convolution operations, and ReLU operations. The prototype uses the masked average pooling method used in traditional research.

3. A small sample semantic segmentation method based on query autocorrelation enhancement, 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 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 purpose 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 taken as a training task.

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

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

6. A small sample semantic segmentation method based on query autocorrelation enhancement, characterized in that: Selecting mIOU and FB-IoU as evaluation indicators during the training process as described in step 4 can effectively evaluate the accuracy and quality of the algorithm and measure the role of the segmentation network.

Images