Small sample medical image segmentation method and system based on query feature enhancement and image attention mechanism

Abstract

The invention discloses a small sample medical image segmentation method and system based on query feature enhancement and a graph attention mechanism. The method is realized by designing a query feature enhancement module and a multi-prototype collaborative segmentation module. The query feature enhancement module provides comprehensive and highly distinguished query features for subsequent prototype generation and matching by capturing a global long-distance dependency relationship, aggregating local corner regions and integrating support sample information. The multi-prototype collaborative segmentation module is composed of two prediction branches. The branch I is used for matching query prototypes and query features which belong to the same object, and the branch II is used for promoting aggregation and mutual enhancement among a plurality of prototypes by adopting a graph attention mechanism, and dynamically allocating the most relevant prototype for each pixel in a query image. And more comprehensive segmentation prediction is realized by fusing prediction results of the two branches. According to the method, challenges caused by anatomical structure heterogeneity are effectively handled, and the accuracy and adaptability of segmentation are improved.

Description

Technical Field

[0001] The present invention relates to the field of medical image processing technology, and specifically to a small sample medical image segmentation method and system based on query feature enhancement and graph attention mechanism. Background Art

[0002] Medical image segmentation is an important step in medical image analysis and processing, and is one of the prerequisites for effective diagnosis and treatment of diseases. With the development of deep learning, deep learning models with full supervision training for specific organs or tissues have achieved good results in many medical image segmentation tasks. However, due to the high cost of medical image annotation, patient privacy issues and the scarcity of samples due to some rare diseases, it is very expensive and difficult to obtain a large amount of labeled data required for full supervision training. In recent years, many small sample image segmentation methods have been proposed to effectively alleviate the burden of data annotation. Such methods can train models with good generalization ability, so that they can quickly adapt to new categories and achieve segmentation prediction under the guidance of a small number of labeled samples. In particular, the prototype learning-based method achieves pixel-level classification by constructing a prototype for each category and calculating the similarity between the query image pixel features and the prototype. This method uses limited samples to extract representative features of the category, showing excellent generalization ability and efficiency, and has therefore attracted much attention.

[0003] However, biological individual differences, dynamic changes in disease states, and differences in imaging technology and equipment configurations lead to significant heterogeneity in anatomical structures between and within individuals. For example, the hearts of different individuals may differ in size, shape, and position; different regions of the liver of the same individual may show different vascular distribution and tissue density. This heterogeneity poses an inherent challenge to prototype learning-based small-sample medical image segmentation methods, increasing the difficulty of extracting features from limited samples and generalizing to new samples. Specifically, this challenge is mainly reflected in two aspects: 1) The inability to generate representative prototypes may lead to insufficient ability to accurately capture complex anatomical structures; 2) Insufficient capture of local details during the matching process between prototypes and query features may weaken the model's ability to adapt to individual differences. Therefore, how to generate comprehensive and representative prototypes and ensure accurate matching with query features has become a core issue in the study of small-sample medical image segmentation based on prototype learning.

[0004] Some methods construct prototypes for the foreground region (target segmentation region) and background region (region outside the target segmentation region) of the support image separately. However, due to the limited number of support image slices, modeling the features of complex background regions is often insufficient. Other methods construct a single prototype only for the foreground region of the support image through mask average pooling. This method is effective in some scenarios (such as when the foreground category is relatively uniform and does not contain multiple distinct regions with clear boundaries). However, this method may lead to semantic ambiguity and information loss when dealing with complex foreground regions. For example, when the left ventricular blood pool and left ventricular myocardium are considered as a single foreground category (left ventricle). In addition, some other methods construct multiple prototypes for the foreground region, but merge these prototypes into a single prototype through global average pooling before calculating the similarity with the query feature. This process may lead to the loss of unique information and details in each prototype, making it difficult for the model to adapt to individual differences when representing foreground regions with highly heterogeneous structures. Overall, none of these methods effectively address the challenges caused by anatomical heterogeneity.

[0005] In addition, current prototype learning-based small-sample medical image segmentation methods generally use pre-trained ResNet networks as feature extractors. However, since ResNet is designed to focus more on general feature extraction, it has certain limitations in capturing global contextual information and local details when processing medical image tasks. At the same time, the key information in the support samples is not fully utilized, resulting in a weak correlation between the query features and the support sample features, and the support samples cannot fully play their guiding role in the segmentation task. Some methods attempt to introduce self-attention mechanisms and cross-attention mechanisms to capture the correlation between the support image and the query image, so that the model focuses on more valuable foreground information. However, the use of traditional cross-attention mechanisms may lead to mismatches between the background features of the query image and the foreground features of the support image, thereby affecting the segmentation performance of the model. Summary of the invention

[0006] In order to overcome the shortcomings of the above-mentioned prior art, the present invention provides a small sample medical image segmentation method and system based on query feature enhancement and graph attention mechanism, which realizes accurate segmentation of medical images through the mutual cooperation of query feature enhancement and multi-prototype collaborative segmentation.

[0007] According to one aspect of the present invention, a small sample medical image segmentation method based on query feature enhancement and graph attention mechanism is provided, comprising:

[0008] Acquire a medical image to be tested;

[0009] Input the medical image to be tested into the trained segmentation model and output the segmentation result; wherein the training of the segmentation model includes:

[0010] Step 1, obtain the training data set;

[0011] Step 2: Use the pre-trained ResNet101 network as a feature extractor to extract features from the support image and the query image to obtain support features and query features.

[0012] Step 3, using the designed query feature enhancement module to enhance the extracted query features;

[0013] Step 4, using the designed multi-prototype collaborative segmentation module to perform segmentation prediction, the multi-prototype collaborative segmentation module uses the query prototype to match the enhanced query features to generate a query prospect prediction; and, combining the single prospect prototype, the support prototype set generated based on k-means clustering, and the query prototype to construct a multi-prototype set, and enhance its representation ability through the graph attention mechanism, and then performing fine-grained matching with the query features to generate a multi-prototype prospect prediction; the generated query prospect prediction is fused with the multi-prototype prospect prediction to obtain the final segmentation prediction;

[0014] Step 5: Construct a loss function for model training and output the trained segmentation model.

[0015] As a further technical solution, a designed query feature enhancement module is used to enhance the extracted query features, including:

[0016] By using the displayed visual center method, the global long-distance dependencies in the query features are captured by a lightweight multi-layer perceptron, and the local angular regions are aggregated by a learnable visual center to integrate the global long-distance dependencies and the local angular regions to obtain enhanced query features.

[0017] As a further technical solution, after obtaining the enhanced query features, it also includes:

[0018] Self-calibrated criss-cross attention is adopted to align query features with foreground regions of supporting features, generating comprehensively enhanced query features.

[0019] As a further technical solution, query prototypes are matched with enhanced query features to generate query prospect predictions, including:

[0020] Using the support features and support image masks provided by the support set, the support prototype is obtained using the mask average pooling operation;

[0021] Calculate the cosine similarity between the supporting prototype and the enhanced query features, and select the prediction result with high confidence as the query prediction;

[0022] Aggregate query features using high-confidence query predictions and combine with supporting prototypes to obtain query prototypes;

[0023] The part of the support prototype used in the above query prototype generation process is replaced with the query prototype, and the query prototype generation process is repeated twice to obtain a more accurate query prototype;

[0024] A foreground prototype portion of the more accurate query prototype is extracted, and a soft threshold prediction method is used to obtain a foreground prediction of the query prototype.

[0025] As a further technical solution, a single foreground prototype, a support prototype set generated based on k-means clustering, and a query prototype are combined to construct a multi-prototype set, and its representation ability is enhanced through the graph attention mechanism. Then, fine-grained matching is performed with the query features to generate multi-prototype foreground predictions, including:

[0026] A single foreground prototype is extracted from the obtained supporting prototypes, and a set of supporting prototypes is constructed for the foreground region of the supporting image using k-means clustering;

[0027] The single foreground prototype, the supporting prototype set and the query prototype are spliced ​​together to assemble a multi-prototype set;

[0028] A single-layer graph attention network is introduced to enhance the representation ability of each prototype in the multi-prototype set, and the enhanced multi-prototype set is obtained;

[0029] Based on the enhanced multi-prototype set, fine-grained matching is performed with query features to generate multi-prototype prospect predictions.

[0030] As a further technical solution, a loss function is constructed to train the model and output the trained segmentation model, including:

[0031] The foreground prediction and background prediction are upsampled to the set image size, and the weighted sum of the binary cross entropy loss and the prototype alignment regularization loss is used as the total loss function;

[0032] The total loss function is used to perform model training and output a trained segmentation model.

[0033] According to one aspect of the present invention, a small sample medical image segmentation system based on query feature enhancement and graph attention mechanism is provided, comprising:

[0034] An image data module, used for acquiring medical images to be tested;

[0035] An image segmentation module is used to input the medical image to be tested into a trained segmentation model and output a segmentation result; wherein the training of the segmentation model includes:

[0036] Step 1, obtain the training data set;

[0037] Step 2: Use the pre-trained ResNet101 network as a feature extractor to extract features from the support image and the query image to obtain support features and query features.

[0038] Step 3, using the designed query feature enhancement module to enhance the extracted query features;

[0039] Step 4, using the designed multi-prototype collaborative segmentation module to perform segmentation prediction, the multi-prototype collaborative segmentation module uses the query prototype to match the enhanced query features to generate a query prospect prediction; and, combining the single prospect prototype, the support prototype set generated based on k-means clustering, and the query prototype to construct a multi-prototype set, and enhance its representation ability through the graph attention mechanism, and then performing fine-grained matching with the query features to generate a multi-prototype prospect prediction; the generated query prospect prediction is fused with the multi-prototype prospect prediction to obtain the final segmentation prediction;

[0040] Step 5: Construct a loss function for model training and output the trained segmentation model.

[0041] According to one aspect of the present invention specification, a small sample medical image segmentation device based on query feature enhancement and graph attention mechanism is provided, comprising a memory and a processor, wherein the memory stores program instructions executed by the processor, and the processor calls the program instructions to execute the steps of the small sample medical image segmentation method based on query feature enhancement and graph attention mechanism.

[0042] According to one aspect of the present specification, a non-transitory computer-readable storage medium is provided, wherein the non-transitory computer-readable storage medium stores computer instructions, wherein the computer instructions enable the computer to execute the steps of the small sample medical image segmentation method based on query feature enhancement and graph attention mechanism.

[0043] Compared with the prior art, the present invention has the following beneficial effects:

[0044] 1. The present invention improves the feature representation capability through a query feature enhancement module, and constructs a comprehensive and representative prototype through a multi-prototype collaborative segmentation module, thereby accurately matching query features, thereby effectively improving segmentation accuracy.

[0045] 2. The multi-prototype collaborative segmentation module of the present invention enhances the prototype's ability to represent complex anatomical structures through the graph attention mechanism, making the model better adapted to different anatomical regions. At the same time, the fine-grained feature matching mechanism dynamically allocates prototypes according to individual anatomical differences, further improving the model's adaptability to different individuals.

[0046] 3. The query feature enhancement module of the present invention enhances the perception of global context and local details by capturing global long-distance dependencies and aggregating local angular areas, and effectively filters noise and background interference by aligning the foreground area of ​​the query image with the support image, thereby improving the robustness of the model in noisy, blurred or missing images. BRIEF DESCRIPTION OF THE DRAWINGS

[0047] In order to more clearly illustrate the embodiments of the present invention or the technical solutions in the prior art, a brief introduction is given below to the drawings used in the embodiments or the description of the prior art. Obviously, the drawings described below are some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying any creative work.

[0048] Figure 1 A schematic diagram of a process flow of a small sample medical image segmentation method based on query feature enhancement and graph attention mechanism provided by an embodiment of the present invention;

[0049] Figure 2 (a)-(b) are schematic diagrams of the architecture of the segmentation model and its multi-prototype collaborative segmentation module provided by an embodiment of the present invention, wherein (a) is a schematic diagram of the architecture of the segmentation model, and (b) is a schematic diagram of the architecture of the multi-prototype collaborative segmentation module;

[0050] Figure 3 A schematic diagram of the architecture of a query feature enhancement module provided for an example of the present invention. DETAILED DESCRIPTION

[0051] The present invention aims to solve the problems in existing small sample medical image segmentation methods, such as the lack of representativeness in prototype construction and insufficient capture of local details in the process of matching with query features, which leads to the inability to effectively deal with the heterogeneity of anatomical structures, thereby affecting the accuracy and adaptability of segmentation. In addition, the present invention also solves the problems of existing feature extraction methods in global and local information processing and supporting sample information utilization, thereby improving the accuracy and robustness of segmentation.

[0052] In order to solve the above technical problems, the present invention discloses a small sample medical image segmentation method based on query feature enhancement and graph attention mechanism, which is implemented by designing a query feature enhancement module and a multi-prototype collaborative segmentation module.

[0053] The present invention designs a multi-prototype collaborative segmentation module, which consists of two prediction branches. Among them, branch one effectively alleviates the anatomical differences between support samples and query samples by matching the query prototypes belonging to the same object with the query features. Branch two uses a graph attention mechanism to promote the aggregation and mutual enhancement of multiple prototypes, thereby improving the prototype's ability to represent the complex anatomical structure of the target area. Subsequently, fine-grained feature matching is achieved by dynamically assigning the most relevant prototype to each pixel in the query image. This matching strategy can be flexibly adjusted according to individual differences to better adapt to the variability of different anatomical structures. Finally, by fusing the prediction results of the two branches to achieve a more comprehensive segmentation prediction, the challenges brought by the heterogeneity of anatomical structures are effectively addressed.

[0054] In addition, the present invention designs a query feature enhancement module, which strengthens the representation ability of the query feature by capturing global long-distance dependencies and aggregating local angular areas. At the same time, the introduction of self-calibrated cross attention effectively aligns the query feature with the foreground area of ​​the supporting feature, fully utilizes the target category prior information provided by the supporting feature and its mask, and guides the query feature to focus on the more valuable foreground area, thereby filtering noise and background interference and improving the robustness of the model. And it effectively avoids the problem of incorrect matching of the query image background features and the supporting image foreground features when using traditional cross attention.

[0055] In order to make the purpose, technical solutions and advantages of the embodiments of the present invention clearer, the technical solutions in the embodiments of the present invention will be clearly and completely described below in conjunction with the drawings in the embodiments of the present invention. Obviously, the described embodiments are part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments in the present invention, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention. In addition, the technical features in the various embodiments or single embodiments provided by the present invention are arbitrarily combined with each other to form a new technical solution. This combination is not restricted by the sequence of steps and / or the structural composition mode, but must be based on the ability of ordinary technicians in this field to achieve. When the combination of technical solutions is contradictory or cannot be achieved, it should be considered that this combination of technical solutions does not exist and is not within the scope of protection required by the present invention.

[0056] An embodiment of the present invention provides a small sample medical image segmentation method based on query feature enhancement and graph attention mechanism. First, a medical image to be tested is obtained; then, a trained segmentation model is loaded, the medical image to be tested is input into the segmentation model, and a segmentation result is output.

[0057] Figure 2(a) is the overall architecture of the segmentation model, which is divided into three parts: feature extraction, query feature enhancement module, and multi-prototype collaborative segmentation module. First, the support image I s and query image I q Input model, use pre-trained ResNet-101 as feature extractor to extract image features and get support features F s and query feature F q , this part is feature extraction. Then, the query feature F q , support feature F s And support image mask M s Input the designed query feature enhancement module for processing to obtain the fully enhanced query features Finally, the enhanced query features Support feature F s And support image mask M s Input the designed multi-prototype collaborative segmentation module to obtain segmentation prediction

[0058] Figure 2 (b) is the specific architecture diagram of the multi-prototype collaborative segmentation module. This module generates segmentation prediction results through two branches and fuses them. Branch 1 first uses the support image mask M s Extract support features F s Get support prototype P s , then P s and Match and filter out high-confidence query predictions through thresholds Reuse extract Get the query prototype. Finally, the query prototype and Matching, generating query prospect predictions Branch 2 will be a single foreground prototype Support prototype set generated based on k-means clustering And query prototype Combine and construct multiple prototype sets P multi , and enhance its representation ability through a single-layer graph attention network Then with the query feature Perform fine-grained matching to generate multi-prototype prospect predictions Finally, branch one With branch 2 Fusion to obtain the final segmentation prediction

[0059] Figure 3 is the specific architecture diagram of the query feature enhancement module. First, the query feature Fq The input shows the visual center, which is processed in parallel by a lightweight multi-layer perceptron and a learnable visual center to obtain a comprehensive and discriminative feature representation. Afterwards, and F s The foreground area is divided into multiple query tiles Q i and support foreground tiles S i , and each Q i The corresponding S i Alignment, generate aligned support tiles Then the self-calibrated cross-attention mechanism is used to obtain Finally, Reassemble into complete query features and combine with Perform element-by-element addition to obtain fully enhanced query features

[0060] Figure 1 A flow chart of a small sample medical image segmentation method based on query feature enhancement and graph attention mechanism provided in an embodiment of the present invention includes four steps: medical image data acquisition and preprocessing, model construction, model training, and model testing.

[0061] Step 1: Collect medical image data and perform data preprocessing.

[0062] This embodiment uses a public medical image dataset as training and test data. All input three-dimensional images are reformatted into two-dimensional axial images and the resolution is adjusted to 256×256 pixels. In addition, each two-dimensional slice is replicated three times along the channel dimension to form three-channel input data to adapt the network structure.

[0063] This example follows the convention in the field of small sample medical image segmentation and adopts a scenario training strategy to construct an N-class K-sample segmentation task. In each scenario, N represents the number of categories involved (for example, the liver and spleen can be considered as two different categories), and K represents the number of samples in each category. train and the test set D test Each scenario consists of multiple scenarios, each of which contains a support set S and a query set Q. In the training phase, train Randomly sample pairs of support and query sets in [(I s , M s ), (I q , M q )], where I s and M s Represent the support image and its mask, I q and M qdenote the query image and its mask respectively. The model is constructed by s , M s ) to learn the category semantic representation to complete the query set I q At the same time, the mask M of the query set q is used as a supervisory signal and compared with the model prediction results to calculate the loss function, thereby guiding the model parameter optimization. In the test phase, the model only relies on D test There are K supporting images and their mask pairs (I s , M s ), for unlabeled query images of the same category I q This embodiment uses a setting of one class and one sample to complete the experiment.

[0064] Step 2: Construct a small sample medical image segmentation model based on query feature enhancement and graph attention mechanism. The model is divided into three parts: feature extraction, query feature enhancement module, and multi-prototype collaborative segmentation module.

[0065] Step 2.1: Perform feature extraction.

[0066] First, we use the ResNet-101 pre-trained on the MS-COCO dataset as the feature extractor fθ, and perform s and query image I q Extract 64×64 scale features to obtain support features F s and query feature F q . This process can be expressed as:

[0067]

[0068] Where C, H, and W represent the number of channels, height, and width of the feature map, respectively, where C is 512, and H and W are 64.

[0069] Step 2.2: Enhance the query features through the query feature enhancement module.

[0070] First, the query feature F is captured by using the display visual center method. q The query feature is then aligned with the foreground region of the supporting features using self-calibrated cross attention to generate a fully enhanced query feature.

[0071] Specifically, the query feature enhancement module uses the exposed visual center method to capture the global long-range dependencies in the query features through a lightweight multi-layer perceptron and aggregate the local angular regions through a learnable visual center. The results of the two are integrated to obtain a comprehensive and discriminative feature representation.

[0072] like Figure 3 As shown, first, the query feature F q Perform feature smoothing to obtain

[0073]

[0074] Among them, Conv 7×7 (·) represents a 7×7 convolution with a stride of 1, BN(·) represents a batch normalization layer, and ReLU(·) represents a ReLU activation function. The lightweight multi-layer perceptron and learnable visual center processing are used in parallel. The lightweight multi-layer perceptron consists of deep convolution and channel multi-layer perceptron, and its processing result is recorded as The process can be expressed as:

[0075]

[0076] Among them, GN(·) represents group normalization, DConv(·) represents depthwise convolution, and CMLP(·) represents channel multi-layer perceptron.

[0077] The learnable visual center uses a learnable codebook and a set of smoothing factors to encode the input features, dynamically adjusts the association strength between each pixel in the input features and each codeword in the codebook, and thus aggregates the local angular region. Specifically, the learnable visual center first passes a set of convolutional layers (including 1×1 convolution, 3×3 convolution, and 1×1 convolution) on the input features. Then, after 1×1 convolution, batch normalization and ReLU activation function processing, the feature Where E = H × W is the number of pixels in the feature map, and C is the number of channels. Next, these encoded features are fed into a learned codebook B = {b1, b2, ..., b D ), the codebook contains D codewords, D is selected from 64. By comparing each codeword b in the codebook d Calculate the matching degree and combine the learned smoothing factor S = {s1, s2, ..., s D ), weight the feature vector to obtain each codeword b d The weighted feature representation e d , the calculation formula is:

[0078]

[0079] in, is the i-th feature vector of the input feature, is the L2 distance between the feature vector and the codeword, s d is the smoothing factor. All ed The global feature e is obtained by fusion through fusion functions (including batch normalization layer, ReLU activation and mean layer). Then, e is mapped to influence factors using a fully connected layer and converted to weights using a Sigmoid function. Finally, these weights are combined with the input features Multiply each channel to get the weighted feature Z, which is then multiplied by the input feature Perform element-by-element addition to obtain the final output of the learnable visual center Then, the results processed by the lightweight multi-layer perceptron and the result processed by the learnable visual center Concatenate along the channel dimension to form an output with explicit visual center

[0080]

[0081] Among them, Conc 1×1 (·) represents a 1×1 convolution. Finally, the enhanced query features are obtained by element-wise addition.

[0082]

[0083] Furthermore, in order to fully utilize the support sample information to capture the correlation between the query image and the support image, we first and F s The foreground area is divided into multiple query tiles and support foreground tiles, which are represented as and where i∈{0,...,N 2 -1), N 2 represents the number of tiles, K×K represents the tile size, and both N and K are 8. i The corresponding S i Align, generate aligned support tiles, denoted as On this basis, a self-calibrated cross-attention mechanism is applied to Q i and Processing, query tile Q i Positioned as Q in the attention framework and aligned to the support tiles With Q iThey are grouped together as K and V. Self-calibrated cross attention calculates self-attention and cross attention in parallel, avoiding the problems caused by sequential calculation in traditional methods. Under this parallel calculation, the foreground features of the query image are matched with the foreground features of the support image through cross attention, while the background features of the query image are matched with the background features of the query image itself through self-attention, thus solving the problem of incorrect matching of the background features of the query image and the foreground features of the support image when cross attention is calculated separately in traditional methods. The process is as follows:

[0084]

[0085] Among them, d k is the dimension scaling factor used to adjust the size of the self-attention score. Subsequently, Q is aggregated using Att i and Information and Q i Fusion, to obtain enhanced query feature tiles

[0086]

[0087] Among them, FFN is a feed-forward neural network. Finally, all the enhanced query feature blocks are Reassemble into complete query features and combine with Perform element-by-element addition to obtain fully enhanced query features

[0088] Step 2.3: Perform segmentation prediction via the multi-prototype co-segmentation module.

[0089] like Figure 2 As shown in (b), the segmentation prediction of branch 1 is based on the principle that “pixels of the same object are more similar than pixels of different objects in the same category”. By matching the query prototype belonging to the same object with the query feature, the anatomical difference between the support sample and the query sample (different objects in the same category) is alleviated. First, the support feature F provided by the support set is s and support image mask M s , using the masked average pooling operation MAP(·) to obtain the support prototype P s =MAP(M s , F s ),in Contains a single foreground prototype and a single background prototype. Then calculate the support prototype P s With enhanced query features The cosine similarity of the two queries is calculated, and the prediction results with high confidence are selected as the query prediction results.

[0090]

[0091] in, is the indicator function, a1=20 is the scaling factor, and the mask threshold τ is used to control the sampling range of the query feature. The foreground and background query masks are set to {τ fg =0.7,τ bg =0.6}. Then, using the high confidence query prediction Aggregate query features Combined with the support prototype P s Get query prototype

[0092]

[0093] Then, in order to obtain a more accurate query prototype, the query prototype generation process above uses the support prototype P s The part is replaced by the query prototype P q , repeat this query prototype generation process twice to obtain a more accurate query prototype, denoted as Contains the foreground prototype and the background prototype, and only extracts the foreground prototype part, which is recorded as Finally, the soft threshold prediction method is used to obtain the prospect prediction of the query prototype

[0094]

[0095] Where a2=20 is the scaling factor, Sigmoid(·) represents the Sigmoid activation function, and T represents the soft threshold learned from the query image, which is obtained by the two fully connected layers on the top of the feature extractor.

[0096] The segmentation prediction of branch 2, first, from the support prototype P s Extracting a single foreground prototype Then, k-means clustering is used to construct a set of support prototypes for the foreground area of ​​the support image. G is selected as 5. The clustering process is to divide the foreground area into G subsets in the feature space Each subset represents a feature subspace. The initial support prototype is generated by averaging the features of each subset, namely: i=1,2,...,G,f n Represents the feature vector of the foreground area. Next, the single foreground prototype Support prototype set And query prototype Splice and assemble to get a multi-prototype set This multi-prototype set can fully capture the foreground features of the support image and the query image. In order to strengthen the relevance between prototypes and optimize their semantic relationships, a single-layer graph attention network is introduced to enhance P multi Specifically, P multi Each prototype in is regarded as an independent node, and an adjacency matrix is ​​constructed using the Spearman rank correlation coefficient to represent the connection relationship between nodes, which is defined as:

[0097]

[0098] Among them, A i，j Indicates P i and P j The edge weight between two prototypes, d i，j For the corresponding feature in P i and P j , where n is the number of features in each prototype. multi As the input feature matrix, calculate the node pair (i.e. prototype pair) P i and P j The attention coefficient e ij , and through A i,j Shield unrelated pairs of nodes:

[0099]

[0100] Among them, A i，j represents the uniformly initialized weight matrix, which is used to perform linear transformation on the input features; a3 and a4 are different parts of the uniformly initialized parameter vector used for attention calculation, and LeakyReLU(·) is the LeakyReLU activation function. Then, e is conditioned by the Softmax function. ij Normalize and get the final attention coefficient α ij :

[0101]

[0102] in, represents the set of neighbor nodes of node i. Then, the attention coefficient α is used ij Update the feature representation of each node:

[0103]

[0104] in, represents the updated feature vector of node i. Finally, the enhanced multi-prototype set is obtained in t represents the fusion coefficient of the prototype, and its value is 0.5. Different from the existing method that matches a single prototype with the query feature for similarity, the embodiment of the present invention dynamically assigns the most relevant prototype to each pixel in the query image, thereby making full use of the information of multiple prototypes to obtain foreground prediction.

[0105]

[0106] Among them, max(·) means selecting the highest score among all prototype scores as the final prediction value of the pixel position. Finally, based on the query prototype foreground prediction of branch 1 obtained above, Multi-prototype prospect prediction of branch 2 Conduct branch one With branch 2 A more comprehensive forecast of the future

[0107]

[0108] Among them, λ represents the fusion coefficient, which is set to 0.5. Then, the background prediction is obtained

[0109] Step 3: Input the medical image to be trained into the model for training and save the model training weights.

[0110] In the training phase, from D train Paired data of sampled support set and query set [(I s , M s ), (I q , M q )] is input into the model described in step 2, and the future forecast is obtained after being processed by the model and background prediction Afterwards, and Upsample to image size (256, 256) and use binary cross entropy loss to calculate the predicted mask With the real mask The difference between . The loss function is expressed as follows:

[0111]

[0112] Where H and W are both 256. Then, the prototype alignment regularization method is adopted to use the prediction results of the query image as the new support set during the training process to segment the original support image, thereby promoting the alignment of the support prototype with the query prototype. Its loss function is expressed as follows:

[0113]

[0114] Finally, the total loss function used to train the model is as follows:

[0115] L=L CE +L PAR

[0116] This example uses Python language and PyTorch framework and is implemented on NVIDIA RTX 3090 GPU. The batch size is set to 1, and the stochastic gradient descent (SGD) optimizer with momentum of 0.9 is used. The initial learning rate is 0.001, and it decays by 0.98 every 1000 iterations. A total of 60,000 iterations are performed to achieve the best performance, and the weight of the model at the best performance is saved.

[0117] Step 4: Input the medical image to be tested into the trained model to obtain the segmentation result.

[0118] In the testing phase, the test data is divided into support images and query images, where the support image is only an image containing the labeled target segmentation area, and the query image is all unlabeled images. The model and its trained weights are loaded, and the support image and query image are input into the model to generate segmentation masks for all query images.

[0119] Specifically, during the testing phase, from D test A pair of support images and their masks are randomly sampled from the support set (I s , M s ), and from D test The query image I is sampled from the query set q . Will I q ,I s and M s Input them together into the model described in step 2, and load the best model weights saved in step 3. The model outputs the foreground prediction of the query image Then Upsample to image size (256, 256) to obtain the final segmentation prediction result of the query image.

[0120] The implementation basis of each embodiment of the present invention is to implement programmed processing through a device with a processor function. Therefore, in engineering practice, the technical solutions and functions of each embodiment of the present invention are encapsulated into various modules. Based on this reality, on the basis of the above embodiments, an embodiment of the present invention provides a small sample medical image segmentation system based on query feature enhancement and graph attention mechanism, which is used to execute the small sample medical image segmentation method based on query feature enhancement and graph attention mechanism in the above method embodiment.

[0121] The system includes: an image data module, which is used to obtain a medical image to be tested; an image segmentation module, which is used to input the medical image to be tested into a trained segmentation model and output a segmentation result; wherein the training of the segmentation model includes: obtaining a training data set; using a pre-trained ResNet101 network as a feature extractor to extract features of a support image and a query image respectively to obtain support features and query features; using a designed query feature enhancement module to enhance the extracted query features; using a designed multi-prototype collaborative segmentation module to perform segmentation prediction, the multi-prototype collaborative segmentation module uses the query prototype to match the enhanced query features to generate a query prospect prediction; and combining a single prospect prototype, a support prototype set generated based on k-means clustering, and a query prototype to construct a multi-prototype set, and enhancing its representation ability through a graph attention mechanism, and then performing fine-grained matching with the query feature to generate a multi-prototype prospect prediction; fusing the generated query prospect prediction with the multi-prototype prospect prediction to obtain a final segmentation prediction; constructing a loss function to perform model training, and outputting a trained segmentation model.

[0122] The small sample medical image segmentation system based on query feature enhancement and graph attention mechanism provided by the embodiment of the present invention aims at the current situation of unsatisfactory performance of medical image segmentation. It adopts the above-mentioned several modules and realizes accurate segmentation of medical images through the mutual cooperation of query feature enhancement and multi-prototype collaborative segmentation.

[0123] It should be noted that the system embodiments provided by the present invention are not only used to implement the methods in the above-mentioned method embodiments, but also used to implement the methods in other method embodiments provided by the present invention. The only difference lies in the setting of corresponding functional modules, and the principles thereof are basically the same as the principles of the above-mentioned system embodiments provided by the present invention. As long as technical personnel in this field refer to the specific technical solutions in other method embodiments on the basis of the above-mentioned system embodiments, obtain corresponding technical means and technical solutions composed of these technical means by combining technical features, and on the premise of ensuring the practicality of the technical solutions, they will improve the modules in the above-mentioned system embodiments to obtain corresponding system class embodiments for implementing the methods in other method class embodiments.

[0124] It should be understood that the implementation of the system embodiment of the present invention can be achieved by referring to the aforementioned method embodiment, and the present invention will not be elaborated here.

[0125] Based on the same inventive concept as the aforementioned embodiment, an embodiment of the present invention also provides a small sample medical image segmentation device based on query feature enhancement and graph attention mechanism, comprising a memory and a processor, wherein the memory stores program instructions executed by the processor, and the processor calls the program instructions to execute the steps of the small sample medical image segmentation method based on query feature enhancement and graph attention mechanism.

[0126] Based on the same inventive concept as the aforementioned embodiment, an embodiment of the present invention also provides a non-transitory computer-readable storage medium, which stores computer instructions, and the computer instructions enable the computer to execute the steps of the small sample medical image segmentation method based on query feature enhancement and graph attention mechanism.

[0127] In summary, the present invention discloses a small sample medical image segmentation method based on query feature enhancement and graph attention mechanism, which is implemented by designing a query feature enhancement module and a multi-prototype collaborative segmentation module. Among them, the multi-prototype collaborative segmentation module consists of two prediction branches. Branch one effectively alleviates the anatomical differences between the support sample and the query sample by matching the query prototype and the query feature belonging to the same object. Branch two adopts the graph attention mechanism to promote the aggregation and mutual enhancement of multiple prototypes, and enhances the prototype's ability to represent the complex anatomical structure of the target area. Subsequently, fine-grained feature matching is achieved by dynamically assigning the most relevant prototype to each pixel in the query image. This matching strategy can be flexibly adjusted according to individual differences and better adapt to the variability of different anatomical structures. Finally, by fusing the prediction results of the two branches to achieve a more comprehensive segmentation prediction, the challenges brought by the heterogeneity of anatomical structures are effectively addressed. In addition, the query feature enhancement module enhances the representation ability of the query feature by capturing global long-distance dependencies and aggregating local angular areas. At the same time, the introduction of self-calibrated cross attention effectively aligns the query feature with the foreground area of ​​the supporting feature, fully utilizes the target category prior information provided by the supporting feature and its mask, guides the query feature to focus on the more valuable foreground area, thereby filtering out noise and background interference and improving the robustness of the model. It also effectively avoids the problem of incorrect matching between the background features of the query image and the foreground features of the supporting image that occurs when using traditional cross attention.

[0128] The terms "including" and "having" and any variations thereof in the specification and claims of the present invention and the above-mentioned drawings are intended to cover non-exclusive inclusions, for example, a process, method, system, product or apparatus comprising a series of steps or units is not necessarily limited to the steps or units explicitly listed, but may include other steps or units not explicitly listed or inherent to these processes, methods, products or apparatuses.

[0129] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention, rather than to limit it. Although the present invention has been described in detail with reference to the aforementioned embodiments, those skilled in the art should understand that they can still modify the technical solutions described in the aforementioned embodiments, or replace some or all of the technical features therein with equivalents. However, these modifications or replacements do not deviate the essence of the corresponding technical solutions from the technical solutions of the embodiments of the present invention.

Claims

1. A small sample medical image segmentation method based on query feature enhancement and graph attention mechanism, characterized in that: include: Acquire a medical image to be tested; Input the medical image to be tested into the trained segmentation model and output the segmentation result; wherein the training of the segmentation model includes: Step 1, obtain the training data set; Step 2: Use the pre-trained ResNet101 network as a feature extractor to extract features from the support image and the query image to obtain support features and query features. Step 3, using the designed query feature enhancement module to enhance the extracted query features; Step 4, using the designed multi-prototype collaborative segmentation module to perform segmentation prediction, the multi-prototype collaborative segmentation module uses the query prototype to match the enhanced query features to generate a query prospect prediction; and, combining the single prospect prototype, the support prototype set generated based on k-means clustering, and the query prototype to construct a multi-prototype set, and enhance its representation ability through the graph attention mechanism, and then performing fine-grained matching with the query features to generate a multi-prototype prospect prediction; the generated query prospect prediction is fused with the multi-prototype prospect prediction to obtain the final segmentation prediction; Step 5: Construct a loss function for model training and output the trained segmentation model.

2. According to claim 1, the small sample medical image segmentation method based on query feature enhancement and graph attention mechanism is characterized in that: The extracted query features are enhanced using the designed query feature enhancement module, including: By using the displayed visual center method, the global long-distance dependencies in the query features are captured by a lightweight multi-layer perceptron, and the local angular regions are aggregated by a learnable visual center to integrate the global long-distance dependencies and the local angular regions to obtain enhanced query features.

3. According to claim 2, the small sample medical image segmentation method based on query feature enhancement and graph attention mechanism is characterized in that: The enhanced query features also include: Self-calibrated criss-cross attention is adopted to align query features with foreground regions of supporting features, generating comprehensively enhanced query features.

4. The small sample medical image segmentation method based on query feature enhancement and graph attention mechanism according to claim 1, characterized in that: Utilize query prototypes and enhanced query feature matching to generate query prospect predictions, including: Using the support features and support image masks provided by the support set, the support prototype is obtained using the mask average pooling operation; Calculate the cosine similarity between the supporting prototype and the enhanced query features, and select the prediction result with high confidence as the query prediction; Aggregate query features using high-confidence query predictions and combine with supporting prototypes to obtain query prototypes; The part of the support prototype used in the above query prototype generation process is replaced with the query prototype, and the query prototype generation process is repeated twice to obtain a more accurate query prototype; A foreground prototype portion of the more accurate query prototype is extracted, and a soft threshold prediction method is used to obtain a foreground prediction of the query prototype.

5. According to claim 4, the small sample medical image segmentation method based on query feature enhancement and graph attention mechanism is characterized in that: The single foreground prototype, the support prototype set generated based on k-means clustering, and the query prototype are combined to construct a multi-prototype set, and its representation ability is enhanced through the graph attention mechanism. Then, fine-grained matching is performed with the query features to generate multi-prototype foreground predictions, including: A single foreground prototype is extracted from the obtained supporting prototypes, and a set of supporting prototypes is constructed for the foreground region of the supporting image using k-means clustering; The single foreground prototype, the supporting prototype set and the query prototype are spliced ​​together to assemble a multi-prototype set; A single-layer graph attention network is introduced to enhance the representation ability of each prototype in the multi-prototype set, and the enhanced multi-prototype set is obtained; Based on the enhanced multi-prototype set, fine-grained matching is performed with query features to generate multi-prototype prospect predictions.

6. The small sample medical image segmentation method based on query feature enhancement and graph attention mechanism according to claim 1, characterized in that: Construct a loss function for model training and output the trained segmentation model, including: The foreground prediction and background prediction are upsampled to the set image size, and the weighted sum of the binary cross entropy loss and the prototype alignment regularization loss is used as the total loss function; The total loss function is used to perform model training and output a trained segmentation model.

7. A small sample medical image segmentation system based on query feature enhancement and graph attention mechanism, characterized by: include: An image data module, used for acquiring medical images to be tested; An image segmentation module is used to input the medical image to be tested into a trained segmentation model and output a segmentation result; wherein the training of the segmentation model includes: Step 1, obtain the training data set; Step 2: Use the pre-trained ResNet101 network as a feature extractor to extract features from the support image and the query image to obtain support features and query features. Step 3, using the designed query feature enhancement module to enhance the extracted query features; Step 4, using the designed multi-prototype collaborative segmentation module to perform segmentation prediction, the multi-prototype collaborative segmentation module uses the query prototype to match the enhanced query features to generate a query prospect prediction; and, combining the single prospect prototype, the support prototype set generated based on k-means clustering, and the query prototype to construct a multi-prototype set, and enhance its representation ability through the graph attention mechanism, and then performing fine-grained matching with the query features to generate a multi-prototype prospect prediction; the generated query prospect prediction is fused with the multi-prototype prospect prediction to obtain the final segmentation prediction; Step 5: Construct a loss function for model training and output the trained segmentation model.

8. A small sample medical image segmentation device based on query feature enhancement and graph attention mechanism, characterized in that: It includes a memory and a processor, the memory stores program instructions executed by the processor, and the processor calls the program instructions to execute the steps of the small sample medical image segmentation method based on query feature enhancement and graph attention mechanism as described in any one of claims 1 to 6.

9. A non-transitory computer-readable storage medium, characterized in that: The non-transitory computer-readable storage medium stores computer instructions, which enable the computer to execute the steps of the small sample medical image segmentation method based on query feature enhancement and graph attention mechanism as described in any one of claims 1 to 6.

Images