Small sample image semantic segmentation method based on prototype network

Abstract

The invention discloses a small sample image semantic segmentation method based on a prototype network, and belongs to the technical field of image semantic segmentation. The problem of excessive image space information loss caused by mask average pooling operation of a traditional small sample image semantic segmentation method in the prior art is solved. According to the method, features of an input support image and an input query image are extracted through a pre-trained neural network, support image features and query image features are obtained, picture features are extracted from a specified small sample image semantic segmentation data set, and generalized category labels are generated through a KMeans algorithm; extracting a generalized prototype vector, and enhancing the query image features through the generalized prototype vector to obtain final enhanced features; and introducing a segmentation model, inputting the final enhanced features for decoding, generating a feature image, inputting the feature image into a classifier, and obtaining a segmentation result. The small sample image semantic segmentation method effectively improves the accuracy of small sample image semantic segmentation, and can be applied to small sample image semantic segmentation.

Description

Technical Field

[0001] The present invention relates to a small sample image semantic segmentation method, in particular to a small sample image semantic segmentation method based on a prototype network, and belongs to the technical field of image semantic segmentation. Background Art

[0002] The semantic segmentation task itself refers to classifying each pixel in an image and dividing the image into regions with specific semantic meanings. However, in many practical applications, such as medical image analysis and robot vision, obtaining a large amount of accurately labeled data is time-consuming and expensive.

[0003] Semantic segmentation of small-sample images is an important topic in the field of computer vision. It aims to overcome the dependence of traditional semantic segmentation tasks on a large amount of labeled data. The goal of semantic segmentation of small-sample images is to achieve efficient image segmentation using a very small number of labeled samples, which can well deal with the problem of small training samples. Some researchers have proposed OSLSM, a model for labeling and segmenting specific objects, which uses a dual-branch structure containing support branches and query branches to transfer annotation information and interact between extracted features. At the same time, influenced by the prototypical network (PrototypicalNet) in few-shot learning (Few-Shot Learning), the OSLSM algorithm can be used to implement semantic segmentation of small-sample images. Inspired by Deep Learning, many previous works use masked average pooling operations to obtain class prototypes from support information to represent support features, and use prototypes to represent specific categories; Probabilistic Mixture Models (PMMs) generate multiple prototypes for the same class, PFENet is a deep learning model for small sample learning, which creates a priori masks on high-level features without trainable parameters, and then interacts the prior masks with cross-scale information to enrich query features. Compared with prototype-based methods, matching-based methods aim to extract correspondences between query images and support annotations, and use pixel-level features to enhance complex details in support contexts; CyCTR is a learning method for video understanding and temporal data analysis, It introduces a recurring consistent attention mechanism to equip query features with relevant support information; BAM is a mechanism for enhanced feature learning, which introduces a new approach for image segmentation in small sample cases (FSS), using additional branches of the supervised model trained on the base class to clearly identify base class targets (i.e., areas that do not need to be segmented), and the supervised model predicts the base class from the query image and helps the meta-learner suppress erroneous predictions; the above technical means have jointly promoted the development of the field of semantic segmentation of small sample images, and improved the robustness and accuracy of the algorithm in the case of limited data, but the use of masked average pooling operations can enrich features and improve model performance, but it leads to excessive loss of image spatial information, resulting in large errors in the output results.

[0004] In summary, a small sample image semantic segmentation method based on prototype networks is needed. Summary of the invention

[0005] A brief overview of the present invention is provided below in order to provide a basic understanding of certain aspects of the present invention. It should be understood that this overview is not an exhaustive overview of the present invention. It is not intended to identify key or important parts of the present invention, nor is it intended to limit the scope of the present invention. Its purpose is merely to present certain concepts in a simplified form as a prelude to a more detailed description discussed later.

[0006] In view of this, in order to solve the problem of excessive image spatial information loss caused by mask average pooling operation in traditional small sample image semantic segmentation methods in the prior art, the present invention provides a small sample image semantic segmentation method based on a prototype network.

[0007] The technical solution is as follows: A small sample image semantic segmentation method based on a prototype network includes the following steps:

[0008] S1. extracting features of the input support image and query image through a pre-trained neural network to obtain support image features and query image features;

[0009] S2. Use a pre-trained neural network to extract image features from a specified small sample image semantic segmentation dataset, and generate generalized category labels using the KMeans algorithm based on the image features;

[0010] S3. construct a generalized prototype enhancement module, input supporting image features and generalized category labels, extract generalized prototype vectors, enhance query image features through generalized prototype vectors, and obtain the final enhanced features;

[0011] S4. Input the final enhanced features into the segmentation model, decode the final enhanced features through the decoder of the segmentation model, generate a feature image and input it into the classifier of the segmentation model to obtain the segmentation result.

[0012] Furthermore, in S1, the support image is the input support set The query image is I q ,in, h is the image height, w is the image width, M s is the binary segmentation mask, M s ∈{0,1} h×w , K represents the number of samples used for training in each category. The pre-trained neural network is used as the backbone network to extract the corresponding features of the support image and the query image, which are respectively recorded as support image features F s and query image features F q ,in, c represents the number of channels of the feature map extracted by the pre-trained neural network, h′ is the height of the feature map, and w′ is the width of the feature map.

[0013] Furthermore, in S2, for a specified small sample image semantic segmentation dataset, there are m categories, and the original category label set is defined as C, where C = {c1, c2, c3, ..., c m}, i = 1, 2, ..., m, for each original category c i , c i ∈C, the image set corresponding to the original category label set C is D i , randomly select a subset D′ of the image set i , The subset size is denoted by |D′ i |=n, for each image x ij , x ij ∈D′ i , use the pre-trained neural network to extract image features F ij , Among them, j = 1, 2, ..., n, h' is the height of the feature map, w' is the width of the feature map, and the image feature F ij Flattened image features suitable for KMeans algorithm According to the KMeans algorithm, the flattened image features Cluster into K clusters and find the cluster center and cluster labels l ij , l ij ∈{1,2,...,K};

[0014] Cluster label l ij It is expressed as:

[0015]

[0016] Among them, argmin means to make the function reach the minimum value;

[0017] For each original category c i In the clustering cluster labels corresponding to all feature vectors under the category, find the clustering cluster label with the highest frequency, and define the clustering cluster label with the highest frequency as the original category c i The generalized category label of .

[0018] Furthermore, the step S3 includes the following steps:

[0019] S31. Extracting a generalized prototype vector through a generalized prototype enhancement module according to supporting image features and generalized category labels;

[0020] S32. Enhance the query image features through the generalized prototype vector, and enhance the enhanced features F′ q Processing is performed to obtain the final enhanced features;

[0021] In S31, a generalized class prototype library P is defined, where P = {P1, ..., P t ,...,P num}, where num represents the size of the generalized class prototype library, t is the broad category index, and the binary segmentation mask M is adjusted by bilinear interpolation. s The size of the binary segmentation mask M is converted using spatial multiplication s Applied to the support features, the activated regions are averaged to calculate the initial prototype, and the initial prototype P is obtained. o ;

[0022] Initial Prototype P o It is expressed as:

[0023]

[0024] Among them, |M s | represents the binary segmentation mask M s The total number of elements with a median value of 1, x is the index of the width w of the traversed image or feature map, and y is the index of the height h of the traversed image or feature map;

[0025] In S32, the generalized prototype vector P in the generalized class prototype library P is updated according to the generalized class label. t , and get the enhanced feature F′ q ;

[0026] The update process is expressed as:

[0027] P t =(1-γ)l2(p o )+γl2(p t )

[0028]

[0029] Among them, γ is the adjustment factor for balancing the introduction of old information and new information, f(·) is a 1×1 convolution, Indicates the splicing operation along the channel dimension, represents broadcast multiplication, l2 represents L2 normalization;

[0030] For the enhanced feature F′ q Perform self-attention operation, the enhanced feature F′ q The final enhanced features are obtained by compressing through the down-projection linear layer, using the GeLU activation function to increase nonlinearity, and finally restoring the feature dimension through the up-projection linear layer.

[0031] The self-attention operation process is expressed as:

[0032]

[0033] Where [Q; K; V] = [W q F′ q ; W k F′ q ; W v F′ q ], is the learnable parameter matrix, d is the dimension;

[0034] Final enhanced features It is expressed as:

[0035]

[0036] Among them, W down is the down-projection linear layer, W up is the up-projection linear layer, α represents the scaling ratio.

[0037] The beneficial effects of the present invention are as follows: the present invention comprises a pre-trained neural network, a self-built generalized prototype enhancement module and a segmentation model to form a small sample image semantic segmentation model based on a prototype network. The small sample image semantic segmentation model can achieve high-precision image semantic segmentation while ensuring that information is not lost. The present invention is mainly divided into four steps: feature extraction, generation of generalized category labels, prototype enhancement and feature segmentation through a decoder. The generalized category labels are generated through the KMeans algorithm, and the generalized prototypes are updated and matched in real time through the prototype enhancement module, thereby enhancing the features of the query image and further improving the accuracy of the segmentation model. BRIEF DESCRIPTION OF THE DRAWINGS

[0038] The drawings described herein are used to provide a further understanding of the present invention and constitute a part of the present invention. The exemplary embodiments of the present invention and their descriptions are used to explain the present invention and do not constitute an improper limitation of the present invention. In the drawings:

[0039] Figure 1 It is a flowchart of a small sample image semantic segmentation method based on prototype network;

[0040] Figure 2 Schematic diagram of the implementation framework of a small sample image semantic segmentation model based on a prototype network. DETAILED DESCRIPTION

[0041] In order to make the technical solutions and advantages of the embodiments of the present invention more clearly understood, the exemplary embodiments of the present invention are further described in detail below in conjunction with the accompanying drawings. Obviously, the described embodiments are only part of the embodiments of the present invention, rather than an exhaustive list of all the embodiments. It should be noted that the embodiments of the present invention and the features in the embodiments can be combined with each other without conflict.

[0042] refer to Figure 1-Figure 2 The present embodiment is described in detail, a small sample image semantic segmentation method based on a prototype network, specifically comprising the following steps:

[0043] S1. extracting features of the input support image and query image through a pre-trained neural network to obtain support image features and query image features;

[0044] S2. Use a pre-trained neural network to extract image features from a specified small sample image semantic segmentation dataset, and generate generalized category labels using the KMeans algorithm based on the image features;

[0045] S3. construct a generalized prototype enhancement module, input supporting image features and generalized category labels, extract generalized prototype vectors, enhance query image features through generalized prototype vectors, and obtain the final enhanced features;

[0046] S4. Input the final enhanced features into the segmentation model, decode the final enhanced features through the decoder of the segmentation model, generate a feature image and input it into the classifier of the segmentation model to obtain a segmentation result;

[0047] Specifically, refer to Figure 2 , a small sample image semantic segmentation model based on prototype network is composed of pre-trained neural network, generalized prototype enhancement module and segmentation model, conv represents pre-trained neural network, MAP represents mask average pooling, Cat represents concatenation, FCL represents fully connected layer, and FEM represents feature enhancement module;

[0048] In step S4, the segmentation features are decoded by the decoder of the segmentation model, and the correlation graph distillation technology is used to enable the shallow network to approximate the semantic relevance of the deep network, thereby improving context awareness and enhancing segmentation quality. The feature image generated by the decoder is then input into the classifier in the segmentation model to obtain the segmentation result of the labeled category.

[0049] Furthermore, in S1, the support image is the input support set The query image is I q ,in, h is the image height, w is the image width, M s is the binary segmentation mask, M s ∈{0,1}h×w , K represents the number of samples used for training in each category. The pre-trained neural network (ResNet-50) is used as the backbone network to extract the corresponding features of the support image and the query image, which are respectively recorded as support image features F s and query image features F q ,in, c represents the number of channels of the feature map extracted by the pre-trained neural network, h′ is the height of the feature map, and w′ is the width of the feature map. ResNet-50 freezes parameters during model training, which improves the generalization of the model to unseen categories.

[0050] Furthermore, in S2, the small sample image semantic segmentation dataset adopts the Pascal-5i dataset, which has 20 categories, m=20, and the original category label set is defined as C, C={c1,c2,c3,...,c m}, i = 1, 2, ..., m, for each original category c i , c i ∈C, the image set corresponding to the original category label set C is D i , randomly select a subset D′ of the image set i , The subset size is denoted by |D′ i |=n, n=100, for each image x ij , x ij ∈D′ i , use pre-trained neural network to extract image features Among them, j = 1, 2, ..., n, h' is the height of the feature map, w' is the width of the feature map, and the image feature F ij Flattened image features suitable for KMeans algorithm According to the KMeans algorithm, the flattened image features Cluster into K clusters and find the cluster center and cluster labels l ij , l ij ∈{1,2,…,K};

[0051] Cluster label l ij It is expressed as:

[0052]

[0053] Among them, argmin means to make the function reach the minimum value;

[0054] For each original category c iIn the clustering cluster labels corresponding to all feature vectors under the category, find the clustering cluster label with the highest frequency, and define the clustering cluster label with the highest frequency as the original category c i The generalized category label of .

[0055] Furthermore, the step S3 includes the following steps:

[0056] S31. Extracting a generalized prototype vector through a generalized prototype enhancement module according to supporting image features and generalized category labels;

[0057] S32. Enhance the query image features through the generalized prototype vector, and enhance the enhanced features F′ q Processing is performed to obtain the final enhanced features;

[0058] In S31, a generalized class prototype library P is defined, where P = {P1, ..., P t ,...,P num}, where num represents the size of the generalized class prototype library, in the present invention, num=6, t is the broad category index, and the binary segmentation mask M is adjusted by bilinear interpolation s Size, M s ∈{0,1} h×w , use spatial multiplication to convert the binary segmentation mask M s Applied to the support features, the activated regions are averaged to calculate the initial prototype, and the initial prototype P is obtained. o ;

[0059] Initial Prototype P o It is expressed as:

[0060]

[0061] Among them, |M s | represents the binary segmentation mask M s The total number of elements with a median value of 1, x is the index of the width w of the traversed image or feature map, y is the index of the height h of the traversed image or feature map, and the above index is used to specify each pixel on the mask image or feature map;

[0062] In S32, the generalized prototype vector P in the generalized class prototype library P is updated according to the generalized class label. i ;

[0063] The update process is expressed as:

[0064] P t =(1-γ)l2(p o )+γl2(p t )

[0065]

[0066] Among them, γ is a regulatory factor that balances the introduction of old information and new information, which is a trainable parameter, f(·) is a 1×1 convolution, Indicates the splicing operation along the channel dimension, represents broadcast multiplication so that the generalized prototype P i Can be extended to F q The same spatial size, l2 represents L2 normalization;

[0067] Furthermore, in the testing phase of the small sample image semantic segmentation model, by traversing the support image prototype P o and the prototype vector P in the generalized prototype library k The similarity between them selects the target class prototype and enhances the query feature to obtain the enhanced feature F′ q ;

[0068] Enhanced feature F′ q It is expressed as:

[0069]

[0070]

[0071] Among them, P o represents the initial prototype, P b Represents the initial prototype P o The generalized prototype vector with the largest cosine similarity, T is the matrix transpose;

[0072] For the enhanced feature F′ q Perform self-attention operation, the enhanced feature F′ q Through the downward projection linear layer compression, the compressed feature dimension uses the GeLU activation function to increase nonlinearity to capture the nonlinear relationship between data. Finally, the feature dimension is restored through the upward projection linear layer to capture the nonlinear relationship between data and obtain the final enhanced feature.

[0073] The self-attention operation process is expressed as:

[0074]

[0075] Where [Q; K; V] = [W q F′ q ; W k F′ q ; W v F′ q ], is the learnable parameter matrix, d is the dimension, d=8;

[0076] Final enhanced features It is expressed as:

[0077]

[0078] Among them, W down is the down-projection linear layer, W up is the up-projection linear layer, α represents the scaling ratio.

[0079] Specifically, the generalized prototype vector P i It helps the segmentation model to generalize better when dealing with new, unseen images. After extracting the prototypes, it expands the single prototype to the same size as the query feature and concatenates the two together.

[0080] The self-attention operation aims to further refine the feature representation by emphasizing relevant information and suppressing irrelevant details, thereby enhancing the segmentation model's ability to process key information in the input data;

[0081] Subsequent operations such as down-projection linear layer compression, adding nonlinearity, and up-projection linear layer restoration of feature dimensions can reduce the risk of overfitting.

[0082] Although the present invention has been described according to a limited number of embodiments, it will be apparent to those skilled in the art, with the benefit of the above description, that other embodiments may be envisioned within the scope of the invention thus described. In addition, it should be noted that the language used in this specification is selected primarily for readability and teaching purposes, rather than for explaining or defining the subject matter of the present invention. Therefore, many modifications and variations will be apparent to those skilled in the art without departing from the scope and spirit of the appended claims. The disclosure of the present invention is illustrative, not restrictive, with respect to the scope of the present invention, which is defined by the appended claims.

Claims

1. A small sample image semantic segmentation method based on prototype network, characterized in that: The following steps are involved: S1. extracting features of the input support image and query image through a pre-trained neural network to obtain support image features and query image features; S2. Use a pre-trained neural network to extract image features from a specified small sample image semantic segmentation dataset, and generate generalized category labels using the KMeans algorithm based on the image features; S3. construct a generalized prototype enhancement module, input supporting image features and generalized category labels, extract generalized prototype vectors, enhance query image features through generalized prototype vectors, and obtain the final enhanced features; S4. Input the final enhanced features into the segmentation model, decode the final enhanced features through the decoder of the segmentation model, generate a feature image and input it into the classifier of the segmentation model to obtain the segmentation result.

2. According to the method of semantic segmentation of small sample images based on prototype network in claim 1, it is characterized in that: In S1, the support image is the input support set The query image is I q ,in, h is the image height, w is the image width, M s is the binary segmentation mask, M s ∈{0,1} h×w , K represents the number of samples used for training in each category. The pre-trained neural network is used as the backbone network to extract the corresponding features of the support image and the query image, which are respectively recorded as support image features F s and query image features F q ,in, c represents the number of channels of the feature map extracted by the pre-trained neural network, h′ is the height of the feature map, and w′ is the width of the feature map.

3. The method for semantic segmentation of small sample images based on prototype network according to claim 2, characterized in that: In S2, for the specified small sample image semantic segmentation dataset, there are m categories, and the original category label set is defined as C, C = {c1, c2, c3, ..., c m }, i = 1, 2, ..., m, for each original category c i , c i ∈C, the image set corresponding to the original category label set C is D i , randomly select a subset D of the image set i ′, The subset size is represented by |D i ′|=n, for each image x ij , x ij ∈D i ′, use the pre-trained neural network to extract image features F ij , Among them, j = 1, 2, ..., n, h' is the height of the feature map, w' is the width of the feature map, and the image feature F ij Flattened image features suitable for KMeans algorithm According to the KMeans algorithm, the flattened image features Cluster into K clusters and find the cluster center and cluster labels l ij , l ij ∈{1,2,…,K}; Cluster label l ij It is expressed as: Among them, argmin means to make the function reach the minimum value; For each original category c i In the clustering cluster labels corresponding to all feature vectors under the category, find the clustering cluster label with the highest frequency, and define the clustering cluster label with the highest frequency as the original category c i The generalized category label of .

4. The method for semantic segmentation of small sample images based on prototype network according to claim 3, characterized in that: The S3 comprises the following steps: S31. Extracting a generalized prototype vector through a generalized prototype enhancement module according to supporting image features and generalized category labels; S32. Enhance the query image features through the generalized prototype vector, and enhance the enhanced features F q ' is processed to obtain the final enhanced features; In S31, a generalized class prototype library P is defined, where P = {P1, ..., P t ,...,P num }, where num represents the size of the generalized class prototype library, t is the broad category index, and the binary segmentation mask M is adjusted by bilinear interpolation. s The size of the binary segmentation mask M is converted using spatial multiplication s Applied to the support features, the activated regions are averaged to calculate the initial prototype, and the initial prototype P is obtained. o ; Initial Prototype P o It is expressed as: Among them, |M s | represents the binary segmentation mask M s The total number of elements with a median value of 1, x is the index of the width w of the traversed image or feature map, and y is the index of the height h of the traversed image or feature map; In S32, the generalized prototype vector P in the generalized class prototype library P is updated according to the generalized class label. t , and get the enhanced feature F q ′; The update process is expressed as: P.S t S(1-γ)l2(p o )+γl2(p t ) Among them, γ is the adjustment factor for balancing the introduction of old information and new information, f(·) is a 1×1 convolution, Indicates the splicing operation along the channel dimension, represents broadcast multiplication, l2 represents L2 normalization; For the enhanced feature F′ q Perform self-attention operation, the enhanced feature F′ q The final enhanced features are obtained by compressing through the down-projection linear layer, using the GeLU activation function to increase nonlinearity, and finally restoring the feature dimension through the up-projection linear layer. The self-attention operation process is expressed as: Where [Q; K; V] = [W q F q ′;W k F q ′;W v F q ′], is the learnable parameter matrix, d is the dimension; Final enhanced features It is expressed as: Among them, W down is the down-projection linear layer, W up is the up-projection linear layer, α represents the scaling ratio.

Images