Medical image segmentation multi-model aggregation method based on learnable token

Abstract

The invention relates to a medical image segmentation multi-model aggregation method based on a learnable token, and the method comprises the following steps: S1, freezing the parameters of teacher models, inputting CT image data into each teacher model, and obtaining a pseudo tag of each single organ segmentation model; s2, taking the pseudo labels and the real new-class labels as a pseudo label set; s3, inputting the CT image data into the student model, calculating an overall loss function based on the output of the student model, and optimizing parameters of the student model based on the loss function until iterative convergence to obtain a medical image segmentation model; and S4, obtaining an actual segmentation task and a CT image, and inputting the CT image into the medical image segmentation model to obtain a segmentation result corresponding to the segmentation task. Compared with the prior art, the method has the advantages of improving the segmentation effect of the aggregation task of the multi-organ segmentation model of the medical image and the like.

Description

Technical Field

[0001] The present invention relates to the technical field of medical images, and in particular to a medical image segmentation multi-model aggregation method based on learnable tokens. Background Art

[0002] Multi-task learning (MTL) uses a shared backbone feature extraction network to store knowledge of multiple tasks at the same time. Since there is no need to save parameter copies for each task and multiple tasks can be inferred at the same time, MTL greatly improves the storage and inference efficiency of the model.

[0003] Existing MTL methods can be roughly divided into two categories: online joint training methods and passive data contact model aggregation methods. Online joint training relies on the paradigm of collecting multi-task data in advance and then conducting joint training, that is, using labeled data of multiple tasks to train a feature extraction network at the same time. Since the gradient return of each task will participate in the update of the parameters of the backbone feature extraction network, the final feature extraction network can effectively extract features from multiple task data. However, this type of method needs to store the training data of each task, and when a new task arrives, all old task models need to be retrained, and the scalability is poor.

[0004] Model aggregation methods can be divided into two types according to the incremental steps: continuous learning and model merging. The former expands the number of processing tasks through multi-step increments; while the model merging method merges all task models in a single step to perform MTL. Continuous learning, also referred to as incremental learning or lifelong learning, is a method that focuses on acquiring knowledge in sequence. As a cutting-edge hotspot in the field of deep learning, continuous learning methods break through the constraints of one-time learning and greatly improve the generalization ability of neural network-based models. In contrast, traditional machine learning methods are usually built on closed sets, that is, they can only process a fixed number of predefined classes, and all data need to be presented to the model in a single-step training. However, models often face the challenge of continuous data increment in applicable scenarios. Therefore, how to make the model continuously adapt to new data or tasks constitutes a common challenge. However, the typical model update design retrains new data or applies transfer learning technology, which will cause catastrophic forgetting. That is, when the gradient backpropagation-based method learns new task knowledge, the performance of the model on old tasks will drop sharply. The primary goal of continuous learning is to achieve the best balance in the stability-plasticity dilemma under the constraints of limited computing and storage resources, where stability refers to the ability to retain source domain knowledge and plasticity refers to the ability to integrate target domain knowledge.

[0005] Model merging emphasizes the direct merging of multiple carefully trained or fine-tuned proprietary models, as no additional training is required or only a small amount of unlabeled data is required for training. In addition to executing MTL, model merging can also solve the performance of subtasks through model knowledge fusion. In the existing technology, for the multi-organ segmentation model aggregation task, a semi-supervised method that only uses the pseudo-label part and a knowledge distillation method that relies on the output layer can be used. However, the semi-supervised method that only uses the pseudo-label part has a serious class imbalance in the dataset because the label part only contains new organoid annotations, resulting in poor segmentation results. The knowledge distillation method that relies on the output layer cannot fully transfer knowledge, and the insufficient knowledge transfer also leads to poor segmentation results for the multi-organ segmentation model aggregation task. Summary of the invention

[0006] The purpose of the present invention is to provide a medical image segmentation multi-model aggregation method based on learnable tokens in order to improve the segmentation effect of the multi-organ segmentation model aggregation task of medical images.

[0007] The purpose of the present invention can be achieved by the following technical solutions:

[0008] A multi-model aggregation method for medical image segmentation based on learnable tokens, characterized in that the method comprises the following steps:

[0009] S1. Collect the single organ segmentation models obtained from T training sets As the teacher model, the t-th training set corresponds to the t-th teacher model, and the t-th teacher model includes the encoder and decoder Each teacher model performs a segmentation task, the t-th teacher model can only identify the foreground organ category t corresponding to the segmentation task, and the remaining categories are identified as background categories, the parameters of the teacher model are frozen, the CT image data are input into each teacher model, the logical value output by each teacher model is converted into a category probability, and the pseudo label of each single organ segmentation model is obtained based on the category probability;

[0010] S2. Sort the pseudo labels based on the segmentation capabilities of different models in the single-organ segmentation model set, integrate the sorted pseudo labels to obtain a teacher model pseudo label set, and use the teacher model pseudo label set and the true new class label together as a pseudo label set;

[0011] S3. Transform the output of the student model so that the output vectors of the student model and the teacher model match, perform distillation of the student model and the teacher model, and input the CT image data into the student model. The student model includes an encoder and decoder The encoder A learnable token is set for each feature layer in the training set, and an overall loss function is calculated based on the output of the student model. The overall loss function includes a distillation part loss, a feature distillation / alignment part loss, and a segmentation part loss. The parameters of the student model are optimized based on the loss function until the iteration converges to obtain a medical image segmentation model.

[0012] S4. Obtain an actual segmentation task and a CT image, and input the CT image into a medical image segmentation model to obtain a segmentation result corresponding to the segmentation task.

[0013] Furthermore, the category probability output by the teacher model is:

[0014]

[0015] Among them, t represents the t-th teacher model, It represents the output logic value of the t-th teacher model for identifying organ category i. i is 1, which means that i is identified as the foreground organ category corresponding to the t-th teacher. i is 0, which means that i is identified as the background category. represents the output logic value of the t-th teacher model for organ category k. k is 1, which means that k is identified as a foreground organ category, and k is 0, which means that k is identified as a background category. T represents the total number of teacher models. Represents the category probability of organ category k output by the tth teacher model.

[0016] Furthermore, the pseudo label of each single organ segmentation model is:

[0017]

[0018] Among them, t represents the t-th teacher model, M t represents the pseudo label value of the t-th teacher model, represents the category probability of the organ category k identified as the background category output by the t-th teacher, Represents the category probability of the t-th teacher output identifying the organ category k as the foreground organ category.

[0019] Furthermore, the output probability of the student model matching each teacher model is:

[0020]

[0021] in, represents the output logical value of the student model S for identifying the foreground organ category t identified by the t-th teacher model, p S (t) represents the probability of the foreground organ category t being identified as the foreground organ category by the output of the student model S, that is, the output of the student model when it is not transformed, 1-p S(t) represents the category probability output by the student model S for identifying the foreground organ category t as the background category.

[0022] Furthermore, the segmentation loss is:

[0023] L pseudo =L ce (p S ,g pseudo )

[0024] Among them, L pseudo represents the segmentation loss, L ce represents the cross entropy loss function, g pseudo is the pseudo label set, p S Represents the output of the student model when it is not transformed.

[0025] Furthermore, the pseudo label set is:

[0026] g pseudo =[M T ,M S ]

[0027] Among them, M T is the pseudo label set of T single organ segmentation models, M S is the real new class label.

[0028] Further, the distillation part loss is:

[0029]

[0030] Where N represents the number of training set samples, KL represents KL divergence, represents the teacher model corresponding to the t-th training set, represents the student model, x i represents the i-th training set sample.

[0031] Furthermore, the feature distillation / alignment part loss is:

[0032]

[0033] Where N represents the number of training set samples. represents the encoder of the teacher model corresponding to the t-th training set, represents the encoder of the student model, φ t Represents the feature conversion function, that is, by setting additional learnable tokens, the features encoded by the student model are converted to the feature space corresponding to the tth teacher model encoding.

[0034] Furthermore, the overall loss function is specifically:

[0035] L total =L pseudo +λ kd L kd +λ feat L feat

[0036] Among them, L total is the overall loss function, L pseudo To split the partial loss, L kd Represents the distillation loss, L feat represents the feature distillation / alignment loss, λ kd represents the weight of the loss term of the output layer distillation part, λ feat Represents the weight of the feature distillation / alignment loss term.

[0037] Furthermore, the pseudo labels are sorted based on the segmentation capabilities of different models in the single organ segmentation model set, and the sorted pseudo labels are integrated to obtain the teacher model pseudo label set. The specific steps are:

[0038] For each teacher pseudo label M f …M T Each pixel on the segmentation map is sorted according to the confidence level corresponding to each segmentation map. Finally, the T pseudo labels are merged to obtain the teacher model pseudo label set M. T .

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

[0040] The present invention synthesizes the pseudo-labels of the outputs of multiple teacher models to obtain a pseudo-label map, and combines it with the real new class label map to form a fully supervised label map of the student model to calculate the segmentation part loss. For the intermediate features of the student model and multiple teacher models, the student model feature distribution is transformed using task-specific learnable tokens, and feature layer distillation is achieved by minimizing the difference between the converted feature distribution and the teacher model feature distribution. The KL divergence of the student model output probability vector and the output probability vectors of different teacher models after background probability conversion is calculated to obtain the output layer distillation loss. The entire network adjusts the model weights by minimizing the segmentation part, feature layer distillation loss and output layer distillation loss, inherits the domain knowledge of multiple proprietary models, and achieves the purpose of processing multiple participating tasks at the same time while outputting more accurate segmentation results for each segmentation task. It has good scalability without having to access the training data of previous multiple models. BRIEF DESCRIPTION OF THE DRAWINGS

[0041] Figure 1 This is a model training framework diagram of the present invention. DETAILED DESCRIPTION

[0042] The present invention is described in detail below in conjunction with the accompanying drawings and specific embodiments. This embodiment is implemented based on the technical solution of the present invention, and provides a detailed implementation method and specific operation process, but the protection scope of the present invention is not limited to the following embodiments.

[0043] The present invention aims to study the model aggregation problem in the field of medical image segmentation. Based on the learnable token method, a multi-model aggregation method for medical image segmentation based on learnable tokens is proposed. The method adopts a novel feature layer and output layer joint distillation model aggregation framework, which can align the feature distribution and output distribution of multiple models, thereby realizing knowledge fusion more efficiently.

[0044] The method proposed by the present invention comprises the following steps:

[0045] S1. Collect the single organ segmentation models obtained from T training sets As the teacher model, the t-th training set corresponds to the t-th teacher model, and the t-th teacher model includes the encoder and decoder Each teacher model performs a segmentation task, the t-th teacher model can only identify the foreground organ category t corresponding to the segmentation task, and the remaining categories are identified as background categories, the parameters of the teacher model are frozen, the CT image data are input into each teacher model, the logical value output by each teacher model is converted into a category probability, and the pseudo label of each single organ segmentation model is obtained based on the category probability;

[0046] S2. Sort the pseudo labels based on the segmentation capabilities of different models in the single-organ segmentation model set, integrate the sorted pseudo labels to obtain a teacher model pseudo label set, and use the teacher model pseudo label set and the true new class label together as a pseudo label set;

[0047] S3. Transform the output of the student model so that the output vectors of the student model and the teacher model match, perform distillation of the student model and the teacher model, and input the CT image data into the student model. The student model includes an encoder and decoder The encoder A learnable token is set for each feature layer in the training set, and an overall loss function is calculated based on the output of the student model. The overall loss function includes a distillation part loss, a feature distillation / alignment part loss, and a segmentation part loss. The parameters of the student model are optimized based on the loss function until the iteration converges to obtain a medical image segmentation model.

[0048] S4. Obtain an actual segmentation task and a CT image, and input the CT image into a medical image segmentation model to obtain a segmentation result corresponding to the segmentation task.

[0049] Defining Neural Networks parameter Input Image Output Where n represents the number of parameters, d represents the dimension of the input data, and c represents the dimension of the output result. The model set obtained by fine-tuning is For the segmentation task, Usually the encoder and decoder It consists of two parts, namely The purpose of model merging is to merge the above model parameters to obtain a unified model parameter θ um , The aforementioned T tasks can be executed simultaneously.

[0050] The model aggregation scenario of the present invention combines the previous continuous learning and model merging scenarios and expands to process T+1 tasks. That is, when training the T+1th model, the training data of the task However, it is impossible to obtain the existing T model training data. During or after training, Merge It has the ability to solve T+1 tasks at the same time. To avoid subsequent ambiguity, and represents the same model. It is worth noting that the scenario setting of this study covers single-step continuous learning and the recent model merging scenario, which has stronger practical application value. When we train a new model on a new task dataset, we expect the model to inherit all the task capabilities that have been mastered. In this way, the final model can not only solve new tasks, but also handle all previously learned tasks. In addition, we also consider the fusion of model knowledge, and model aggregation may improve the performance of each task through knowledge growth.

[0051] In the present invention, the image data is first input into multiple teacher models and student models The teacher model weights are frozen. After forward propagation and softmax processing at the output layer, the category probability value of each pixel in the image is obtained. The outputs of multiple teacher models are synthesized through pseudo-labels to obtain a pseudo-label map, which is combined with the real new class label map to form a fully supervised label map of the student model to calculate the classification loss. For the intermediate features of the student model and multiple teacher models, the task-specific learnable tokens are used to transform the student model feature distribution, and the feature layer distillation is achieved by minimizing the difference between the transformed feature distribution and the teacher model feature distribution. The KL divergence of the student model output probability vector and the output probability vectors of different teacher models after background probability conversion is calculated to obtain the output layer distillation loss. The entire network adjusts the model weights by minimizing the classification loss, feature layer distillation loss, and output layer distillation loss, inherits the domain knowledge of multiple proprietary models, and achieves the purpose of being able to handle multiple participating tasks at the same time.

[0052] like Figure 1 As shown, the present invention inherits the model knowledge by aligning the unified model and the output of each proprietary model, which is generally achieved through the knowledge distillation method. Considering multiple proprietary models as teacher models and the unified model as the student model, the knowledge transfer can be achieved by minimizing the KL divergence (Kullback-Leibler Divergence) between the distributions of the two. The KL divergence calculation formula is:

[0053]

[0054] Where P represents the probability distribution of the teacher model output, and Q represents the probability distribution of the student model output.

[0055] Since different segmentation tasks output different segmentation results and different tasks have different semantic understandings of the image background area, certain processing is required to perform correct output layer knowledge distillation. The pseudo-label generation and background probability conversion methods used in the present invention are described in detail below.

[0056] In S1, a pseudo label set is obtained based on the output of the teacher model. First, the probability output of the single organ segmentation model is obtained: under the premise of freezing the network parameters of multiple teacher models, the CT image after data enhancement is input, and the output logical value (Logits) is processed by formula (2) to obtain the foreground category and background category probabilities, and then the category output of each single organ segmentation model is obtained by formula (3).

[0057]

[0058] Among them, M t represents the pseudo label map corresponding to the teacher model t, Represents the logical output of its model, k represents the segmented organ category, 1 represents the foreground, and 0 represents the background.

[0059] Merge multiple organ categories: Sorting the order of integrating pseudo labels by the segmentation capabilities of multiple single-organ segmentation models can often achieve better results. That is, the model with strong segmentation ability can judge the foreground more accurately, and the pseudo labels obtained in this way are more certain.

[0060] M T =[M 1 ,,…,M T ]#(4)

[0061] [·] indicates sorting by segmentation ability, that is, the pixel category is preferentially assigned to the category predicted by the model with strong segmentation ability. Specifically, for each segmentation map M 1 …M T Each pixel on the segmentation map is sorted according to the confidence level corresponding to each segmentation map, and finally T segmentation maps are merged to obtain M T . M S It is the annotated segmentation map contained in the current dataset (student model). When merging multiple teacher models, it contains a labeled dataset M S To train the new categories of the student model, this dataset does not contain any annotation information for training the teacher model. S That is, it represents the real new class label.

[0062] Finally, the new dataset is labeled with the new class M S and the above M T To integrate, that is, g pseudo =[M T ,M S ] Get the final pseudo label set g pseudo After the pseudo-label synthesis, the label space of the unified student model S is complete, and it can be used as a supervision signal to train the fully supervised network. Then, the classification loss is calculated by the following formula:

[0063] L pseudo =L ce (p S ,g pseudo )#(5)

[0064] Among them, L ce Represents CrossEntropyLoss, a commonly used loss function for classification tasks, which has been built into Pytorch. S represents the output probability vector of the student model, and its specific definition is shown in formula (7).

[0065] Given a single organ segmentation model, for each pixel output, the teacher model output, that is, the predicted output probability for category k is defined as:

[0066]

[0067] Where t represents the t-th teacher model, Represents the output logical value of the teacher model for class i, 1 represents the foreground organ category, and 0 represents the background category.

[0068] For the student model output, the probability of its predicted category k is expressed as:

[0069]

[0070] Since the student model needs to be able to segment the organs of the teacher model, its categories include the first T teacher categories, the current new category T+1, and the background category 0. From equations (6) and (7), it can be seen that the dimensions of the teacher model output and the student model output vector are inconsistent, and knowledge distillation cannot be performed directly. In addition, since there is a background semantic offset between the teacher model and the student model, it is necessary to convert the background probability output by the student model to correctly transfer the knowledge of the old model to the new model. Specifically, Figure 1 As shown, assuming that knowledge transfer is currently being performed between the teacher model t and the student model S, the teacher model t can only identify the foreground organ category t, and other categories are regarded as background categories. If the student model wants to inherit the capabilities of the teacher model, it needs to "imitate" its behavior. Then for the student model, the following processing can be done: for the probability vector output for each pixel, the probability of the channel of category t remains unchanged, and the remaining probabilities are superimposed on the background probability. In this way, the semantics of the teacher model and the student model are consistent in the probability space, and correct knowledge transfer can be performed. The formal expression is:

[0071]

[0072] After the background probability conversion in formula (8), the distribution deviation between teacher t and student S can be calculated.

[0073] Based on the above analysis, knowledge distillation of the output layer in multi-segmentation model aggregation is performed.

[0074] For the scenario of the present invention, for the data set A data in The set of multiple proprietary model outputs to be aggregated is Unified model output It needs to be aligned to each proprietary model output. The loss function for this part is:

[0075]

[0076] Where N represents the number of samples in the dataset.

[0077] Since data from different fields There is usually a large domain deviation, and the training The features extracted from the same image input may also vary greatly due to different training data, that is, there is a gap in the direct migration of model knowledge between different data. In order to solve the above feature offset problem, feature alignment methods are usually used to correct model knowledge bias. Specifically, for the dataset A data in The embeddings extracted by different models are The goal is to minimize the deviation between different embeddings. Figure 1 As shown in Figure 2, we want to concretize the feature alignment task into optimizing an additional lightweight learnable token. The added learnable token is input to the model through The features between the layers interact, changing the output feature distribution, and achieving feature alignment by correcting the output feature distribution deviation. The loss of this part is:

[0078]

[0079] where φ t (·) represents the additional learnable tokens attached to each task, which change the unified model through interaction The feature distribution of the network is adapted to the feature output of each individual model. It is worth noting that the learnable token is added at each feature layer of the network, which means that we have interacted with features of more scales, which is different from previous similar works that only transform features at the highest semantic feature level. For medical image segmentation, the use of multi-scale features plays a great role in improving segmentation results.

[0080] In the inference phase, when executing task t(t∈1,2,…,T), we use the unified model Medium Encoder Based on this, add task-specific tokenφ t (·) Get the model The data corresponding to task t can be input to obtain the corresponding segmentation result.

[0081] In summary, the student model, that is, the training model The overall loss function of (unified model) is:

[0082] L total =L pseudo +λ kd L kd +λ feat L feat #(11)

[0083] where λ kd represents the weight of the loss term of the output layer distillation part, λ feat Represents the weight of the feature distillation / alignment loss term.

[0084] The main contributions of the present invention are as follows:

[0085] 1. A novel framework for multi-model aggregation in medical image segmentation is proposed, which can effectively integrate the knowledge of proprietary models from multiple institutions into a unified model;

[0086] 2. A novel feature alignment method based on learnable tokens is proposed to learn a set of independent parameters for each proprietary model to map the unified model features to the proprietary model feature distribution;

[0087] 3. Through background probability conversion, the output distribution of multiple models is aligned at the output end to strengthen the fusion of model knowledge.

[0088] The present invention conducts aggregation experiments of multiple models on WORD and AMOS datasets. The experimental results show that the present invention can achieve more efficient multi-model knowledge fusion, the use of output distribution alignment can bring obvious effect improvement, and the comprehensive use of feature alignment and output distribution alignment can achieve the best aggregation model segmentation performance improvement.

[0089] In order to verify the effectiveness of the present invention, we conducted model aggregation experiments within and between AMOS and WORD datasets, aiming to simulate the effect evaluation of the fusion of proprietary model knowledge of multiple different types of institutions in the real world. The information of the benchmark dataset used in the present invention is shown in Table 1.

[0090] Table 1 Benchmark dataset information

[0091]

[0092] The method of the present invention was compared with related methods on AMOS and WORD datasets. Including background modeling MIB (Modeling the Background), task vector (Task Arithmetic), weight average (WeightAverage) and semi-supervised learning (Semi-Supervised Learning). It is worth noting that the MIB method was originally used to solve the problem of continuous learning of a single model. Here we extend it to the model aggregation scenario for a fairer comparison. The weight merging method directly interpolates the unified model parameters in the multi-model weight space, and the task vector selects the appropriate interpolation parameters by measuring the effect of the merged model on the validation set. The semi-supervised learning method uses the logits output of multiple trained models in the currently available dataset as pseudo-labels to train the current unified model to inherit its model knowledge.

[0093] Considering that the model has many parameters and the original CT image is large in size, the original CT image is first cropped of its largest foreground area before being sent to the network. On this basis, all CT slices are sampled into 96x96x96 pixel blocks. Then it is normalized, and the pixels with density values ​​between -175 and 250 are mapped to the interval [0,1]. The part exceeding 250 is truncated to 1, and the part below -175 is truncated to 0. In order to improve data diversity, data enhancement strategies such as random horizontal flipping are adopted. The experiment uses the Adamw optimizer to optimize the network, and the weight decay is set to 1e-4 and the momentum is set to 0.9. The initial learning rate is 3e-4, and the cosine learning rate decay strategy is used to obtain a better segmentation effect. The batch size is set to 3, and the model is iterated for a total of 300 rounds. Five rounds of experiments are carried out by changing the random seed, and the experimental results are taken as the average of the five rounds.

[0094] The Dice comparison of each model on the AMOS dataset, {{Spleen (spleen)}, {Right Kidney (right kidney)}, {Left Kidney (left kidney)}} + {Gallbladder (gall bladder)} tasks is shown in Table 2.

[0095] Table 2 Dice comparison of various models on the AMOS dataset, {{Spleen (spleen)}, {Right Kidney (right kidney)}, {Left Kidney (left kidney)}} + {Gallbladder (gallbladder)} tasks

[0096]

[0097] As can be seen from Table 2, the weight-merging method performs poorly in the multi-organ segmentation model aggregation task, and the merged model basically forgets all organ segmentation capabilities. The reason for the analysis is that for the classification task, the model features are mainly concentrated in the encoder part. Using a similar weight space linear weighting strategy can maximize the combination of feature extraction capabilities of different models. For the segmentation task, since it contains a considerable proportion of the decoder part and often fuses multi-scale information through jump connections, its features are not as clear as the classification task. Simply interpolating in the weight space cannot achieve the purpose of integrating the capabilities of different models. The semi-supervised method Semi-Supervised (0.828) that only uses the pseudo-label part has improved the segmentation effect compared with the above method (0.101, 0.025). However, relying solely on pseudo-labels for knowledge transfer depends largely on the accuracy of pseudo-labels, and therefore cannot achieve efficient knowledge transfer. The MIB method uses a pure output layer distillation method to realize the knowledge transfer between the teacher model and the student model. Since the background probability conversion has been completed, it can use the paradigm of knowledge distillation to inherit the teacher model capabilities to the student model as much as possible. However, since the label part only contains the annotations of new organoids, the dataset has serious class imbalance, resulting in a decrease in the segmentation performance of the new gallbladder. The method of the present invention combines pseudo-labeling and knowledge distillation strategies to achieve the best average segmentation effect (0.908).

[0098] The Dice comparison of each model on the WORD dataset, {{Liver(liver)}, {Stomach(stomach)}, {Pancreas(pancreas)}}+{Bladder(bladder)} tasks is shown in Table 3.

[0099] Table 3 Dice comparison of each model on the WORD dataset, {{Liver(liver)}, {Stomach(stomach)}, {Pancreas(pancreas)}}+{Bladder(bladder)} tasks

[0100]

[0101] As can be seen from the table, the method of the present invention shows better performance. Under the premise of ensuring the new class segmentation performance (about 0.925), the knowledge fusion is well completed without touching the annotated dataset used to train the old model. Compared with MIB, which only uses the distillation method, the pseudo-label synthesis strategy proposed by the present invention can improve by 0.9%.

[0102] To simulate the real-world scenario of proprietary model aggregation from multiple institutions, we selected three single-organ segmentation models on the AMOS and WORD datasets, and selected a gallbladder segmentation task on AMOS as the current task. The experimental results are shown in the table:

[0103] Table 4 Aggregation experimental results of WORD and AMOS multi-organ segmentation models

[0104]

[0105] Note that the comparison methods Weight Average and Task Arithmetic performed very poorly in this experimental setting (0.068, 0.000). This paper presents this result here to illustrate that it is very difficult to achieve multi-segmentation model aggregation without training, and it is necessary to improve and optimize it in the field of medical image segmentation. Since this type of method has achieved certain success in the field of natural image classification, and it is also a high-level semantic task, segmentation can be understood as a pixel-level classification task. There is a great correlation between the two. It should be possible to apply this type of method to the field of multi-segmentation model aggregation through certain transformations. This is also an important direction for future exploration in this paper. As can be seen from the table, cross-domain multi-segmentation model aggregation is more difficult than previous experiments. The domain offset between data sets has caused great interference to the transfer of knowledge, and the experimental methods have a large degradation performance. Similarly, the method based on model weight merging performed the worst, and the merged model still basically did not have the ability to segment foreground organs. The MIB method that only uses the distillation method to transfer knowledge has a 2.6% gap compared to the best method in this paper, indicating that knowledge distillation that only relies on the output layer cannot fully transfer knowledge. Since the Semi-Supervised method only uses pseudo labels, the knowledge transfer is also insufficient, and the segmentation effect of the old organoids is significantly reduced compared with the method of the present invention. When feature distillation is not used, the organ segmentation results of the WORD dataset are poor. After adding this module, the performance is significantly improved, which shows the effectiveness of feature distillation for cross-domain model knowledge inheritance. When combining pseudo labels, output layer distillation and feature layer distillation, the method of the present invention can inherit multi-domain and multi-model knowledge to the greatest extent, and the unified model achieves the best results in the teacher model segmentation category.

[0106] The preferred specific embodiments of the present invention are described in detail above. It should be understood that a person skilled in the art can make many modifications and changes based on the concept of the present invention without creative work. Therefore, any technical solution that can be obtained by a person skilled in the art through logical analysis, reasoning or limited experiments based on the concept of the present invention on the basis of the prior art should be within the scope of protection determined by the claims.

Claims

1. A multi-model aggregation method for medical image segmentation based on learnable tokens, characterized in that The method comprises the following steps: S1. Collect the single organ segmentation models obtained from T training sets As the teacher model, the t-th training set corresponds to the t-th teacher model, and the t-th teacher model includes the encoder and decoder Each teacher model performs a segmentation task, the t-th teacher model can only identify the foreground organ category t corresponding to the segmentation task, and the remaining categories are identified as background categories, the parameters of the teacher model are frozen, the CT image data are input into each teacher model, the logical value output by each teacher model is converted into a category probability, and the pseudo label of each single organ segmentation model is obtained based on the category probability; S2. Sort the pseudo labels based on the segmentation capabilities of different models in the single-organ segmentation model set, integrate the sorted pseudo labels to obtain a teacher model pseudo label set, and use the teacher model pseudo label set and the true new class label together as a pseudo label set; S3. Transform the output of the student model so that the output vectors of the student model and the teacher model match, perform distillation of the student model and the teacher model, and input the CT image data into the student model. The student model includes an encoder and decoder The encoder A learnable token is set for each feature layer in the training set, and an overall loss function is calculated based on the output of the student model. The overall loss function includes a distillation part loss, a feature distillation / alignment part loss, and a segmentation part loss. The parameters of the student model are optimized based on the loss function until the iteration converges to obtain a medical image segmentation model. S4. Obtain the actual segmentation task and CT image, and input the CT image into the medical image segmentation model to obtain the segmentation result corresponding to the segmentation task.

2. The multi-model aggregation method for medical image segmentation based on learnable tokens according to claim 1, characterized in that: The class probability output by the teacher model is: Among them, t represents the t-th teacher model, It represents the output logic value of the t-th teacher model for identifying organ category i. i is 1, which means that i is identified as the foreground organ category corresponding to the t-th teacher. i is 0, which means that i is identified as the background category. represents the output logic value of the t-th teacher model for organ category k. k is 1, which means that k is identified as a foreground organ category, and k is 0, which means that k is identified as a background category. T represents the total number of teacher models. Represents the category probability of organ category k output by the tth teacher model.

3. The multi-model aggregation method for medical image segmentation based on learnable tokens according to claim 2, characterized in that: The pseudo label of each single organ segmentation model is: Among them, t represents the t-th teacher model, M t represents the pseudo label value of the t-th teacher model, p t T (0) represents the probability of the t-th teacher output identifying the organ category k as the background category, p t T (1) represents the category probability of the t-th teacher output identifying the organ category k as the foreground organ category.

4. The method for multi-model aggregation of medical image segmentation based on learnable tokens according to claim 3, characterized in that: The output probability of the student model matching each teacher model is: in, represents the output logical value of the student model S for identifying the foreground organ category t identified by the t-th teacher model, p S (t) represents the probability of the foreground organ category t being identified as the foreground organ category by the output of the student model S, that is, the output of the student model when it is not transformed, 1-p S (t) represents the category probability output by the student model S for identifying the foreground organ category t as the background category.

5. The method for multi-model aggregation of medical image segmentation based on learnable tokens according to claim 4, characterized in that: The segmentation loss is: L pseudo =L ce (p S ,g pseudo ) Among them, L pseudo represents the segmentation loss, L ce represents the cross entropy loss function, g pseudo is the pseudo label set, p S Represents the output of the student model when it is not transformed.

6. The multi-model aggregation method for medical image segmentation based on learnable tokens according to claim 1, characterized in that: The pseudo-label set is: g pseudo =[M T ,M S ] Among them, M T is the pseudo label set of T single organ segmentation models, M S is the real new class label.

7. The multi-model aggregation method for medical image segmentation based on learnable tokens according to claim 1, characterized in that: The distillation part loss is: Where N represents the number of training set samples, KL represents KL divergence, represents the teacher model corresponding to the t-th training set, represents the student model, x i represents the i-th training set sample.

8. The multi-model aggregation method for medical image segmentation based on learnable tokens according to claim 1, characterized in that: The feature distillation / alignment part loss is: Where N represents the number of training set samples. represents the encoder of the teacher model corresponding to the t-th training set, represents the encoder of the student model, φ t Represents the feature conversion function, that is, by setting additional learnable tokens, the features encoded by the student model are converted to the feature space corresponding to the tth teacher model encoding.

9. The method for multi-model aggregation of medical image segmentation based on learnable tokens according to claim 1, characterized in that: The overall loss function is specifically: L total =L pseudo +λ kd L kd +λ feat L feat Among them, L total is the overall loss function, L pseudo To split the partial loss, L kd Represents the distillation loss, L feat represents the feature distillation / alignment loss, λ kd represents the weight of the loss term of the output layer distillation part, λ feat Represents the weight of the feature distillation / alignment loss term.

10. The method for multi-model aggregation of medical image segmentation based on learnable tokens according to claim 1, characterized in that: The specific steps of sorting the pseudo labels based on the segmentation capabilities of different models in the single organ segmentation model set and integrating the sorted pseudo labels to obtain the teacher model pseudo label set are as follows: For each teacher pseudo label M1…M T Each pixel on the segmentation map is sorted according to the confidence level corresponding to each segmentation map. Finally, the T pseudo labels are merged to obtain the teacher model pseudo label set M. T .

Images