Semi-supervised medical image segmentation method based on model self-distillation and prototype learning

Abstract

The application discloses a semi-supervised medical image segmentation method based on model self-distillation and prototype learning, and relates to the technical field of image signal processing. The semi-supervised medical image segmentation method based on model self-distillation and prototype learning comprises the following steps: S1, establishing a semi-supervised medical image segmentation dataset; S2, constructing a network architecture; S3, designing a semi-supervised medical image segmentation scheme, and building a semi-supervised medical image segmentation model according to the designed scheme; S4, training the semi-supervised medical image segmentation model by using a deep learning Pytorch framework; and S5, inputting a medical image to be segmented into the model to obtain a medical image segmentation result. The semi-supervised medical image segmentation performance is improved to a new height by using the proposed double-flow memory bank architecture, the self-distillation method based on an image block affinity matrix and the prototype synthesis method based on context matching.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of image signal processing, and in particular to a semi-supervised medical image segmentation method based on model self-distillation and prototype learning. BACKGROUND

[0002] Medical image segmentation aims to segment organs, lesions, blood vessels or other clinically concerned anatomical structures from medical images of different modalities, which can be applied to cancer diagnosis, lesion monitoring and many other practical problems. In recent years, data-driven deep learning methods have made great progress in medical image segmentation tasks. Fully supervised deep learning image segmentation methods have achieved better segmentation performance on multiple public datasets. However, the performance of fully supervised image segmentation methods depends on a large number of dense pixel labels, and the segmentation labels of medical images require time-consuming and laborious labor of professionals, so it is difficult to obtain a large number of them. In order to solve this problem, researchers propose semi-supervised medical image segmentation, which aims to train a model using a small amount of labeled data and a large amount of unlabeled data to obtain segmentation results comparable to fully supervised methods.

[0003] Existing semi-supervised medical image segmentation methods can be roughly divided into consistency regularization methods, entropy minimization methods, self-supervised methods and adversarial learning methods. These methods generally pay attention to the overall quality of the segmentation results or extract global features of the semantic relationship, ignoring the characteristics of the pixel level or local area, so that inaccurate segmentation results are generated in the edge or fuzzy area of the segmentation target.

[0004] On the other hand, some researchers have applied prototype learning methods to the field of medical image segmentation. Such methods extract the class center (prototype) of a class or an instance from the feature space, and then use it as the target for subsequent model training. The model trained by this method has a more compact intra-class feature distribution and a more discriminative inter-class feature distribution, greatly improving the performance of the model under limited data labels. However, in the field of medical image segmentation, existing methods usually only calculate one overall prototype in a class, ignoring the different characteristics of the same class of pixels in the context of semantics, resulting in inaccurate segmentation results. Therefore, it is of research value to explore intra-class prototype generation methods that can perceive context information. In addition, self-distillation methods can learn to distinguish different local semantics without using labels, which is an effective way to improve segmentation performance. SUMMARY

[0005] The purpose of the present application is to generate high-quality medical image segmentation results using a small amount of labels, and on this basis, a semi-supervised medical image segmentation method based on model self-distillation and prototype learning is proposed.

[0006] In order to achieve the above purpose, the present application adopts the following technical solutions:

[0007] The semi-supervised medical image segmentation method based on model self-distillation and prototype learning comprises the following steps:

[0008] S1, establishing a semi-supervised medical image segmentation dataset: given a medical image q, the segmentation label y is consistent with the resolution of the original image, and different regions are marked by different colors or different gray values to indicate organ regions, lesion positions, blood vessel networks or other clinically concerned anatomical structures, thereby constructing a semi-supervised medical image segmentation dataset containing labeled images and unlabeled images;

[0009] S2, constructing a network framework: designing a student model comprising a backbone network F, a self-distillation encoder E sd , and two independent linear branches G and H; at the same time, a teacher model identical in architecture to the student model is constructed, and the above-mentioned student model and teacher model together constitute a teacher-student model framework; the parameters φ s of the student model are optimized through the gradient backpropagation of the loss function during the training process; the parameters φ t of the teacher model do not calculate gradients during the training process, and the parameters of the teacher model are updated through the formula φ t ' = 0.99 * φ t + 0.01 * φ s after each training iteration ends;

[0010] S3, designing a scheme and building a model: based on model self-distillation and prototype learning, combining the dataset constructed in S1 and the network framework constructed in S2, designing a semi-supervised medical image segmentation scheme, and building a semi-supervised medical image segmentation model according to the designed scheme, the scheme specifically comprises the following contents:

[0011] ① Medical image data augmentation: randomly cropping the medical image q into a fixed size image, then performing two independent random spatial transformations T s and T t , obtaining the augmented images q s and q t , and inputting them into the student model and the teacher model constructed in S2, respectively;

[0012] ② Obtaining image feature maps: inputting the augmented images into the backbone network F to obtain image feature maps consistent with the resolution of the image, denoted as f; inputting the image feature maps f into the self-distillation encoder E sd to obtain high-dimensional low-resolution self-distillation feature maps, denoted as f sd , and each spatial position in the feature map represents the features of an image block at the corresponding position in the original image; inputting the image feature maps f into two linear branches G and H to obtain contrast feature maps and prediction probability maps consistent with the resolution of the image, denoted as fco and f p ;

[0013] ③Obtain pseudo-labels and calculate image block prototypes: For unlabeled images, use the prediction probability map f p output by the teacher model to calculate the pseudo-labels of the images, denoted as y p ; Use the student model described in ② to compare the feature map f co and the image label y or the pseudo-label y p , and calculate the image block prototypes corresponding to each image block for each class in units of image blocks;

[0014] ④Construct an image block feature memory bank: After each training iteration, for labeled images in the input images, flatten the self-distillation feature map f sd output by the teacher model in spatial order and store it in the image block feature memory bank, denoted as M f ; The image block feature memory bank M f stores the same number of image blocks as all labeled image blocks in the dataset; After the memory bank is full, update the memory bank in the order of a first-in, first-out queue;

[0015] ⑤Construct an image block prototype memory bank: After each training iteration, for labeled images in the input images, use the teacher model described in ② to compare the feature map f co and the corresponding segmentation label y of the image, calculate the corresponding image block prototypes for each class in units of image blocks, and save them to the image block prototype memory bank, denoted as M η ; The image block prototype memory bank M η has a one-to-one correspondence with the image block features in the image block feature memory bank M f described in ④, and the number of stored image blocks and the memory bank update method are consistent;

[0016] ⑥Construct an image block affinity matrix: Combine the self-distillation feature map f sd described in ② and the image block feature memory bank M f described in ④ to calculate the image block affinity matrix; The student model and the teacher model calculate their respective image block affinity matrices, denoted as A s and A t , respectively;

[0017] ⑦Prototype synthesis with context awareness: Use the student model image block affinity matrix A s described in ⑥ to perform clustering operations, placing image blocks with similar features in the image and the memory bank into the same cluster, and combining the affinity matrix A s , the image block prototypes described in ③, and the image block prototype memory bank Mη compute the context-aware prototype of each class for each cluster;

[0018] ⑧Loss function design: for labeled images, compute the supervised loss using the student model prediction probability map f p and image label y; for unlabeled images, compute the consistency loss using the student model prediction probability map f p and pseudo label y p ; compute the self-distillation loss using the image block affinity matrix A s and A t ; compute the contrastive loss based on context-aware prototype using the student model contrastive feature map f co and the context-aware prototype described in ⑦;

[0019] S4, model training: train the model using the deep learning Pytorch framework; first, use the supervised loss described in ⑧ to traverse the small amount of labeled data in the semi-supervised medical image segmentation dataset constructed in S1 until the student network supervised loss function is preliminarily converged; then use the total loss function described in ⑧ to traverse all the data in the semi-supervised medical image segmentation dataset until the final stable model is obtained;

[0020] S5, output results: input the data in the semi-supervised medical image segmentation dataset obtained in S1 into the stable model to obtain the medical image segmentation results.

[0021] Preferably, the scheme ② further includes the following contents:

[0022] 2.1) The backbone network F is a fully symmetric encoder-decoder network, wherein the encoder network contains 4 encoding sub-modules, each of which can downsample the feature map resolution to half of the original; the decoder contains 4 decoding sub-modules, each of which can upsample the feature map resolution to twice the original, and then input the output feature map to the next module after splicing with the feature map of the same resolution in the encoding sub-module;

[0023] 2.2) The network structure of the self-distillation encoder E sd is consistent with that of the encoder network of the backbone network F, but the gradient of the feature map in the self-distillation encoder E sd does not return to the backbone network F; the linear branch G is a 1x1 convolutional layer with an output channel number of 16; the linear branch H is a 1x1 convolutional layer plus a softmax layer, and the output channel number is equal to the number of classes of the medical image label.

[0024] Preferably, the scheme ③ specifically includes the following contents:

[0025] for the i-th image block pi , use the following formula to calculate the image block prototype η of its category c i,c :

[0026]

[0027] In the formula, |·| represents the number of elements in the set; Represents the indicator function; the superscript (i, j) represents the j-th pixel in the i-th image block; y represents the label or pseudo-label of the image.

[0028] Preferably, the solution ⑥ specifically includes the following contents:

[0029] For the self-distilled feature map f output by the teacher model sd Before calculating the affinity matrix of the image block, its spatial position must be transformed by T s (T t -1 (f sd )), thereby ensuring the spatial consistency of the self-distillation feature maps output by the teacher model and the student model; combined with the self-distillation feature map f sd and image patch feature memory M f , calculate the image block affinity matrix A according to the following formula:

[0030]

[0031] Where, Represents matrix multiplication.

[0032] Preferably, the solution ⑦ further includes the following contents:

[0033] 7.1) Use the k-means method to calculate the image block affinity matrix A s The behavioral units in are clustered so that the image blocks in the memory bank are assigned to K clusters;

[0034] 7.2) For each cluster, sum and normalize all rows in the cluster to get the similarity list S = [s1, s2, ..., s i ], the number s in S i The image patch p represents the input image i Similarity between clusters; if the image block p of the input image i The similarity s i Greater than the average of all numbers in S, or image block p i Similarity s with the current cluster i If the similarity between the image block and all other clusters is greater than that between the image block and all other clusters, the image block is assigned to the current cluster.

[0035] 7.3) For each cluster, according to the input image patches contained in it extracting input image patch prototypes from the memory bank image patch prototypes contained therein from the memory bank image patch prototypes contained therein η extracting corresponding memory bank image patch prototypes

[0036] 7.4) For a class c in a cluster k, the context-aware prototype is computed using the following equation:

[0037]

[0038] where λ q =∑ i s i ; λ m denotes the total number of memory bank image patch prototypes belonging to class c contained in cluster k; γ denotes a hyper-parameter balancing the weights between input image patch prototypes and memory bank image patch prototypes.

[0039] Preferably, the scheme (vii) further comprises the following:

[0040] 8.1) For labeled images, the student model predicts a probability map f p and the image label y, the supervised loss is computed using the following equation:

[0041]

[0042] where L dice (·) denotes the dice segmentation loss; L ce (·) denotes the cross-entropy segmentation loss;

[0043] 8.2) For unlabeled images, the student model predicts a probability map f p and the image pseudo-label y p , the consistency loss is computed using the following equation:

[0044]

[0045] 8.3) Using the image patch affinity matrix A s and A t , the self-distillation loss is computed using the following equation:

[0046] L sd = KL(Softmax(A t / τ t )||Softmax(A s / τ s ))

[0047] where KL(·) denotes the KL divergence, τ t and τs temperature hyper-parameters representing the teacher model and the student model;

[0048] 8.4) For class c in cluster k, the student model is used to calculate the contrastive feature map f co and the context-aware prototype, the context-aware prototype-based contrastive loss is calculated using the following formula:

[0049]

[0050] where τ represents the temperature hyper-parameter, representing the selected pixels in the input image block belonging to cluster k that are classified as c by the image label or pseudo label, and the pixels with higher uncertainty are selected to participate in the calculation; if some input image blocks exist in all clusters, these image blocks are divided into a new cluster containing all memory bank image blocks and do not participate in the calculation of the contrastive loss in the original cluster;

[0051] The total context-aware prototype-based contrastive loss is calculated using the following formula:

[0052]

[0053] 8.5) The total loss of the method is calculated using the following formula:

[0054] L = L sup + L con + aL sd + bL caco

[0055] where a and b represent the self-distillation loss weight and the context-aware prototype-based contrastive loss weight, respectively.

[0056] Compared with the prior art, the present application provides a semi-supervised medical image segmentation method based on model self-distillation and prototype learning, which has the following beneficial effects:

[0057] (1) The present application proposes a semi-supervised medical image segmentation method based on model self-distillation and prototype learning; model self-distillation learns to distinguish different local semantic information under unsupervised conditions through teacher-student models; prototype learning synthesizes context-aware prototypes by matching image blocks with similar contexts and obtains rich and diverse local semantic information in images using prototype-based contrastive loss.

[0058] (2) The present application proposes a dual-flow memory bank architecture that can match image blocks with similar local information within the full data set range, enhance the ability of class prototypes to capture global context information of medical images, and fully utilize the reliable information of a small amount of labeled data to enhance the correct recognition ability of the network for edge and fuzzy regions.

[0059] (3) The experiment based on the present application shows that the proposed method is superior to the current mainstream semi-supervised medical image segmentation method; through the research and exploration of the present application, more researches using local semantic information in visual tasks can be inspired. BRIEF DESCRIPTION OF DRAWINGS

[0060] Figure 1 The overall framework diagram of the semi-supervised medical image segmentation method based on model self-distillation and prototype learning proposed by the present application (omitting the linear branch G, the supervised loss and the consistency loss) ;

[0061] Figure 2 The context-aware prototype diagram proposed by the present application, the upper part is a prototype generation diagram of the past method, and the lower part is a context-aware prototype generation diagram adopted by the present method. DETAILED DESCRIPTION

[0062] The technical solutions in the embodiments of the present application will be clearly and completely described below with reference to the drawings in the embodiments of the present application. Obviously, the described embodiments are only a part of the embodiments of the present application, rather than all the embodiments.

[0063] Embodiment 1:

[0064] Please refer to Figure 1 , the present application proposes a semi-supervised medical image segmentation method based on model self-distillation and prototype learning, which comprises the following steps:

[0065] S1, establishing a semi-supervised medical image segmentation dataset: given a medical image q, its segmentation label y is consistent with the resolution of the original image, but marks different colors or different gray values of different regions, which are used to indicate organ regions, lesion positions, blood vessel networks or other clinically concerned anatomical structures; the semi-supervised medical image segmentation dataset contains a small amount of labeled images and a large amount of unlabeled images;

[0066] S2, constructing a network framework: designing a student model containing a backbone network F, a self-distillation encoder E sd and two independent linear branches G and H; at the same time, a teacher model which is completely consistent with the architecture of the student model is constructed. The two networks together constitute a teacher-student model framework. The parameters φ s of the student model are optimized through the gradient back propagation of the loss function in the training process; the parameters φ t of the teacher model are not calculated in the training process, and the parameters φ t of the teacher model are updated through the formula φ t ’=0.99*φ s +0.01*φ

[0067] S3, design scheme, build model: based on model self-distillation and prototype learning, combined with the data set described in S1 and the network framework described in S2, design a semi-supervised medical image segmentation scheme, and build a semi-supervised medical image segmentation model according to the designed scheme;

[0068] S4, training model: training the model using the deep learning Pytorch framework; first, using supervised loss, traversing the small amount of labeled data in the semi-supervised medical image segmentation data set constructed in S1, until the student network supervision loss function is preliminarily converged; then using all loss functions, traversing all data in the semi-supervised medical image segmentation data set, until the final stable model is obtained;

[0069] S5, output result: input the data in the semi-supervised medical image segmentation data set obtained in S1 into the stable model to obtain the medical image segmentation result.

[0070] Further to the above method, further includes the following contents:

[0071] 3.1 teacher-student model

[0072] In the teacher-student model, the teacher model does not accept gradient back propagation in the training process, but accepts the parameter sharing of the student model and integrates it in time sequence, so it is considered as a more stable and reliable model, which can assist the training of the student model and make the training process more stable.

[0073] In order to make full use of the teacher model and increase the robustness of the model training, the input medical image q is randomly cropped into a fixed size image, and then two independent random spatial transformations T s and T t are performed to obtain enhanced images q s and q t , which are input into the student model and the teacher model respectively.

[0074] The enhanced images are input into the backbone network F to obtain an image feature map consistent with the image resolution, denoted as f; the image feature map f is input into the self-distillation encoder E sd to obtain a high-dimensional low-resolution self-distillation feature map, denoted as f sd , and each spatial position in the feature map represents the features of an image block at the corresponding position in the original image; the image feature map f is input into two linear branches G and H respectively to obtain a contrast feature map and a prediction probability map consistent with the image resolution, denoted as f co and f p ; for the same medical image, the student model and the teacher model output different types of feature maps described above. Since the self-distillation encoder E sdFour times of downsampling are performed, and each time the feature map is downsampled to 1 / 2 of the original resolution, so f sd The spatial resolution of f sd is 1 / 16 of the original image, that is, f sd The features at each spatial position in f f represent a 16*16 image block (for a 2D dataset) on the original image. The image block of this size is also the basic unit for subsequent image block affinity learning and image block prototype calculation.

[0075] The teacher-student model is introduced for the following three reasons: first, the teacher model can generate high-quality pseudo-labels y p for unlabeled images, thereby guiding the calculation of subsequent consistency loss and contrast loss based on upper and lower perceptual prototypes; second, the teacher model can generate robust feature maps, helping to build subsequent image block feature memory and image block prototype memory; third, the teacher-student model, as a common distillation learning strategy, can well connect the subsequent self-distillation learning, enabling the model to learn image block affinity unsupervisedly.

[0076] 3.2 Building a double-flow memory

[0077] In order to effectively utilize the reliable information of a small amount of labeled data and simultaneously make the feature distribution of unlabeled data approach the labeled data, the present application builds a double-flow memory framework to store the high-dimensional global information and low-dimensional local information of the labeled data.

[0078] After each training iteration, for the labeled images in the input images, the self-distillation feature maps f sd output by the teacher model are flattened in spatial order and stored in the image block feature memory, denoted as M f ; at the same time, the contrast feature maps f co output by the teacher model and the segmentation labels y corresponding to the images are used to calculate the image block prototypes corresponding to each class in units of image blocks, and the image block prototypes are saved in the image block prototype memory, denoted as M η The number of image blocks stored in the two memories is equal to the number of all labeled image blocks in the dataset, and the image blocks stored in the two memories have a one-to-one correspondence; when the memories are full, the memories are updated in the order of first-in-first-out. Such a design enables the model to utilize not only the current batch of data but also the labeled data features in the entire dataset during training.

[0079] 3.3 Self-distillation learning based on image block affinity

[0080] In order to enable the model to learn to distinguish the local semantic information of images unsupervisedly, the present application proposes self-distillation learning based on image block affinity. For the self-distillation feature maps f sd, first transform the spatial position of the image block p s (T t -1 (f sd )) to ensure the spatial consistency of the self-distillation feature maps output by the teacher model and the student model; then combine the self-distillation feature map f sd and the image block feature memory M f , and calculate the image block affinity matrix A according to the following formula:

[0081]

[0082] In the formula, denotes matrix multiplication. The student model and the teacher model calculate their respective image block affinity matrices, denoted as A s and A t respectively; then use the two affinity matrices to calculate the self-distillation loss using the following formula:

[0083] L sd =KL(Softmax(A t / τ t )||Softmax(A s / τ s ))

[0084] In the formula, KL(·) denotes the KL divergence, and τ t and τ s are the temperature hyperparameters of the teacher model and the student model, respectively, and are set to 0.1 and 0.01.

[0085] 3.4 Contrastive learning based on context-aware prototypes

[0086] The segmentation result of a medical image depends on specific context information, especially in the edge region and the fuzzy region of the target in the image. The present application designs a set of context-aware prototype generation methods, so that the prototypes as the learning target of the model can have specific context information, thereby helping the model to learn the segmentation of difficult sample regions.

[0087] In order to obtain the prototype of the smallest unit region, it is necessary to calculate the image block prototype. For the i-th image block p i in the image, the image block prototype η i,c of class c is calculated using the following formula:

[0088]

[0089] In the formula, |·| denotes the number of elements in the set, denotes the indicator function, and the superscript (i,j) denotes the j-th pixel in the i-th image block, and y denotes the label or pseudo-label of the image.

[0090] To match similar regions into the same cluster, and thus generate context-aware prototypes and compute prototype-based contrastive loss, one needs to utilize the image patch affinity matrix A s The image patches are clustered. The k-means method is used to cluster the behavior units in the image patch affinity matrix A s , so that the image patches in the memory bank are assigned into K clusters; for each cluster, sum all the rows in the cluster and normalize, to get a similarity list S = [s1, s2,..., s i ], the number s i in S represents the similarity between the image patch p i of the input image and the cluster; if the similarity s i of the image patch p i of the input image is greater than the average of all numbers in S, or the similarity s i of the image patch p i to the current cluster is greater than its similarity to all other clusters, then the image patch is assigned to the current cluster.

[0091] At this time, the clusters contain both memory bank image patches and input image patches, and context-aware prototypes can be calculated accordingly. For each cluster, according to the input image patches p i q contained therein, extract the input image patch prototype η i q ; according to the memory bank image patches p j m contained therein, extract the corresponding memory bank image patch prototype η η j m from the image patch prototype memory bank M. For class c in cluster k, the context-aware prototype is calculated using the following formula:

[0092]

[0093] In the formula, λ q =∑ i s i , λ m is equal to the total number of memory bank image patch prototypes belonging to class c contained in cluster k, and γ is a hyperparameter that balances the weight between input image patch prototypes and memory bank image patch prototypes.

[0094] According to the clustering results and the generated prototypes, context-aware prototype-based contrastive loss can be calculated. For class c in cluster k, using the student model contrastive feature map f co and the context-aware prototype, the context-aware prototype-based contrastive loss is calculated using the following formula:

[0095]

[0096] where τ represents the temperature hyper-parameter, selected pixels in the input image patches belonging to cluster k that are classified as c by the image label or pseudo label, and the pixels with higher uncertainty are selected to participate in the calculation. If some input image patches exist in all clusters, these image patches will be calculated multiple times, which not only affects the training efficiency of the model, but also has a negative effect on the segmentation result of the model. These image patches are usually the main part of the target class area and are easy to be correctly segmented by the model. The repeated calculation of these image patches is equivalent to increasing the weight of these regions in the loss, so that the weight of a small amount of edge regions or fuzzy regions is squeezed, resulting in inaccurate segmentation results. Therefore, the above image patches are divided into a new cluster containing all the memory bank image patches, and do not participate in the calculation of the contrast loss in the original cluster. Then, the total context-aware prototype-based contrast loss is calculated using the following formula:

[0097]

[0098] 3.5 Loss function

[0099] For labeled images, the student model prediction probability map f p and the image label y, the supervised loss is calculated using the following formula:

[0100]

[0101] where L dice (·) is the commonly used dice segmentation loss, L ce (·) is the commonly used cross-entropy segmentation loss;

[0102] For unlabeled images, the student model prediction probability map f p and the image pseudo label y p generated by the teacher model, the consistency loss is calculated using the following formula:

[0103]

[0104] The total loss of the method is calculated using the following formula:

[0105] L = L sup + L con + aL sd + bL caco

[0106] where a and b represent the self-distillation loss weight and the context-aware prototype-based contrast loss weight, respectively.

[0107] Example 2:

[0108] Based on Example 1 but with the difference that:

[0109] The present application selects 8 advanced comparison methods trained on the 2D ACDC dataset, including: UAMT, SASSNet, DTC, URPC, SS-Net, MC-Net, MC-Net+ and BCP-Net. On this dataset, 2D U-Net is used as the backbone network. U-Net is trained using 10% data labels and supervised loss as the lower limit of model performance, denoted as U-Net-L. U-Net is trained using 100% data labels and supervised loss as the upper limit of model performance, denoted as U-Net-F. For specific results, please refer to Table 1. The present application also selects 9 advanced comparison methods trained on the 3D LA dataset, including: UAMT, SASSNet, DTC, URPC, SS-Net, MC-Net, MC-Net+, CAML and BCP-Net. On this dataset, 3D V-Net is used as the backbone network. V-Net is trained using 10%, 20% data labels and supervised loss as the lower limit of model performance, denoted as V-Net-L. V-Net is trained using 100% data labels and supervised loss as the upper limit of model performance, denoted as V-Net-F. For specific results, please refer to Table 2.

[0110] Table 1: ACDC dataset quantitative result comparison table. The best result is shown in bold.

[0111]

[0112] As shown in Table 1, the table shows the quantitative comparison results on the Dice, JC, 95HD and ASD indicators, where Dice and JC are similarity indicators, the larger the value, the better, 95HD and ASD are evaluation indicators for measuring the distance of the segmentation edge, the smaller the value, the better. The optimal result in the table is highlighted in bold. When using 10% labeled data, there is a big gap between the present method and all comparison methods, which exceeds the sub-optimal BCP-Net method by 1.42% / 0.41% / 0.64 / 0.22 in Dice, JC, 95HD and ASD, respectively. These results prove the effectiveness of the self-distillation and context-aware prototype proposed in the present method on 2D medical image segmentation. Especially in the evaluation indicators for measuring the distance of the segmentation edge, the present method achieves better results than the fully supervised U-Net, proving the excellent recognition ability of the present method in the edge and fuzzy areas.

[0113] Table 2: LA dataset quantitative result comparison table. The best result is shown in bold.

[0114]

[0115] As shown in Table 2, the table shows the quantitative comparison results under the settings of 10% and 20% data labels. As can be seen from the table, the method of the present chapter also has superiority in the 3D medical image segmentation task. When using 10% labeled data, the method exceeds the suboptimal BCP-Net method by 0.6% / 0.96% / 0.41 / 0.08 in Dice, JC, 95HD and ASD, respectively. When using 20% labeled data, the method exceeds the suboptimal MC-Net+ method by 0.25% / 0.42% / 0.02 in Dice, JC and ASD, respectively, and exceeds the suboptimal BCP-Net method by 0.24 in 95HD. These results prove the effectiveness of the self-distillation and context-aware prototype proposed by the method in 3D medical image segmentation. When using 20% data labels, the segmentation performance of the method basically reaches the level of the fully supervised V-Net using 100% data labels, proving that the method can effectively utilize the reliable features of a small amount of labeled data and a large amount of unlabeled data, so that the segmentation effect can be compared with full supervision.

[0116] The above merely describes preferred specific embodiments of the present application, but the protection scope of the present application is not limited thereto, and any person skilled in the art can make equivalent replacements or changes to the technical solutions and inventive concepts of the present application within the technical scope disclosed by the present application, which should be covered within the protection scope of the present application.

Claims

1. A semi-supervised medical image segmentation method based on model self-distillation and prototype learning, characterized in that, The specific steps include: S1. Establish a semi-supervised medical image segmentation dataset: Given a medical image q, its segmentation label y is consistent with the original image resolution. Different regions are marked with different colors or grayscale values ​​to indicate organ regions, lesion locations, vascular networks, or other anatomical structures of clinical interest. Then, a semi-supervised medical image segmentation dataset containing labeled and unlabeled images is constructed. S2, constructing a network framework: designing a student model containing a backbone network F, a self-distillation encoder E sd and two independent linear branches G and H; at the same time, constructing a teacher model which is completely consistent with the architecture of the student model, and utilizing the above-mentioned student model and teacher model to jointly constitute a teacher-student model framework; the parameters φ s of the student model are optimized through the gradient backpropagation of the loss function during the training process; the parameters φ t of the teacher model are not calculated in the training process, and the parameters of the teacher model are updated through the formula φ t '=0.99*φ t +0.01*φ s after each training iteration is completed. S3. Design a solution and build a model: Based on model self-distillation and prototype learning, combined with the dataset constructed in S1 and the network framework constructed in S2, design a semi-supervised medical image segmentation solution and build a semi-supervised medical image segmentation model based on the designed solution. The solution specifically includes the following: ① Medical image data augmentation: randomly crop the medical image q into a fixed size image, and then perform two independent random spatial transformations T s and T t , to obtain the augmented image q s and q t , and input them into the student model and the teacher model constructed in S2, respectively; ②Obtain image feature map: input the enhanced image into the backbone network F to obtain an image feature map consistent with the image resolution, denoted as f; input the image feature map f into the self-distillation encoder E sd to obtain a high-dimensional low-resolution self-distillation feature map, denoted as f sd ; the high-dimensional vector of each spatial position in the feature map represents the features of an image block at the corresponding position in the original image; input the image feature map f into two linear branches G and H respectively to obtain contrast feature maps and prediction probability maps consistent with the image resolution, denoted as f co and f p ; ③ Obtain pseudo-labels and calculate image block prototypes: For unlabeled images, use the prediction probability map f p Compute the pseudo-labels of the images, denoted as y p ; Use the student model described in ② to compare the feature map f co , and the image label y or the pseudo-label y p Calculate the image block prototype corresponding to each image block for each class in units of image blocks; ④ Constructing the image block feature memory bank: after each training iteration, for the labeled images in the input images, the self-distillation feature map f sd output by the teacher model is flattened in spatial order and stored in the image block feature memory bank, denoted as M f ; the image block feature memory bank M f stores the number of image blocks equal to the number of all labeled image blocks in the data set; after the memory bank is full, the memory bank is updated in the order of the first-in, first-out queue; ⑤ Constructing the image block prototype memory bank: after each training iteration, for the labeled images in the input images, the teacher model described in ② is used to compare the feature map f co and the image corresponding segmentation label y, and the corresponding image block prototype for each class is calculated in units of image blocks, and saved to the image block prototype memory bank, denoted as M η ; the image block prototype memory bank M η has a one-to-one correspondence with the image block features in the image block feature memory bank M f described in ④, and the number of stored image blocks and the memory bank updating mode are consistent. ⑥ Constructing the image block affinity matrix: combining the self-distilled feature map f sd with the image block feature memory M f described in ④, the image block affinity matrix is calculated; the student model and the teacher model calculate their respective image block affinity matrices, respectively denoted as A s and A t ; Prototype synthesis with context awareness: using the student model image patch affinity matrix A from 6 s to perform a clustering operation to place image patches with similar features in the image and memory bank into the same cluster and combine the affinity matrix A s with the image patch prototypes from 3 and the image patch prototype memory bank M from 5 η to compute a context aware prototype for each class of pixels for each cluster; ⑧ Loss function design: For labeled images, use the student model to predict the probability map f p and image label y, calculate the supervision loss; for unlabeled images, use the student model to predict the probability map f p and pseudo-label y p , calculate the consistency loss; use the image block affinity matrix A described in ⑥ s and A t , calculate the self-distillation loss; use the student model to compare the feature map f co and the context-aware prototype described in ⑦, calculating the contrast loss based on the context-aware prototype; S4. Training the model: Use the deep learning Pytorch framework to train the model. First, use the supervised loss function described in ⑧ to traverse a small amount of labeled data in the semi-supervised medical image segmentation dataset constructed in S1 until the student network supervised loss function initially converges. Then, use all the loss functions described in ⑧ to traverse all the data in the semi-supervised medical image segmentation dataset until a final stable model is obtained. S5. Output results: Input the data in the semi-supervised medical image segmentation dataset obtained in S1 into the stable model to obtain the medical image segmentation results.

2. The model-based self-distillation and prototype learning based semi-supervised medical image segmentation method of claim 1, wherein, The solution ② further includes the following contents: 2.1) The backbone network F is a fully symmetric encoder-decoder network. The encoder network contains four encoding submodules, each of which can downsample the feature map resolution to half of its original resolution. The decoder network contains four decoding submodules, each of which can upsample the feature map resolution to twice its original resolution. The output feature map is concatenated with the feature map of the same resolution from the encoding submodule and then input into the next module. 2.2) Self-distillation encoder E sd The network structure of the self-distillation encoder E is consistent with the encoder network of the backbone network F, but the gradient of the feature map in the self-distillation encoder E sd does not feed back to the backbone network F; the linear branch G is a 1x1 convolutional layer with an output channel number of 16; the linear branch H is a 1x1 convolutional layer plus a Softmax layer, and the output channel number is equal to the class number of the medical image label.

3. The model-based self-distillation and prototype learning based semi-supervised medical image segmentation method of claim 1, wherein, The said plan ③ specifically includes the following contents: For the i-th image patch p in the image i , the image patch prototype η for its class c is computed using the following formula i,c : In the formula, |·| represents the number of elements in the set; denotes an indicator function; the superscript (i,j) denotes the jth pixel in the ith image block; y denotes the label or pseudo-label of the image.

4. The model-based self-distillation and prototype learning based semi-supervised medical image segmentation method of claim 1, wherein, The said scheme ⑥ specifically includes the following contents: For the self-distillation feature map f output by the teacher model sd Before calculating the image block affinity matrix, the spatial position thereof is transformed Thus, the spatial consistency of the self-distillation feature maps output by the teacher model and the student model is ensured; in combination with the self-distillation feature map f sd and the image block feature memory library M f The image block affinity matrix A is calculated according to the following formula: In the formula, denotes matrix multiplication.

5. The model-based self-distillation and prototype learning based semi-supervised medical image segmentation method of claim 1, wherein, The scheme ⑦ further includes the following contents: 7.1) Use the k-means method to calculate the image block affinity matrix A s The behavioral units in are clustered so that the image blocks in the memory bank are assigned to K clusters; 7.2) For each cluster, sum all rows in the cluster and normalize to get a list of similarities S = [s1, s2,..., s i ], where s i in S represents the similarity between the image patch p i and the cluster. If the similarity s i of the image block p i of the input image is greater than the average of all numbers in S, or the similarity s i of the image block p i to the current cluster is greater than its similarity to all other clusters, then the image block is assigned to the current cluster; 7.3) for each cluster, extract from the input image blocks contained therein extract input image block prototypes from the memory bank image blocks contained therein from the image block prototype memory bank M η extract corresponding memory bank image block prototypes 7.4) For category c in cluster k, calculate the context-aware prototype using the following formula: where λ q =∑ i s i ; λ m denotes the total number of memory image patch prototypes belonging to class c contained in cluster k; γ denotes a hyperparameter balancing the weights between input image patch prototypes and memory image patch prototypes.

6. The model-based self-distillation and prototype learning based semi-supervised medical image segmentation method of claim 1, wherein, The scheme ⑧ further includes the following contents: 8.1) For labeled images, predict a probability map f with the student model p and the image label y, the supervised loss is computed using the following formula: In the formula, L dice (·) represents the dice segmentation loss; L ce (·) represents the cross-entropy segmentation loss; 8.2) For the no-label images, predict a probability map f using the student model p and the image pseudo-label y p The consistency loss is computed using the following formula: 8.3) Utilizing an image patch affinity matrix A s and A t The self-distillation loss is computed using the following equation: L sd = KL(Softmax(A t / τ t )||Softmax(A s / τ s )) where KL(·) denotes the KL-divergence, τ t and τ s denote the temperature hyperparameters of the teacher and student models, respectively. 8.4) For class c in cluster k, contrastive features f are utilized by the student model co And the context-aware prototype, the context-aware prototype-based contrastive loss is calculated using the following formula: where τ denotes the temperature hyper-parameter, represents the selected pixels in the input image patches belonging to cluster k that are classified as c by the image labels or pseudo labels, and the pixels with higher uncertainty are selected to participate in the calculation; if some input image patches exist in all clusters, these image patches are divided into a new cluster containing all the memory bank image patches, and do not participate in the calculation of the contrast loss in the original cluster. The total context-aware prototype-based contrast loss is calculated using the following formula: 8.5) Calculate the total loss of the method using the following formula: L = L sup + L con + aL sd + βL caco Where α and β represent the weight of the self-distillation loss and the contrastive loss weight based on context-aware prototypes, respectively.

Images