Prototype learning based self-supervised instance adaptive medical image segmentation method

Abstract

This invention discloses a self-supervised instance-adaptive medical image segmentation method based on prototype learning. 1) Images from a publicly available 3D medical image segmentation dataset are preprocessed to obtain standardized images; 2) The standardized images are input into a 3Dbackbone; 3) An adaptive prototype learning module is constructed to generate the final segmentation prediction result, and the segmentation loss is calculated to supervise the model; 4) A self-supervised mask generation module is constructed to generate pseudo-labels; 5) A self-supervised filtering mechanism is constructed, and an auxiliary loss is introduced; 6) The overall target loss is calculated based on the segmentation loss and the auxiliary loss; 7) The model is trained based on the loss function from step 6). This invention introduces an SMG module to better model intra-class diversity and designs a self-supervised filtering mechanism to supervise randomly initialized target queries; furthermore, this invention can also utilize an IPL module to make adaptive adjustments for different inputs.

Description

Technical Field

[0001] This invention belongs to the field of three-dimensional medical image segmentation technology, specifically constructing a self-supervised instance adaptive medical image segmentation method based on prototype learning. Background Technology

[0002] Medical image segmentation (MIS) aims to divide an image into parts representing different semantic categories, playing a crucial role in cancer diagnosis, surgical robot navigation, and radiation therapy dosage control. Early MIS methods were primarily based on contours and traditional machine learning algorithms. However, with the rapid development of computer technology, most current MIS models are based on deep CNNs and transformer methods, such as Res-UNet, U-Net++, Swin-unet, UNETR, and UNETR++. These methods are highly effective in generating representative pixel-level embeddings; however, their fixed categories limit their ability to capture class diversity.

[0003] Prototype learning refers to learning a set of prototypes representing classes, mapping the pixel-level embeddings output by the model network to the corresponding prototype centers. Networks based on prototype learning have achieved significant results in few-shot or zero-shot learning due to their direct output layer and rapid generalization to new classes, such as Matching Networks and Relation Networks. However, while current image segmentation methods based on prototype learning have made significant improvements, primarily generating segmentation masks by comparing pixels with prototypes, they still lack flexibility in adapting prototypes to different inputs and in capturing class diversity.

[0004] Given the limitations of both approaches, we find that introducing prototype learning into MIS can compensate for their respective shortcomings. This involves viewing MIS as a segmentation method that clusters image pixels into corresponding semantic categories. To implement this method, we must address two issues: how to enable the model to capture intra-class diversity and how to make the prototype adaptable to different inputs. Summary of the Invention

[0005] To address the problems in the background art, this invention provides a self-supervised instance-adaptive medical image segmentation method based on prototype learning. The purpose of this invention is to model class richness by constructing a self-supervised mask generation (SMG) module, where target queries represent classes and their subclasses, and the number of queries is greater than the number of classes; simultaneously, it constructs an adaptive prototype learning (IPL) module so that prototypes can adaptively generate based on different inputs.

[0006] The technical solution adopted in this invention is as follows, specifically including the following steps:

[0007] I. A Self-Supervised Instance Adaptive Medical Image Segmentation Method Based on Prototype Learning

[0008] Includes the following steps:

[0009] 1) The three-dimensional medical images are standardized through data preprocessing operations to reduce regional differences between different images; the data preprocessing operations include isotropic spacing and uniform intensity scale;

[0010] 2) Input the standardized 3D image obtained in step 1) into the 3D backbone, generate an intermediate feature map F through the encoder, and then generate pixel-level spatial embeddings F through the pixel decoder. o ;

[0011] 3) Generate pseudo-labels M using the self-supervised mask generation SMG module:

[0012] Initialize N target queries Q, and use the multi-scale (L-layer) features generated by the pixel decoder to guide the learning and updating of target queries (Queries) to generate pseudo-labels M;

[0013] 4) Input the intermediate feature map F generated in step 2) and the pseudo-label M generated in step 3) into the adaptive prototype learning IPL module to generate the final segmentation prediction result.

[0014] Step 3) specifically refers to:

[0015] 3.1) Initialize N target queries Q, where N>K, and K represents the number of target categories (when N>K, subclass information can be modeled);

[0016] 3.2) Update of target query Q:

[0017] Based on the image features of the current layer l The target query Q of the previous layer, i.e., layer l-1 l-1 Guide the target query Q at the current layer l l Update:

[0018] Where l∈L, L represents the number of multi-scale feature layers generated by the pixel decoder; the image features of the current l-th layer. Generated by the pixel decoder in step 2);

[0019] Similarity is calculated using matrix multiplication, and the target query Q with the highest similarity is selected. l The index value where (cluster) is located:

[0020]

[0021] in, This represents the matrix multiplication operation;

[0022] Then, through the multi-head self-attention mechanism and the image features generated by the current layer l, Update the target query Q at the current layer l. l :

[0023]

[0024] in, MHSA() represents the matrix multiplication operation and multi-head self-attention mechanism.

[0025] 3.3) Based on the updated query embeddings Q from step 3.2) l The image features generated by the current layer l The pixel is mapped to its nearest cluster, generating a pixel-to-cluster assignment map. The formula is expressed as follows:

[0026]

[0027] in, H represents height, W represents width, Z represents depth, and N represents channels; sim() represents similarity calculation; clustering is performed in step 3.2) after updating the target query Q. l ;

[0028] 3.4) Based on the updated query embedding Q from step 3.2) l This is achieved by mapping the clusters to their corresponding categories using a multilayer perceptron (MLP), generating a cluster-to-class projection matrix for the current layer l.

[0029]

[0030] in, K+1 indicates that in addition to the known categories, there is also a background class; MLP represents a linear layer.

[0031] 3.5) The product generated in step 3.3) And generated in step 3.4) The pseudo-label M of the l-th layer is generated by matrix multiplication. l ;

[0032]

[0033] Among them, M l ∈(K+1)×H×W×Z, where K+1 indicates that in addition to the known categories, there is also a background class. This represents the matrix multiplication operation;

[0034] 3.6) Repeat steps 3.2) to 3.5) until all multi-scale features of layer L are updated;

[0035] 3.7) Normalize the obtained pseudo-labels M using the softmax function; where M = {M 1 ,···,M l ,···,M L}

[0036] Step 4) specifically involves:

[0037] 4.1) Fuse the pseudo-label M generated in step 3) with the intermediate feature map F generated in step 2) to generate instance-specific prototype proposals (IPPs) P. i ;

[0038]

[0039] in, The expression represents element-wise multiplication, and GAP() represents global average pooling operation.

[0040] 4.2) Randomly initialize K common prototype suggestions P c (CPPs), This indicates that there are K categories in total;

[0041] The instance-specific prototype suggestions generated in step 4.1) With public prototype suggestions The components are merged to generate the final prototype G;

[0042]

[0043] Among them, g k ∈G, and g k ∈R d×1 d represents the embedding dimension; This indicates a concatenation operation; MLP() represents a fully connected layer.

[0044] 4.3) The prototype g generated in step 4.2) k Compared with the pixel-level spatial embedding features F generated in step 2), o The data is then fused to obtain the final segmentation prediction result.

[0045]

[0046] Where ρ represents the sigmoid activation function. This represents matrix multiplication.

[0047] This invention relates to a self-supervised instance-adaptive medical image segmentation method based on prototype learning:

[0048] A: By constructing a self-supervised filtering mechanism (SF), the N target queries Q randomly initialized in step 3) are guided to better model category information and avoid background bias; specifically:

[0049] Since the N target queries Q initialized in step 3) may cause background-based bias during model training without proper guidance, we introduce a self-supervised filtering mechanism (SF) to avoid this situation.

[0050] 5.1) Calculate the intersection of the mask of the (l-1)th layer cluster and the mask of each category, that is, calculate the overlap score (IoU) between the currently activated region and the foreground region of each category, which is the maximum ratio between the intersection of the cluster and the category and the related category region.

[0051]

[0052] in, A specified graph representing pixels to clusters The nth query, The pseudo-label M represents the previous layer, i.e., layer l-1. l-1 The kth category;

[0053] 5.2) Set the threshold using a warm-up method:

[0054]

[0055] Where epoch represents the epoch index during training. When the epoch is less than 50, the threshold increases linearly between 0.1 and 0.5. When the epoch is greater than 50, the threshold remains unchanged at 0.5.

[0056] Clusters with overlap scores greater than a threshold are used for updates, while clusters with overlap scores less than a threshold are filtered out.

[0057] B: By introducing an auxiliary loss to supervise the updated target query in step 3), the generated pseudo-labels M are ensured. l The quality; the auxiliary loss for:

[0058]

[0059] Among them, Y k,i , K represents the true label of the k-th class corresponding to the i-th pixel and the prediction result of the SMG model; K+1 indicates that the calculated analogy loss also includes a background class.

[0060] C: Calculate the segmentation loss based on the prediction results obtained in step 4), which is used to supervise the segmentation of the adaptive prototype learning module in step 4.

[0061] The formula for calculating the segmentation loss is as follows:

[0062]

[0063] Among them, Y k,i , Let represent the true label and IPL model prediction result for the k-th category corresponding to the i-th pixel, respectively. Here, the loss for the k categories is calculated.

[0064] D: Calculate the overall target loss based on the segmentation loss and auxiliary loss, and train the instance_segmentation model based on the overall target loss;

[0065] The instance_segmentation model includes a 3D backbone module, an SMG module, and an IPL module.

[0066] The overall target loss is the segmentation loss L. seg and auxiliary losses The sum:

[0067]

[0068] Here, α is a hyperparameter, set to 0.05.

[0069] II. A type of terminal

[0070] Including memory and processor;

[0071] The memory is used to store computer programs;

[0072] The processor is configured to implement the above-described self-supervised instance adaptive medical image segmentation method based on prototype learning when executing the computer program.

[0073] III. A computer-readable storage medium

[0074] The storage medium stores a computer program, which, when executed by a processor, implements the aforementioned self-supervised instance-adaptive medical image segmentation method based on prototype learning.

[0075] The beneficial effects of this invention are:

[0076] The instance_segmentation network model constructed in this invention mainly comprises an instance adaptive prototyping (IPL) module and a self-supervised mask generation (SMG) module, surpassing current state-of-the-art (SOTA) models and achieving the best segmentation results. The instance adaptive prototyping (IPL) module created in this invention allows the model to adaptively adjust the prototype based on the current input. Furthermore, the self-supervised mask generation (SMG) module designed in this invention can also model class diversity. The combination of these two features enables the model to achieve SOTA results on the abdominal multi-tissue segmentation dataset (BTCV), surpassing other existing models. Attached Figure Description

[0077] Figure 1 This is a flowchart of the method of the present invention;

[0078] Figure 2 This is a schematic diagram of the basic structure of the key module of this invention - the Adaptive Prototype Learning (IPL) module;

[0079] Figure 3 This is a schematic diagram of the basic structure of the key module of the present invention - the Self-Supervised Mask Generation (SMG) module. Detailed Implementation

[0080] The present invention will be further described in detail below with reference to the accompanying drawings and specific embodiments.

[0081] like Figure 1 As shown, the present invention provides a self-supervised instance-adaptive medical image segmentation method based on prototype learning, comprising the following steps:

[0082] Step 1): Standardize the 3D medical images through a series of data preprocessing operations to reduce regional differences between different images, such as isotropic spacing and uniform intensity scale.

[0083] Step 2): Take the image V∈R that has undergone data preprocessing in Step 1) H×W×ZThe encoder input to the 3D backbone generates intermediate feature maps. Then it is input into the pixel encoder to generate per-pixel embeddings F o ∈R H×W×Z×d ;

[0084] Step 3): As Figure 3 As described above, a self-supervised mask generation (SMG) module is constructed. The kMax decoder uses the multi-scale (L-layer) features generated by the pixel decoder to guide the learning and updating of the target Queries Q, as described in steps 3.1) to 3.6).

[0085] 3.1) Update of the target Queries Q:

[0086] First, initialize N target queries Q, based on the image features generated by the current layer l of the pixel decoder. Queries Q of the previous layer (l-1 layer) l-1 To guide the update of the current layer l target Queries Q, the index position of the cluster with the highest similarity is selected:

[0087]

[0088] in This represents the matrix multiplication operation;

[0089] Then, through the multi-head self-attention mechanism and the image features generated by the current layer l, Update the current layer l target Queries Q l :

[0090]

[0091] in MHSA() represents the matrix multiplication operation and multi-head self-attention mechanism.

[0092] 3.2) Based on the updated query embeddings Q from step 3.1) l The image features generated by layer l This maps pixels to their nearest cluster, generating a pixel-to-cluster assignment map.

[0093]

[0094] in

[0095] 3.3) Based on the updated query embeddings Q from step 3.1) l The cluster-to-class projection matrix of layer l is generated by mapping the data to the corresponding class using a multilayer perceptron (MLP).

[0096]

[0097] in K+1 indicates that in addition to the known categories, there is also a background category, and MLP represents a fully connected layer;

[0098] 3.4) The product generated in step 3.2) And generated in step 3.3) The pseudo-label M for the current layer l is generated through matrix multiplication. l ;

[0099]

[0100] Where M l ∈(K+1)×H×W×Z, This represents the matrix multiplication operation;

[0101] 3.5) Repeat steps 3.1) to 3.4) until all multi-scale features of layer L are updated;

[0102] 3.6) Use the softmax function to normalize the obtained pseudo-labels M;

[0103] Step 4): In step 3), N target queries are initialized. Without proper guidance during model training, the Q will cause background-based bias. Therefore, we introduce a self-supervised filtering mechanism (SF) to avoid this situation. The specific operation is as follows: Step 4.1) to Step 4.3).

[0104] 4.1) Calculate the intersection of the current cluster mask and the mask of each class, that is, calculate the overlap score (IoU) between the currently active region and the foreground region of each class, and generate the maximum ratio between the intersection of the cluster mapping to the class and the related cluster region.

[0105]

[0106] in This represents a pixel-to-cluster assignment map. The nth cluster, The pseudo-tag M represents the previous layer (l-1 layer). l-1 The kth category;

[0107] 4.2) Set a threshold to use clusters with IoU scores greater than the threshold for updates, and set the threshold using a warm-up approach;

[0108]

[0109] Where epoch represents the epoch index during training. When the epoch is less than 50, the threshold increases linearly between 0.1 and 0.5. When the epoch is greater than 50, the threshold remains unchanged at 0.5.

[0110] 4.3) Set a self-supervised auxiliary loss update Query to ensure that the generated pseudo-labels M l The quality;

[0111]

[0112] Where Y k,i , This represents the true label and prediction result corresponding to the i-th pixel. Here, the loss for K+1 categories is calculated, which includes one background class.

[0113] Step 5): As Figure 2 As described above, an adaptive prototype learning (IPL) module is constructed. Based on the intermediate feature map F generated in step 2) and the pseudo-label M generated by the SMG module in step 3), the IPL module (such as...) is used to... Figure 2 Generate the final segmentation prediction results, as detailed in steps 5.1) to 5.4).

[0114] 5.1) Use the pseudo-label M of the kth category generated in step 3) k The intermediate feature map F generated in step 2) is fused with the map to generate the instance-specific prototype proposal (IPP) for the k-th category.

[0115]

[0116] Where M k ∈M, The expression represents element-wise multiplication, and GAP() represents global average pooling operation.

[0117] 5.2) Randomly initialize a series of common prototype proposals (CPPs) This indicates that there are K categories in total, and the instance-specific prototype proposal generated in step 5.1) will be used. With public prototype proposal The components are merged to generate the final prototype G;

[0118]

[0119] Where g k ∈G, and g k ∈R d×1 , This indicates a concatenation operation; MLP() represents a fully connected layer.

[0120] 5.3) The prototype g generated in step 5.2) k Compared with the per-pixel embeddings F generated in step 2), o The data is then fused to obtain the final segmentation prediction result;

[0121]

[0122] Where ρ represents the sigmoid activation function. Indicates matrix multiplication;

[0123] 5.4) The prediction results obtained in step 5.3) are used to calculate the segmentation loss;

[0124]

[0125] Where Y k,i , These represent the true label and the model prediction result for the k-th category corresponding to the i-th pixel, respectively. Here, the loss for the k categories is calculated.

[0126] Step 6): The overall target loss is the segmentation loss L. seg and auxiliary losses The sum;

[0127]

[0128] Where α is a hyperparameter set to 0.05;

[0129] Step 7): Train the instance_segmentation network model, selecting the sum of the soft dice loss and the binary cross-entropy loss function as the final loss function. Backpropagation is used to adjust the weights and biases of the segmentation network model. Parameters are selected based on the Dice Similarity Coefficient (DSC), ultimately determining a set of parameters that maximizes the image segmentation accuracy (DSC score).

[0130] The image segmentation model of this invention can determine whether each pixel of a three-dimensional image belongs to its corresponding category, and it has a significant effect on multi-category image segmentation.

[0131] The performance of the instance_segmentation model of this invention and some state-of-the-art models on the current BTCV multi-organization segmentation dataset is as follows:

[0132]

[0133] Note:

[0134] (1) UNETR, Swin-UNETR, TransBTS, nnFormer and UNETR++ are models that have previously achieved state-of-the-art performance on the BTCV dataset, and instance_segmentation is the model of this invention.

[0135] (2) Spl, RKid, LKid, Gal, Eso, Liv, Sto, Aor, IVC, PSV, Pan, RAG, and LAG represent the thirteen abdominal tissues: spleen, right kidney, left kidney, gallbladder, liver, stomach, aorta, pancreas, esophagus, inferior vena cava, portal and splenic veins, right adrenal gland, and left adrenal gland.

[0136] (3) The values ​​in the table are Dice similarity coefficients (DSC). The larger the value, the better the segmentation effect.

[0137] (4) Avg: represents the average DSC score of the thirteen tissues. The larger the value, the better the segmentation effect.

[0138] (5) The values ​​in bold black in the table indicate the best segmentation performance;

[0139] (6) Compared with the state-of-the-art models UNETR, Swin-UNETR, TransBTS, nnFormer, and UNETR++, the average DSC(Avg) of the segmentation index on the same dataset in this invention is improved by 11%, 6.56%, 5.69%, 5.38%, and 3.72%, respectively.

Claims

1. A self-supervised instance-adaptive medical image segmentation method based on prototype learning, characterized in that, Includes the following steps: 1) Standardize the 3D medical images through data preprocessing to reduce regional differences between different images; 2) Input the standardized 3D image obtained in step 1) into the 3D backbone, and generate an intermediate feature map through the encoder. Then, the pixel decoder generates pixel-level spatial embedding features. ; 3) Generating pseudo-labels through the self-supervised mask generation SMG module. : Initialize N target queries The kMax decoder uses multi-scale features generated by the pixel decoder to guide the learning and updating of the target query, generating pseudo-labels. ; 4) The intermediate feature map generated in step 2) And the pseudo-tags generated in step 3) Input the adaptive prototype learning IPL module to generate the final segmentation prediction result; Step 4) specifically refers to: 4.1) The pseudo-tags generated in step 3) Compared with the intermediate feature map generated in step 2), Fusion, generating instance-specific prototype suggestions ; in, , }, This indicates element-wise multiplication. This indicates a global average pooling operation; 4.2) Random initialization A public prototype suggestion , , , indicating shared ownership One category; The instance-specific prototype suggestions generated in step 4.1) With public prototype suggestions The final prototype is generated through fusion. ; in, ,and , Indicates the embedding dimension; ; Indicates a splicing operation; This represents a fully connected layer; 4.3) The prototype generated in step 4.2) Compared with the pixel-level spatial embedding features generated in step 2), The data is then fused to obtain the final segmentation prediction result. : in, This represents the sigmoid activation function. This represents matrix multiplication.

2. The medical image segmentation method based on prototype learning and self-supervised instance adaptation according to claim 1, characterized in that, Step 3) specifically refers to: 3.1) Initialization Target query ,in , Indicates the number of target categories; 3.2) Update of target query Q: According to the current Image features of layers and the previous layer Target query at level -1 Guiding the present Target query of layer Update: in, , This indicates the number of multi-scale feature layers generated by the pixel decoder; Similarity is calculated using matrix multiplication to filter out the target queries with the highest similarity. Index value: in, This represents the matrix multiplication operation; Then through the multi-head self-attention mechanism and the current Image features generated by layers Update current Target query of layer : in, This represents the matrix multiplication operation. This indicates a multi-head self-attention mechanism; 3.3) Query embedding updated based on step 3.2) , will the current Image features generated by layers Mapping to its nearest cluster generates a pixel-to-cluster map. The formula is expressed as follows: ) in, , Indicates altitude, Indicates width, Indicates depth, Indicates a channel; () indicates similarity calculation; clustering is the updated target query in step 3.2). ; 3.4) Query embedding updated based on step 3.2) It maps the current data to the corresponding category using a multilayer perceptron (MLP) to generate the current data. Layers are mapping matrices from clusters to categories. : in, , This indicates that in addition to the known categories, there is also a background class; MLP stands for Linear Layer. 3.5) The product generated in step 3.3) And the generated in step 3.4) Generate the first by matrix multiplication Pseudo-tags of layers ; ； in, , This indicates that in addition to the known categories, there is also a background class. This represents the matrix multiplication operation; 3.6) Repeat steps 3.2) to 3.5) until... All multi-scale features have been updated. 3.7) Obtain the pseudo-labels Normalization is performed using the softmax function; where... ={ } 3. The medical image segmentation method based on prototype learning and self-supervised instance adaptation according to claim 2, characterized in that, Step 3) involves randomly initializing N target queries by constructing a self-supervised filtering mechanism. ; By introducing an auxiliary loss to supervise the updated target query in step 3), the generated pseudo-labels are ensured. The quality; The segmentation loss is calculated based on the prediction results obtained in step 4) and used to supervise the segmentation of the adaptive prototype learning module in step 4). The overall target loss is calculated based on the segmentation loss and auxiliary loss, and the instance_segmentation model is trained based on the overall target loss; the instance_segmentation model includes a 3D backbone module, an SMG module and an IPL module.

4. The medical image segmentation method based on prototype learning and self-supervised instance adaptation according to claim 3, characterized in that, The specific method for constructing a self-supervised filtering mechanism is as follows: 5.1) Calculate the first... The intersection of the layer cluster mask and the mask of each category is calculated by calculating the overlap score between the currently activated region and the foreground region of each category, which is the maximum ratio between the intersection of the cluster and the category and the related category region. in, A specified graph representing pixels to clusters The One query, Indicates the previous layer, i.e. Pseudo-tags of layers The One category; 5.2) Set the threshold using a warm-up method. : in, Indicating the training process Index, when When the threshold is less than 50, it increases linearly between 0.1 and 0.

5. When the value is greater than 50, the threshold remains unchanged at 0.5; Clusters with overlap scores greater than a threshold are used for updates, while clusters with overlap scores less than a threshold are filtered out.

5. The medical image segmentation method based on prototype learning and self-supervised instance adaptation according to claim 3, characterized in that, The auxiliary loss for: in, , Indicates the first The corresponding pixel point True labels and SMG model prediction results for each category; The calculated analogy loss also includes a background class.

6. The medical image segmentation method based on prototype learning and self-supervised instance adaptation according to claim 3, characterized in that, Calculate the segmentation loss based on the prediction results obtained in step 4): in, , They represent the first The corresponding pixel point The true labels for each category and the prediction results from the IPL model.

7. The medical image segmentation method based on prototype learning and self-supervised instance adaptation according to claim 3, characterized in that, The overall target loss is the segmentation loss. and auxiliary losses The sum: in, It is a hyperparameter, set to 0.

05.

8. A terminal, characterized in that, Including memory and processor; The memory is used to store computer programs; The processor is configured to implement, when executing the computer program, the self-supervised instance adaptive medical image segmentation method based on prototype learning as described in any one of claims 1 to 7.

9. A computer-readable storage medium, characterized in that, The storage medium stores a computer program that, when executed by a processor, implements the self-supervised instance adaptive medical image segmentation method based on prototype learning as described in any one of claims 1 to 7.

Images