A small-sample medical image segmentation method based on self-attention mechanism and contrastive learning

Abstract

The present invention discloses a small-sample medical image segmentation method based on a self-attention mechanism and contrastive learning. The method first requires preprocessing Magnetic Resonance Imaging (MRI) data and pre-segmenting the processed data. The pre-segmentation result is then masked to obtain a pseudo-mask. The pseudo-mask and its corresponding medical image are fed into a deep learning network model, the model parameters are updated based on a comprehensive loss, and the model training results are saved. The trained model is used to segment new medical image data to obtain a segmentation mask. The present invention improves the accuracy of network model learning by correcting the pseudo-mask. The addition of the self-attention mechanism and contrastive learning enables information transmission within the small-sample training process and suppresses interference from image background information, achieving efficient segmentation of medical images with limited annotated data.

Description

Technical Field

[0001] The present invention relates to a small sample medical image segmentation method based on a self-attention mechanism and contrastive learning, and belongs to the technical field of medical image processing. Background Art

[0002] With the rapid development of medical imaging technology, medical images have become an important tool for diagnosing, treating, and evaluating diseases. Medical image segmentation, as one of the key technologies in image processing, aims to identify and extract anatomical structures or lesions of interest from complex medical images. In recent years, deep learning technology has made significant progress in the field of medical image segmentation, and this advancement has greatly promoted the automation and precision of medical imaging diagnosis. Through deep learning, especially the widespread application of convolutional neural networks and their derivative models, the accuracy and efficiency of medical image segmentation have achieved a qualitative leap, not only reducing the time cost of manual analysis, but also improving the accuracy and consistency of diagnosis.

[0003] However, despite the enormous potential of deep learning technology in the field of medical image segmentation, the difficulty in obtaining labeled samples remains one of the main factors restricting its further development. Medical image labeling is typically performed manually by experienced radiologists or pathologists. This process is not only time-consuming and labor-intensive, but also requires a high level of expertise and skills from the labeler. Furthermore, due to the complexity and diversity of medical images, even the same type of lesion can exhibit vastly different morphologies and characteristics in different patients, further increasing the difficulty and uncertainty of labeling. Therefore, how to reduce the reliance on labeled samples while maintaining model performance is a key issue that deep learning currently needs to address in the field of medical image segmentation.

[0004] The attention mechanism is a powerful tool for processing sequential data. It mimics the way humans allocate attention when reading or understanding, dynamically assigning different weights to different positions in a sequence, thereby more efficiently and accurately capturing information dependencies. In medical image segmentation, the attention mechanism can enhance the model's sensitivity to key features, helping it ignore irrelevant or redundant information, and thus more efficiently utilize limited labeled samples. Contrastive learning is an unsupervised or self-supervised learning method that learns feature representations by comparing the similarities between positive and negative samples. In medical image segmentation, contrastive learning can leverage limited labeled data by constructing positive samples (similar image pairs) and negative samples (dissimilar image pairs) to enhance the model's feature extraction capabilities and improve segmentation performance in small sample sizes. We combine the attention mechanism with contrastive learning. The attention mechanism helps the model focus on key regions, while contrastive learning enhances the model's understanding of lesion features and its generalization capabilities. This combination enables the model to learn effective feature representations and achieve high-accuracy segmentation even with limited labeled samples. Summary of the Invention

[0005] The technical problem to be solved by the present invention is the difficulty in obtaining medical image annotated data. It also overcomes many limitations of existing methods and provides a small-sample medical image segmentation method based on self-attention mechanism and contrastive learning, which can achieve medical image segmentation using only a small number of annotated samples.

[0006] In order to achieve the above objectives, the technical solution of the present invention is:

[0007] A small-sample medical image segmentation method based on self-attention mechanism and contrastive learning includes the following steps:

[0008] (1) MRI data preprocessing and pre-segmentation: To adapt to the model input, the image data is reformatted into two dimensions and resized to 256 × 256 pixels. Each two-dimensional slice is replicated three times in the channel dimension and pre-segmented using an unsupervised image segmentation algorithm to obtain a rough pseudo mask.

[0009] (2) Acquisition of pseudo mask: Correct the rough pseudo mask obtained by pre-segmentation and remove the erroneous areas to obtain the pseudo mask;

[0010] (3) Constructing a deep learning network model: The feature extraction network used is ResNet-101 pre-trained on the MS-COCO dataset, where the classifier is replaced by a 1×1 convolutional layer. A multi-head attention module is added to the back of the model to transfer information between features and build contrastive learning between the foreground and background of the query image;

[0011] (4) Model training process: The total number of iterations of the model is preset. In each iteration, two MRI images are taken out, which are called support images and query images respectively. The two images are fed into the network model to extract features. The prototype of the support feature is extracted according to the pseudo mask of the support image. The query feature is segmented using the support prototype. The query feature segmentation result and the query image pseudo mask are used to calculate the loss back propagation to update the network model. After the iteration, the model parameters are saved.

[0012] (5) Model testing process: The model is loaded with the model parameters obtained in step (4), a new image is taken out as the support image, and the remaining images are used as query images. The support image is sent to the model to obtain the prototype, and the prototype is used to segment all query images to obtain the segmentation results.

[0013] Specifically, the steps for correcting the rough pseudo mask obtained by pre-segmentation in step (2) are as follows: create a Boolean array B of the same size as the pseudo mask to track the visited pixels, and initialize the size of the largest connected area in the mask to C = 0. Traverse each pixel in the pseudo mask, and for each pixel that has not been visited and is marked as 1, perform a breadth-first search on the entire pseudo mask starting from the pixel. If the connected area found is larger than the current maximum value C max , then update the maximum connected area and record the pixel point in B. After traversing all pixels, retain the maximum connected area C max , create a new array M of the same size as the rough pseudo mask, and put the maximum connected area C in M max The corresponding position is marked as 1 to obtain the corrected pseudo mask.

[0014] Specifically, the step (3) of using the multi-head attention model to transfer feature information is as follows: first, the input support image and query image are processed through the feature extraction network respectively to obtain their respective feature representations. This process can be expressed as:

[0015] F s =f θ (I s )

[0016] F q =f θ (I q )

[0017] Among them, I s and I q represent support images and query images respectively, and f θ Represents the mapping function of the feature extraction network, F s and F q are the extracted support features and query features. Next, we query the feature F qThe self-attention mechanism allows the model to dynamically focus on other parts of the feature when processing the query feature, thereby capturing the dependencies between elements. Its calculation formula is as follows:

[0018] Q = Attention(F q W q Q ,F q W q K ,F q W q V )

[0019]

[0020] Where W q Q 、W q K 、W q V is for query feature F q The linear transformation matrix is ​​used to generate the query (Q = F q W q Q ), key (K=F q W q K ) and value (V=F q W q V ) vector. The Attention function is the core calculation process of the self-attention mechanism, and its output Q q represents the enhanced query features, K T and d k Denote the transpose and dimension of the key vector, respectively. Next, we calculate the similarity between the key (K) of the support feature and the query (Q) of the query feature, and apply the softmax function for normalization to generate the attention weights. Here, we use softmax (dim = -2) to ensure that the sum of the similarity scores of each column of the result (corresponding to each support pixel) is 1, allowing the support features to aggregate toward the query features. Subsequently, we multiply the attention weights with the pseudo mask of the support image to filter out irrelevant areas. Finally, the filtered attention weights are multiplied by the value (V) of the support image to achieve feature aggregation. This process can be expressed as:

[0021]

[0022] F' q =attn·V s

[0023] Among them, Ks T represents the transposition of the support key, ⊙ represents element-wise multiplication, M s is a pseudo mask of the support image, used to filter out irrelevant areas, attn is the intermediate result, F q ' is the result of aggregating the support feature information into the query feature, and the subscripts q and s represent the information from the query image and the support image, respectively.

[0024] Finally, we feed the aggregated query features into a module consisting of a residual connection, a multi-layer perceptron, and another residual connection to obtain the final aggregated query features. This process uses a one-way information transfer strategy to efficiently aggregate key feature information from the support set to the query set, avoiding the problem of information misleading caused by inaccurate query image prediction masks. The specific formula is as follows:

[0025] O q =LN(attn·V+Q)

[0026] O out =LN(MLP(O q )+O q )

[0027] Among them, LN represents layer normalization, MLP represents multi-layer perceptron, and O q represents the intermediate query feature, O out is the final query feature,

[0028] Specifically, the query image foreground-background contrast learning steps in step (3) are as follows:

[0029] First, we use the Masked Average Pooling (MAP) technique to extract the support prototype representing the foreground features from the labeled support image. s The calculation formula is as follows:

[0030] P s =MAP(F s ,M s )

[0031] Next, we use the cosine similarity between these support prototypes and query features to generate the query mask and extract the query prototype P q The calculation of cosine similarity and the extraction of query prototypes help us preliminarily screen out the areas in the query image that may be related to the foreground. In order to distinguish the complex background area, we first invert the support image mask to obtain Then take out the background features corresponding to the support features and express them as F s b, using the Voronoi-based method to Divide and obtain the background area mask set Here B n represents the nth background region, R is the number of regions, which is set to 150. The formula for extracting different region representations is as follows:

[0032]

[0033] Among them, P n represents the prototype of the nth region, H and W represent The height and width, and B n,i Denote the support feature and background area corresponding to the i-th pixel respectively. In the contrastive learning framework we constructed, the query prototype is used as the anchor point, the support prototype is used as the positive sample, and multiple partition prototypes of the background area constitute the negative sample set. Each negative sample represents a different part of the background area, thus providing rich background information for the model.

[0034] Specifically, in step (4), the total number of iterations of the model is set to 50,000. The loss of the model includes the use of cross entropy loss to directly measure the similarity between the predicted probability distribution and the true mask distribution. For each pixel, the cross entropy loss is The calculation formula is:

[0035]

[0036] Among them, M q represents the ground-truth mask of the query image, and denote the query image background and foreground masks predicted by the model, respectively.

[0037] In addition, a threshold loss is added to the model Where T is the segmentation threshold, α is the adjustment parameter, which is set to 20 by default. By minimizing the learning threshold, the foreground features are encouraged to be more compact in the embedding space. In addition, the foreground-background contrast loss is proposed. The segmentation performance of the model is optimized by maximizing the similarity between the query prototype and the support prototype while minimizing the similarity with the background negative sample. The mathematical expression is:

[0038]

[0039] Among them, cos(·) is the cosine similarity function. It is worth noting that the corrected mask background area occupies the majority, so The loss can effectively use a large number of negative samples to optimize the model, Add up to get the total loss The calculation formula is as follows:

[0040] Specifically, in step (5), the model testing process is as follows: In the testing phase, a set of medical images that were not used in the training process are used for segmentation. The test image is divided into a support image and a query image, where the support image contains only one sample with accurate annotations. The support and query features are extracted using the model obtained in the training phase, and then the support image features are transferred to the query image features using the attention mechanism. The support image pseudo mask M is used s From the support feature F s Extract the supporting prototype P s Finally, with the help of the extracted supporting feature prototype, the query feature F q Segmentation is performed to obtain the segmentation results. Specifically, the attention mechanism enables the model to accurately focus on key features in the image, while contrastive learning strengthens the model's feature learning and generalization capabilities. The synergy of these two greatly improves the model's representational capabilities. Because the trained model has fully learned the ability to extract key features from a small amount of data and effectively generalize, when faced with new data during the testing phase, the model only needs to rely on a single labeled example to accurately segment the target organ.

[0041] Compared with the prior art, the beneficial effects of the present invention are as follows: the present invention provides a small sample medical image segmentation method based on self-attention mechanism and contrastive learning, which can realize medical image segmentation using a small amount of labeled data; the proposed mask correction method removes the noise and error areas associated with the pseudo mask, providing a more accurate mask for subsequent model learning; the use of the attention mechanism realizes one-way information transfer, that is, the key feature information is efficiently aggregated only from the support feature to the query feature, avoiding the information misleading problem caused by the inaccurate mask prediction of the query image; foreground-background contrastive learning aims to enable the model to learn more discriminative feature representations under limited data by constructing a contrast relationship between positive samples (foreground areas) and negative samples (background areas). This method effectively overcomes the model training difficulties caused by the lack of data annotation, improves the generalization ability and robustness of the model in the small sample medical image segmentation scenario, reduces the interference of image noise and complex background, solves the problem of poor segmentation effect of deep learning models in a small number of sample scenarios, and provides new ideas and directions for research in the field of small sample medical image segmentation. BRIEF DESCRIPTION OF THE DRAWINGS

[0042] Figure 1 It is a schematic diagram of the overall process of the present invention;

[0043] Figure 2 This is a flowchart of the deep learning segmentation network training of the present invention;

[0044] Figure 3 This is a flow chart of the present invention during the testing phase;

[0045] Figure 4 Schematic diagram of the segmentation result of the present invention. DETAILED DESCRIPTION

[0046] In order to enable those skilled in the art to better understand the technical solutions of the present invention, the technical solutions in the embodiments of the present invention will be clearly and completely described below in conjunction with the accompanying drawings of the embodiments of the present invention. Obviously, the described embodiments are only 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 efforts should fall within the scope of protection of the present invention.

[0047] Example: Figure 1 The overall flow chart of a small sample medical image segmentation method based on self-attention mechanism and contrastive learning provided in this embodiment includes the following steps:

[0048] (1) MRI data preprocessing and pre-segmentation: To adapt to the model input, the image data is reformatted into two dimensions and resized to 256 × 256 pixels. Each two-dimensional slice is replicated three times in the channel dimension and pre-segmented using an unsupervised image segmentation algorithm to obtain a rough pseudo mask.

[0049] (2) Acquisition of pseudo mask: Correct the rough pseudo mask obtained by pre-segmentation and remove the erroneous areas to obtain the pseudo mask;

[0050] (3) Constructing a deep learning network model: The feature extraction network used is ResNet-101 pre-trained on the MS-COCO dataset, where the classifier is replaced by a 1×1 convolutional layer. A multi-head attention module is added to the back of the model to transfer information between features and build contrastive learning between the foreground and background of the query image;

[0051] (4) Model training process: The total number of iterations of the model is preset. In each iteration, two MRI images are taken out, which are called support images and query images respectively. The two images are fed into the network model to extract features. The prototype of the support feature is extracted according to the pseudo mask of the support image. The query feature is segmented using the support prototype. The query feature segmentation result and the query image pseudo mask are used to calculate the loss back propagation to update the network model. After the iteration, the model parameters are saved.

[0052] (5) Model testing process: The model is loaded with the model parameters obtained in step (4), a new image is taken out as the support image, and the remaining images are used as query images. The support image is sent to the model to obtain the prototype, and the prototype is used to segment all query images to obtain the segmentation results.

[0053] like Figure 2 As shown in Figure 2, the steps for correcting the rough pseudo mask obtained by pre-segmentation in step (2) are as follows: Create a Boolean array B of the same size as the pseudo mask to track the visited pixels, and initialize the size of the largest connected area in the mask to C = 0. Traverse each pixel in the pseudo mask, and for unvisited pixels marked as 1, perform a breadth-first search on the entire pseudo mask starting from the pixel. If the connected area found is larger than the current maximum value C max , then update the maximum connected area and record the pixel point in B. After traversing all pixels, retain the maximum connected area C max , create a new array M of the same size as the rough pseudo mask, and put the maximum connected area C in M max The corresponding position is marked as 1 to obtain the corrected pseudo mask.

[0054] like Figure 2 As shown, the step (3) of using the multi-head attention model to transfer feature information is as follows: first, the input support image and query image are processed through the feature extraction network respectively to obtain their respective feature representations. This process can be expressed as:

[0055] F s =f θ (I s )

[0056] F q =f θ (I q )

[0057] Among them, I s and I q represent support images and query images respectively, and f θ Represents the mapping function of the feature extraction network, F s and F q are the extracted support features and query features. Next, we query the feature F q The self-attention mechanism allows the model to dynamically focus on other parts of the feature when processing the query feature, thereby capturing the dependencies between elements. Its calculation formula is as follows:

[0058] Q = Attention(F q W q Q ,F q W q K ,F q W q V )

[0059]

[0060] Where W q Q 、W q K 、W q V is for query feature F q The linear transformation matrix is ​​used to generate the query (Q = F q W q Q ), key (K=F q W q K ) and value (V=F q W q V ) vector. The Attention function is the core calculation process of the self-attention mechanism, and its output Q q represents the enhanced query features, K T and d k Denote the transpose and dimension of the key vector, respectively. Next, we calculate the similarity between the key (K) of the support feature and the query (Q) of the query feature, and apply the softmax function for normalization to generate the attention weights. Here, we use softmax (dim = -2) to ensure that the sum of the similarity scores of each column of the result (corresponding to each support pixel) is 1, allowing the support features to aggregate toward the query features. Subsequently, we multiply the attention weights with the pseudo mask of the support image to filter out irrelevant areas. Finally, the filtered attention weights are multiplied by the value (V) of the support image to achieve feature aggregation. This process can be expressed as:

[0061]

[0062] F' q =attn·V s

[0063] in, Indicates the transposition of the support key, ⊙ indicates element-wise multiplication, M s is a pseudo mask of the support image, used to filter out irrelevant areas, attn is the intermediate result, F q ' is the result of aggregating the support feature information into the query feature, and the subscripts q and s represent the information from the query image and the support image, respectively.

[0064] Finally, we feed the aggregated query features into a module consisting of a residual connection, a multi-layer perceptron, and another residual connection to obtain the final aggregated query features. This process efficiently aggregates key feature information from the support set to the query set through a one-way information transfer strategy, avoiding the problem of information misleading caused by inaccurate prediction masks of the query image. The specific formula is as follows:

[0065] O q =LN(attn·V+Q)

[0066] O out =LN(MLP(O q )+O q )

[0067] Among them, LN represents layer normalization, MLP represents multi-layer perceptron, and O q represents the intermediate query feature, O out is the final query feature.

[0068] like Figure 2 As shown, the query image foreground-background contrast learning steps in step (3) are as follows:

[0069] First, we use the Masked Average Pooling (MAP) technique to extract support prototypes representing foreground features from the labeled support images. s The calculation formula is as follows:

[0070] P s =MAP(F s ,M s )

[0071] Next, we use the cosine similarity between these support prototypes and query features to generate the query mask and extract the query prototype P q The calculation of cosine similarity and the extraction of query prototypes help us preliminarily screen out the areas in the query image that may be related to the foreground. In order to distinguish the complex background area, we first invert the support image mask to obtain Then take out the background features corresponding to the support features and express them as F s b The Voronoi-based method is used to Divide and obtain the background area mask set Here B n represents the nth background region, R is the number of regions, and is set to 150. The formula for extracting different region representations is as follows:

[0072]

[0073] Among them, P n represents the prototype of the nth region, H and W represent The height and width, and B n,i Denote the support feature and background region corresponding to the i-th pixel, respectively. In our contrastive learning framework, the query prototype serves as the anchor, the support prototype as the positive sample, and multiple partitioned prototypes of the background region constitute the negative sample set. Each negative sample represents a different part of the background region, providing rich contextual information to the model.

[0074] like Figure 2 As shown in step (4), the model loss includes the use of cross entropy loss to directly measure the similarity between the predicted probability distribution and the true mask distribution. For each pixel, its cross entropy loss is The calculation formula is:

[0075]

[0076] Among them, M q represents the ground-truth mask of the query image, and denote the background and foreground masks of the query image predicted by the model, respectively.

[0077] In addition, a threshold loss is added to the model Where T is the segmentation threshold and α is the adjustment parameter, which is set to 20 by default. By minimizing the learning threshold, the foreground features are encouraged to be more compact in the embedding space. In addition, the foreground-background contrast loss is proposed. The segmentation performance of the model is optimized by maximizing the similarity between the query prototype and the support prototype while minimizing the similarity with the background negative sample. The mathematical expression is:

[0078]

[0079] Among them, cos(·) is the cosine similarity function. It is worth noting that the corrected mask background area occupies the majority, so The loss can effectively use a large number of negative samples to optimize the model. Add up to get the total loss The calculation formula is as follows:

[0080] like Figure 3As shown in Figure 5, in step (5), the model testing process is as follows: In the testing phase, a set of medical images that have not been seen in the training process are used for segmentation. The test image is divided into a support image and a query image, where the support image contains only one sample with accurate annotations. The support and query features are extracted using the model obtained in the training phase, and then the support image features are transferred to the query image features using the attention mechanism. The support image pseudo mask M is used to extract the support and query features. s From the support feature F s Extract the supporting prototype P s Finally, with the help of the extracted supporting feature prototype, the query feature F q Segmentation is performed to obtain the segmentation results. Specifically, the attention mechanism enables the model to accurately focus on the key features in the image, while contrastive learning strengthens the model's feature learning and generalization capabilities. The synergy between the two greatly improves the model's representation capabilities. Since the trained model has fully learned the ability to extract key features from a small amount of data and generalize effectively, when faced with new data in the test phase, the model only needs to rely on one labeled sample to achieve accurate segmentation of the target organ. The detailed segmentation results are as follows: Figure 4 shown.

[0081] It should be noted that the above embodiments are not intended to limit the scope of protection of the present invention, and equivalent changes or substitutions made on the basis of the above technical solutions fall within the scope of protection of the claims of the present invention.

Claims

1. A small sample medical image segmentation method based on self-attention mechanism and contrastive learning, characterized by: The method comprises the following steps: (1) MRI data preprocessing and pre-segmentation: To adapt to the model input, the image data is reformatted into two dimensions and resized to 256 × 256 pixels. Each two-dimensional slice is replicated three times in the channel dimension and pre-segmented using an unsupervised image segmentation algorithm to obtain a rough pseudo mask. (2) Acquisition of pseudo mask: Correct the rough pseudo mask obtained by pre-segmentation and remove the erroneous areas to obtain the pseudo mask; (3) Constructing a deep learning network model: The feature extraction network used is ResNet-101 pre-trained on the MS-COCO dataset, where the classifier is replaced by a 1×1 convolutional layer. A multi-head attention module is added to the back of the model to transfer information between features and build contrastive learning between the foreground and background of the query image; (4) Model training process: The total number of iterations of the model is preset. In each iteration, two MRI images are taken out, which are called support images and query images respectively. The two images are fed into the network model to extract features. The prototype of the support feature is extracted according to the pseudo mask of the support image. The query feature is segmented using the support prototype. The query feature segmentation result and the query image pseudo mask are used to calculate the loss back propagation to update the network model. After the iteration, the model parameters are saved. (5) Model testing process: The model is loaded with the model parameters obtained in step (4), a new image is taken as the support image, and the remaining images are used as query images. The support image is fed into the model to obtain the prototype, and the prototype is used to segment all query images to obtain the segmentation results; Among them, the steps of using the multi-head attention module to transfer feature information in step (3) are as follows: first, the input support image and query image are processed through the feature extraction network respectively to obtain their respective feature representations. This process is expressed as: F s =f θ (I s ) F q =f θ (I q ) Among them, I s and I q represent support images and query images respectively, and f θ Represents the mapping function of the feature extraction network, F s and F q The extracted support features and query features are then used to extract the query features F. q The self-attention mechanism is fed into the self-attention for encoding. It allows the model to dynamically focus on other parts of the feature when processing the query feature, thereby capturing the dependencies between elements. The calculation formula is as follows: in is for query feature F q The linear transformation matrix is ​​used to generate the query key Sum vector, the Attention function is the core calculation process of the self-attention mechanism, and its output Q q represents the enhanced query features, K T and d k Represent the transpose and dimension of the key vector respectively. Next, the similarity between the key K of the support feature and the query Q of the query feature is calculated and normalized by applying the softmax function to generate the attention weight. Softmax(dim=-2) is used to ensure that the sum of the similarity scores of each column of the result is 1, thereby allowing the support feature to be aggregated to the query feature. Subsequently, the attention weight is multiplied by the pseudo mask of the support image to filter out irrelevant areas. Finally, the filtered attention weight is multiplied by the value V of the support image to achieve feature aggregation. This process is expressed as: F' q =attn·V s in, represents the transposition of the support key, ⊙ represents element-wise multiplication, M s is a pseudo mask of the support image, used to filter out irrelevant areas, attn is the intermediate result, F' q It is the result of aggregating the support feature information into the query feature. The subscripts q and s represent the information from the query image and the support image respectively. Finally, the aggregated query features are fed into a module consisting of a residual connection, a multi-layer perceptron, and another residual connection to obtain the final aggregated query features. This process efficiently aggregates key feature information from the support set to the query set through a one-way information transfer strategy, avoiding the information misleading problem caused by inaccurate query image prediction masks. The specific formula is as follows: O q =LN(attn·V+Q) O out =LN(MLP(O q )+O q ) Among them, LN represents layer normalization, MLP represents multi-layer perceptron, and O q represents the intermediate query feature, O out is the final query feature; The steps for learning the foreground-background contrast of the query image in step (3) are as follows: First, Masked Average Pooling (MAP) technology is used to extract the support prototype representing the foreground feature from the labeled support image. s The calculation formula is as follows: P s =MAP(F s ,M s ) Then, the cosine similarity between these support prototypes and query features is used to generate the query mask, and then the query prototype P is extracted. q , the calculation of cosine similarity and the extraction process of query prototype, preliminarily screen out the areas in the query image that may be related to the foreground. In order to distinguish the complex background area, the support image mask is first inverted to obtain Then the background features corresponding to the support features are expressed as The Voronoi-based method is used to Divide and obtain the background area mask set Here B n represents the nth background region, R is the number of regions, which is set to 150. The formula for extracting different region representations is as follows: Among them, P n represents the prototype of the nth region, H and W represent The height and width, and B n,i They represent the support feature and background area corresponding to the i-th pixel respectively. Under the constructed contrastive learning framework, the query prototype is used as the anchor point, the support prototype is used as the positive sample, and multiple partition prototypes of the background area constitute the negative sample set. Each negative sample represents a different part of the background area, thereby providing rich background information for the model.

2. The small sample medical image segmentation method based on self-attention mechanism and contrastive learning according to claim 1, characterized in that: The steps for correcting the rough pseudo mask obtained by pre-segmentation in step (2) are as follows: create a Boolean array B of the same size as the pseudo mask to track the visited pixels, initialize the size of the largest connected area in the mask to C = 0, traverse each pixel in the pseudo mask, and for unvisited pixels and the pixel is marked as 1, perform a breadth-first search on the entire pseudo mask starting from the pixel. If the connected area found is larger than the current maximum value C max , then update the maximum connected area and record the pixel point in B. After traversing all pixels, retain the maximum connected area C max , create a new array M of the same size as the rough pseudo mask, and put the maximum connected area C in M max The corresponding position is marked as 1 to obtain the corrected pseudo mask.

3. The small sample medical image segmentation method based on self-attention mechanism and contrastive learning according to claim 2, characterized in that: In step (4), the model loss includes the use of cross entropy loss to directly measure the similarity between the predicted probability distribution and the true mask distribution. For each pixel, the cross entropy loss L ce The calculation formula is: Among them, M q represents the ground-truth mask of the query image, and denote the query image background and foreground masks predicted by the model, respectively, In addition, the threshold loss L is added to the model T =T / α, where T is the segmentation threshold and α is the adjustment parameter, which is set to 20 by default. By minimizing the learning threshold, the foreground features are encouraged to be more compact in the embedding space, and the foreground-background contrast loss L is proposed. FB , the segmentation performance of the model is optimized by maximizing the similarity between the query prototype and the support prototype while minimizing the similarity with the background negative sample. The mathematical expression is: Among them, cos(·) is the cosine similarity function, and the corrected mask background area occupies most of the area, so L FB Able to effectively use a large number of negative samples to optimize the model and reduce the loss L CE , L T , L FB Add up to get the total loss L total , the calculation formula is as follows Lower: L total =L CE +L T +L FB .

Images