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

Abstract

The present invention discloses a small sample medical image segmentation method and system 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. The query feature enhancement module provides comprehensive and highly discriminative query features for subsequent prototype generation and matching by capturing global long-distance dependencies, aggregating local angular areas, and integrating supporting sample information. The multi-prototype collaborative segmentation module consists of two prediction branches. Branch one matches the query prototypes and query features belonging to the same object, and branch two uses a graph attention mechanism to promote aggregation and mutual enhancement between multiple prototypes, and dynamically assigns the most relevant prototype to each pixel in the query image. A more comprehensive segmentation prediction is achieved by fusing the prediction results of the two branches. The present invention effectively addresses the challenges brought by anatomical structure heterogeneity and improves the accuracy and adaptability of segmentation.

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 a crucial step in medical image analysis and processing, and a prerequisite for effective disease diagnosis and treatment. With the advancement of deep learning, fully supervised deep learning models trained for specific organs or tissues have achieved promising results in many medical image segmentation tasks. However, due to the high cost of medical image annotation, patient privacy concerns, and the scarcity of samples for certain rare diseases, obtaining the large amounts of labeled data required for fully supervised training is both expensive and difficult. In recent years, numerous small-sample image segmentation methods have been proposed, effectively alleviating the burden of data annotation. These methods can train models with excellent generalization capabilities, allowing them to quickly adapt to new categories and achieve segmentation predictions with the guidance of a small number of labeled samples. Prototype learning-based methods, in particular, achieve pixel-level classification by constructing a prototype for each category and calculating the similarity between the query image's pixel features and the prototype. These methods, which extract representative features from a limited number of samples, demonstrate excellent generalization and efficiency, and have therefore attracted considerable 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 heart 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 small-sample medical image segmentation methods based on prototype learning, 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 research of small-sample medical image segmentation based on prototype learning.

[0004] Some methods construct separate prototypes for the foreground region (target segmentation region) and background region (regions outside the target segmentation region) of the support image. However, due to the limited number of support image slices, feature modeling of complex background regions is often insufficient. Other methods construct a single prototype only for the foreground region of the support image through masked average pooling. This approach is effective in certain scenarios (e.g., when the foreground category is relatively homogeneous and does not contain multiple distinct regions with clear boundaries). However, when dealing with complex foreground regions, this approach can lead to semantic ambiguity and information loss. For example, when the left ventricular blood pool and left ventricular myocardium are treated as a single foreground category (left ventricle). Furthermore, some methods construct multiple prototypes for the foreground region but merge these prototypes into a single prototype through global average pooling before calculating similarity with the query feature. This process can result in the loss of unique information and details in each prototype, making it difficult for the model to adapt to individual differences when representing highly heterogeneous foreground regions. Overall, these methods have not effectively addressed the challenges posed 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, because 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 inability to fully play the guiding role of the support samples 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] Obtaining a medical image to be tested;

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

[0010] Step 1: Obtain training data set;

[0011] Step 2: Use the pre-trained ResNet101 network as a feature extractor to extract features from the support image and query image respectively 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: Perform segmentation prediction using the designed multi-prototype collaborative segmentation module. The multi-prototype collaborative segmentation module matches the query prototype with the enhanced query features to generate a query foreground prediction. Furthermore, a single foreground prototype, a support prototype set generated based on k-means clustering, and the query prototype are combined to construct a multi-prototype set, and its representational capability is enhanced through a graph attention mechanism. This is then fine-grainedly matched with the query features to generate a multi-prototype foreground prediction. The generated query foreground prediction is fused with the multi-prototype foreground 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 utilizing the displayed visual center method, a lightweight multi-layer perceptron is used to capture the global long-distance dependencies in the query features. Local angular regions are aggregated through a learnable visual center, and the global long-distance dependencies and local angular regions are integrated to obtain enhanced query features.

[0017] As a further technical solution, after obtaining the enhanced query features, the following is also included:

[0018] Self-calibrated cross-attention is employed to align query features with the 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] Through 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 them with supporting prototypes to obtain query prototypes;

[0023] Replace the part of the support prototype used in the above query prototype generation process with the query prototype, repeat this query prototype generation process 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. The representation capability of the set is enhanced through a graph attention mechanism. Subsequently, 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 support prototypes, and a set of support prototypes is constructed for the foreground region of the support 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 for model training and the trained segmentation model is output, 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 the 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 configured 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 training data set;

[0037] Step 2: Use the pre-trained ResNet101 network as a feature extractor to extract features from the support image and query image respectively 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: Perform segmentation prediction using the designed multi-prototype collaborative segmentation module. The multi-prototype collaborative segmentation module matches the query prototype with the enhanced query features to generate a query foreground prediction. Furthermore, a single foreground prototype, a support prototype set generated based on k-means clustering, and the query prototype are combined to construct a multi-prototype set, and its representational capability is enhanced through a graph attention mechanism. This is then fine-grainedly matched with the query features to generate a multi-prototype foreground prediction. The generated query foreground prediction is fused with the multi-prototype foreground 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, 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 invention, a non-transitory computer-readable storage medium is provided, wherein the non-transitory computer-readable storage medium 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.

[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 the query feature enhancement module and constructs comprehensive and representative prototypes through the multi-prototype collaborative segmentation module, thereby accurately matching the query features and effectively improving the segmentation accuracy.

[0045] 2. The multi-prototype collaborative segmentation module of our invention uses a graph attention mechanism to enhance the prototypes' ability to represent complex anatomical structures, making the model more adaptable to different anatomical regions. Furthermore, a fine-grained feature matching mechanism dynamically assigns prototypes based on 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 and 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 will be 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 the process 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 the query feature enhancement module provided in an example of the present invention. DETAILED DESCRIPTION

[0051] This invention aims to address the problems in existing small-sample medical image segmentation methods, such as the lack of representative prototype construction and insufficient capture of local details during query feature matching. This leads to an inability to effectively address anatomical heterogeneity, which in turn affects segmentation accuracy and adaptability. Furthermore, this invention addresses the challenges of existing feature extraction methods in processing global and local information and leveraging sample information, thereby improving segmentation accuracy and robustness.

[0052] 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 the support sample and the query sample by matching the query prototype with the query feature belonging to the same object. Branch two uses the 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 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.

[0054] In addition, the present invention designs a query feature enhancement module that enhances the representation capability of query features by capturing global long-range dependencies and aggregating local angular regions. Simultaneously, the introduction of self-calibrated cross-attention effectively aligns the query feature with the foreground region of the supporting features, fully utilizing the target category prior information provided by the supporting features and their masks to guide the query feature to focus on the more valuable foreground region, thereby filtering out noise and background interference and improving the robustness of the model. This effectively avoids the problem of mismatching between the query image background features and the supporting image foreground features that occurs 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 of the present invention, all other embodiments obtained by ordinary technicians in this field without making 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 structural composition mode, but must be based on the ability of ordinary technicians in this field to implement it. When the combination of technical solutions is contradictory or cannot be implemented, it should be deemed 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 the segmentation result is output.

[0057] Figure 2(a) is the overall architecture diagram 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 obtain 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, which 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 prototypes Combine and construct multiple prototype sets P multi , and enhance its representation ability by a single-layer graph attention network to obtain Then with the query feature Perform fine-grained matching to generate multi-prototype prospect predictions Finally, branch one With branch 2 Fusion is performed to obtain the final segmentation prediction

[0059] Figure 3 This 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, generating 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 example uses a publicly available medical image dataset as training and test data. All input 3D images are reformatted into 2D axial images and adjusted to a resolution of 256×256 pixels. Furthermore, each 2D slice is replicated three times along the channel dimension to form three-channel input data to accommodate the network architecture.

[0063] This example follows the conventions 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, 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 paired data of support set and query set [(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 It is used as a supervisory signal and compared with the model prediction results to calculate the loss function, thereby guiding the parameter optimization of the model. In the test phase, the model only relies on D test There are K support images and their mask pairs (I s , M s ), for the unlabeled query image I of the same category 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 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, 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 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 utilizes the revealed visual center method to capture global long-range dependencies in query features through a lightweight multi-layer perceptron and aggregate local angular regions through a learnable visual center. The results of these two methods 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 multilayer perceptron and learnable visual center processing are used in parallel. The lightweight multilayer perceptron consists of depth convolution and channel multilayer 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 correlation strength between each pixel in the input features and each codeword in the codebook, and thus aggregates the local corner area. Specifically, the learnable visual center first uses a set of convolutional layers (including 1×1 convolution, 3×3 convolution and 1×1 convolution) to transform 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 as 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 eigenvector 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 channel by 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 results processed by the learnable visual center Splicing 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 of ​​​​is divided into multiple query tiles and support foreground tiles, respectively, represented as and where i∈{0,...,N 2 -1), N 2 Indicates the number of tiles, K×K indicates 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 with the support tiles With Q iThey are grouped together as K and V. Self-calibrated cross-attention parallelizes the calculation of self-attention and cross-attention, 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 between 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 feedforward 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 used to s and support image mask M s , use the mask average pooling operation MAP(·) to get 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 predictions.

[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 query prediction with high confidence 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 is used to support the 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, recorded as Contains foreground prototype and background prototype, only extracts the foreground prototype part, 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 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 region 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 prototypes 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, the representation capability of each prototype in 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 The ranking difference between them is n, 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, the Softmax function is used to transform e ij Normalize and get the final attention coefficient α ij :

[0101]

[0102] in, Represents the set of neighbor nodes of node i. Then, use the attention coefficient α 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. Unlike 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, the query prototype foreground prediction based on the previously obtained branch 1 is and multi-prototype prospect prediction of branch 2 Carry out branch one With branch 2 The integration of

[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] During the training phase, 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 the model processing 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 loss function is expressed as follows:

[0111]

[0112] Where H and W are both 256. Then, the prototype alignment regularization method is adopted. During the training process, the prediction results of the query image are used as the new support set 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 and the PyTorch framework, implemented on an NVIDIA RTX 3090 GPU. A batch size of 1 is set, and a stochastic gradient descent (SGD) optimizer with a momentum of 0.9 is used. The initial learning rate is 0.001, decaying by 0.98 every 1,000 iterations. 60,000 iterations are performed to achieve optimal performance, and the model weights at the optimal performance are saved.

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

[0118] During the testing phase, the test data is divided into support images and query images. The support image is a single image containing the labeled target segmentation region, while the query image is all unlabeled images. The model and its trained weights are loaded, and the support and query images are fed into the model to generate segmentation masks for all query images.

[0119] Specifically, during the testing phase, 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 achieved through programmed processing by a device with processor functionality. 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-mentioned embodiments, an embodiment of the present invention provides a small-sample medical image segmentation system based on query feature enhancement and graph attention mechanism. This system is used to implement the small-sample medical image segmentation method based on query feature enhancement and graph attention mechanism described in the above-mentioned method embodiment.

[0121] The system includes: an image data module for acquiring a medical image to be tested; an image segmentation module for inputting the medical image to be tested into a trained segmentation model and outputting a segmentation result; wherein the training of the segmentation model includes: acquiring 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 foreground prediction; and combining a single foreground 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 features to generate a multi-prototype foreground prediction; fusing the generated query foreground prediction with the multi-prototype foreground prediction to obtain a final segmentation prediction; constructing a loss function for 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 addresses the current situation where medical image segmentation performance is not ideal. It adopts the above-mentioned 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 those skilled in the art 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, 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 described in detail 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, 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 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 support feature, fully utilizing the target category prior information provided by the support feature and its mask to guide 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 mismatching the query image background features with the support image foreground features that occurs when using traditional cross-attention.

[0128] The terms "including" and "having" and any variations thereof in the description 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 that includes 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 by: include: Obtaining a medical image to be tested; Inputting the medical image to be tested into the trained segmentation model and outputting a segmentation result; wherein the training of the segmentation model includes: Step 1: Obtain training data set; Step 2: Use the pre-trained ResNet101 network as a feature extractor to extract features from the support image and query image respectively to obtain support features and query features; Step 3, using the designed query feature enhancement module to enhance the extracted query features, including: using the revealed visual center method to capture the global long-range dependency in the query features through a lightweight multi-layer perceptron, aggregating the local angular regions through a learnable visual center, and integrating the global long-range dependency and the local angular regions to obtain enhanced query features; Step 4: Perform segmentation prediction using the designed multi-prototype collaborative segmentation module. The multi-prototype collaborative segmentation module matches the query prototype with the enhanced query features to generate a query foreground prediction. Furthermore, a single foreground prototype, a support prototype set generated based on k-means clustering, and the query prototype are combined to construct a multi-prototype set, and its representational capability is enhanced through a graph attention mechanism. This is then fine-grainedly matched with the query features to generate a multi-prototype foreground prediction. The generated query foreground prediction is fused with the multi-prototype foreground prediction to obtain the final segmentation prediction. Step 5: Construct a loss function for model training and output the trained segmentation model.

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

3. The small sample medical image segmentation method based on query feature enhancement and graph attention mechanism according to claim 1 is characterized in that: Utilize query prototypes and enhanced query feature matching to generate query prospect predictions, including: Through 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 them with supporting prototypes to obtain query prototypes; Replace the part of the support prototype used in the above query prototype generation process with the query prototype, repeat this query prototype generation process 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.

4. The small sample medical image segmentation method based on query feature enhancement and graph attention mechanism according to claim 3 is characterized in that: The multi-prototype set is constructed by combining a single foreground prototype, a support prototype set generated based on k-means clustering, and a query prototype. Its representation capability is enhanced through the graph attention mechanism. It is then fine-grainedly matched with the query features to generate multi-prototype foreground predictions, including: A single foreground prototype is extracted from the obtained support prototypes, and a set of support prototypes is constructed for the foreground region of the support 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.

5. The small sample medical image segmentation method based on query feature enhancement and graph attention mechanism according to claim 1 is 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 the trained segmentation model.

6. 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 configured 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 training data set; Step 2: Use the pre-trained ResNet101 network as a feature extractor to extract features from the support image and query image respectively to obtain support features and query features; Step 3, using the designed query feature enhancement module to enhance the extracted query features, including: using the revealed visual center method to capture the global long-range dependency in the query features through a lightweight multi-layer perceptron, aggregating the local angular regions through a learnable visual center, and integrating the global long-range dependency and the local angular regions to obtain enhanced query features; Step 4: Perform segmentation prediction using the designed multi-prototype collaborative segmentation module. The multi-prototype collaborative segmentation module matches the query prototype with the enhanced query features to generate a query foreground prediction. Furthermore, a single foreground prototype, a support prototype set generated based on k-means clustering, and the query prototype are combined to construct a multi-prototype set, and its representational capability is enhanced through a graph attention mechanism. This is then fine-grainedly matched with the query features to generate a multi-prototype foreground prediction. The generated query foreground prediction is fused with the multi-prototype foreground prediction to obtain the final segmentation prediction. Step 5: Construct a loss function for model training and output the trained segmentation model.

7. 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 5.

8. 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 5.

Images