Self-distillation image segmentation method based on graffiti marking

Abstract

The invention discloses a self-distillation image segmentation method and system based on graffiti annotation, and belongs to the technical field of image segmentation, images with graffiti annotation are used for training a neural network, a set of self-distillation framework is constructed, a teacher network and a student network have the same network structure but have different network weights, and the self-distillation framework is used for carrying out self-distillation image segmentation on the basis of the graffiti annotation. The weight of the teacher network is from the index moving average of the weight of the student network; according to the method, the student network is distilled from two levels of prediction results and middle layer features by using information in a high-quality noise-robust pseudo tag generated by the teacher network, so that the influence of noise in graffiti labeling is reduced, and the segmentation performance of a student model is improved. According to the method, the problem of low quality of the pseudo-tag generated for auxiliary training in the existing medical image weak supervision segmentation algorithm can be solved, the noise robustness of the pseudo-tag is improved, the quality of the pseudo-tag is further improved, and the performance of training the segmentation model based on graffiti labeling is improved.

Description

Technical Field

[0001] The present invention relates to the technical field of image segmentation, and in particular to a self-distillation image segmentation method based on graffiti annotation. Background Art

[0002] With the continuous development of medical imaging technology, medical imaging has become an important means to assist doctors in making accurate disease diagnosis. The accurate segmentation of target organs or lesions from images can help doctors make quantitative evaluations and provide strong support for accurate diagnosis of diseases and the formulation of treatment plans. However, manual segmentation of medical images is time-consuming and labor-intensive, which will waste the precious time of radiologists. Therefore, automatic segmentation methods have received more and more clinical attention, and automatic segmentation methods of medical images based on deep learning have become mainstream methods and have been clinically applied. However, the success of current deep learning methods mainly depends on training from a large number of annotated medical images, and the annotation of medical images is very difficult. It is difficult for ordinary people to complete accurate annotation, and professional doctors are required to complete it. However, due to the scarcity of professional doctors, the annotation of thousands of large-scale medical images has high time and labor costs, which greatly limits the development and application of deep learning medical image processing algorithms.

[0003] In order to reduce the reliance on densely labeled data, methods such as semi-supervised, self-supervised, unsupervised and weakly supervised are attracting more and more attention. Semi-supervised methods allow only a portion of images (such as 50% of the images) to be labeled, which can greatly reduce the amount of labeling. However, for large-scale medical images, the absolute number of labels required by this method is still very large. Self-supervised methods are a special type of semi-supervised method. This method first learns auxiliary tasks on unlabeled data, and then transfers the learning ability to the target task through a portion of labeled data. Unsupervised algorithms do not require manual labeling during the training phase. They generally cluster the training data to find the distribution pattern in the data. However, due to the lack of supervisory signals, this method is also difficult to achieve satisfactory performance.

[0004] Weakly supervised methods allow annotators to provide only weak annotations for an image, such as image-level labels, bounding boxes, or scribble annotations, which are easier to obtain than dense pixel-level annotations and can greatly reduce the annotation cost. In addition, compared with image-level labels and bounding boxes, scribble annotations can more flexibly annotate objects of various shapes, are more widely used, and the trained models usually have better performance. In order to learn segmentation from scribble annotations, existing methods are mainly divided into two categories: introducing regularization constraints for unlabeled pixels and generating pseudo labels for network training based on scribble annotations.

[0005] An investigation of the current methods for reducing the annotation cost of medical images found that although semi-supervised methods can significantly reduce the amount of image annotation, their annotation amount is still huge and has not completely reduced the image annotation cost. The existing clustering-based unsupervised methods are greatly limited in performance in the absence of supervision and are difficult to use in actual scenarios. In addition, although the current mainstream weak supervision method saves annotation costs by using graffiti annotations, it is easily disturbed by noise in graffiti annotations, resulting in inaccurate generated pseudo-labels, limiting the performance of the model, and ultimately resulting in its performance being lower than that of training using dense annotations, limiting its applicability in clinical practice. To this end, a more effective method is needed to generate pseudo-labels with higher accuracy using only graffiti annotation data, so as to further improve the segmentation performance of the model. Summary of the invention

[0006] In order to solve the above problems, the purpose of the present invention is to provide a self-distillation image segmentation technology based on graffiti annotation, aiming to overcome the problem that the existing automatic segmentation algorithm of medical images based on graffiti annotation is greatly affected by graffiti annotation noise and it is difficult to generate high-quality pseudo labels for model training.

[0007] In order to achieve the above technical objectives, the present application provides a self-distillation image segmentation method based on graffiti annotation, comprising the following steps:

[0008] Obtain a dataset that requires medical image segmentation, provide graffiti annotations for each category to be segmented in each image in the training set, and fully annotate each image in the test set as the gold standard for testing;

[0009] A neural network structure based on the teacher-student framework is constructed, trained with annotated data sets, and the knowledge of the teacher model is transferred to the student model using knowledge distillation. After obtaining the loss function of the student model, the trained student model is used as an image segmentation model for medical image segmentation.

[0010] Preferably, in the process of constructing a neural network structure based on a teacher-student framework, the teacher model and the student model have the same encoder-decoder structure, implemented by UNet, but have different model weights, wherein the exponential moving average of the student model weights is used as the weight of the teacher model.

[0011] Preferably, in the process of obtaining the weight of the teacher model, the weight of the teacher model is expressed as:

[0012]

[0013] In the formula, and are the weights of the student and teacher at the t-th and t-1-th iterations, respectively, and λ is a weight coefficient, which is obtained by the following formula:

[0014]

[0015] Among them, t represents the current iteration number, and 0.99 means that the weight λ of the exponential moving average gradually increases from 0 to 0.99.

[0016] Preferably, in the process of obtaining the loss function of the student model, the loss function of the student model is expressed as:

[0017] L=L pCE +α·L CRF +β t ·(L PD +L FD )

[0018] Where, L pCE represents the partial supervision loss function, L CRF represents the conditional random field loss function, L PD represents the knowledge distillation loss function, L FD represents the feature consistency loss, α represents the weight of the conditional random field loss function, and β t Represents the weight of the distillation loss function.

[0019] Preferably, in the process of obtaining the partial supervision loss function, the partial supervision loss function is expressed as:

[0020]

[0021] Where p is the probability map predicted by the student model, s is the one-hot encoding of the graffiti annotation, and ω s is a set of labeled pixels, and C represents the number of categories to be segmented; and are the probabilities that pixel i belongs to category c in network prediction and graffiti annotation, respectively.

[0022] Preferably, in the process of obtaining the conditional random field loss function, the conditional random field loss function is expressed as:

[0023]

[0024] In the formula, Φ represents the set of pixel pairs in the image, G(i, j) represents the Gaussian kernel bandwidth filter, and H i and I i The table represents the spatial position and spectral value of pixel i, ω is the normalized weight, σ H and σ I is a hyperparameter that controls the size of the Gaussian kernel.

[0025] Preferably, in the process of obtaining the knowledge distillation loss function, the knowledge distillation loss function is expressed as:

[0026]

[0027] In the formula, is the predicted score logical value output of the teacher model on the cth channel at the nth pixel, is the probability value of the soft label on the cth channel; τ is a hyperparameter that controls the smoothness, represents the soft prediction obtained by the student model through τ-Sofmax, N represents the number of pixels in the image, and C represents the number of target categories in the segmentation task.

[0028] Preferably, in the process of obtaining the feature consistency loss, the feature consistency loss is expressed as:

[0029]

[0030] Where D represents the number of scales based on upsampling in the decoder, and the feature maps of the teacher and student models at the dth scale are represented by and To represent, sim represents cosine similarity, γ d is the weight of the d-th level scale.

[0031] Preferably, in the process of obtaining the weight of the distillation loss function, the weight of the distillation loss function is expressed as:

[0032]

[0033] In the formula, t represents the current iteration number, t max represents the maximum number of iterations of the ramp-up function, β * is the control β t The maximum value of the hyperparameter.

[0034] The present invention discloses a self-distillation image segmentation system based on graffiti annotation, which is used to implement the above-mentioned self-distillation image segmentation method based on graffiti annotation. The system includes:

[0035] The data collection and annotation module is used to obtain the data set that needs medical image segmentation, provide graffiti annotation for each category to be segmented in each image in the training set, and fully annotate each image in the test set as the test gold standard;

[0036] The image segmentation module is used to build a neural network structure based on the teacher-student framework. It is trained through annotated data sets, and the knowledge of the teacher model is transferred to the student model by knowledge distillation. After obtaining the loss function of the student model, the trained student model is used as an image segmentation model for medical image segmentation.

[0037] The present invention discloses the following technical effects:

[0038] Compared with the existing weakly supervised training methods, the student model in the present invention learns from the teacher model from two levels: prediction results and intermediate layer features, making full use of the information in the pseudo-labels that are robust to noise, which can greatly improve the segmentation performance of the student model.

[0039] The present invention proposes to use the exponential moving average method in the training process, and use the exponential moving average of the student model weights as the weight of the teacher model. This method avoids additional training of a teacher model while ensuring that the teacher model has high segmentation performance, which can greatly simplify the training process and improve efficiency. BRIEF DESCRIPTION OF THE DRAWINGS

[0040] In order to more clearly illustrate the embodiments of the present invention or the technical solutions in the prior art, the drawings required for use in the embodiments will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying creative work.

[0041] Figure 1 is a comparison diagram of graffiti annotation and pixel-level dense annotation in the present invention, wherein black in the graffiti annotation represents unannotated pixels, and different shades of gray represent different categories of foreground graffiti;

[0042] Figure 2 A schematic diagram of the training method flow in the present invention;

[0043] Figure 3 Comparison of the segmentation results of two segmentation targets (placenta and fetal brain) in magnetic resonance images between the present invention and other methods, wherein gray and white represent segmentation results of two different categories. DETAILED DESCRIPTION

[0044] To make the purpose, technical scheme and advantages of the embodiments of the present application clearer, the technical scheme in the embodiments of the present application will be clearly and completely described below in conjunction with the drawings in the embodiments of the present application. Obviously, the described embodiments are only part of the embodiments of the present application, rather than all of the embodiments. The components of the embodiments of the present application generally described and shown in the drawings here can be arranged and designed in various different configurations. Therefore, the following detailed description of the embodiments of the present application provided in the drawings is not intended to limit the scope of the application claimed for protection, but merely represents the selected embodiments of the present application. Based on the embodiments of the present application, all other embodiments obtained by those skilled in the art without making creative work belong to the scope of protection of the present application.

[0045] like Figure 1-3 As shown, the present invention provides a self-distillation image segmentation method based on graffiti annotation, comprising the following steps:

[0046] Step 1: Collection of training set data based on graffiti annotation:

[0047] For a medical image segmentation task, a batch of images are collected as training sets and another batch as test sets. Each image in the training set is provided with scribble annotations for each category to be segmented, and each image in the test set is fully annotated as the test gold standard. A batch of mid-pregnancy fetal magnetic resonance images are collected, with an image size of 256×256. The training set is annotated with scribbles, and the validation set and test set are annotated with pixel-level dense annotations. Figure 1 A comparison of a set of images and their corresponding scribble annotations and pixel-level dense annotations is shown.

[0048] Step 2: Building the teacher-student segmentation network:

[0049] In order to effectively learn from graffiti annotation, the present invention adopts a neural network structure based on a teacher-student framework. The teacher and student models have the same encoder-decoder structure but different model weights. S and θ T Represent the weights of the student and teacher models respectively, and the exponential moving average of the student model weight is used as the weight of the teacher model, as shown below:

[0050]

[0051] in and are the weights of the teacher and the student at the t-th and t-1-th iterations, respectively. λ is a weight coefficient, obtained by the following formula:

[0052]

[0053] Among them, t represents the current iteration number, and 0.99 means that the weight λ of the exponential moving average gradually increases from 0 to 0.99.

[0054] Step 3: Construction of knowledge distillation loss function:

[0055] For unlabeled pixels, knowledge distillation is used to transfer the knowledge of the teacher model to the student model. First, the temperature-calibrated Softmax (τ-Sofmax) is used to obtain the segmentation soft labels predicted by the teacher model to supervise the training of the student model, as shown below.

[0056]

[0057] in is the predicted score logical value (logit) output of the teacher model on the cth channel at the nth pixel, is the probability value of the soft label on the cth channel. τ is a hyperparameter that controls the smoothness, where a larger τ leads to a smoother output, i.e., the probability of each category is less likely to be 0.0 or 1.0. Similarly, using To represent the soft predictions obtained by the student model through τ-Sofmax. The knowledge distillation loss function is defined as:

[0058]

[0059] Where N represents the number of pixels in the image, and C represents the number of target categories in the segmentation task.

[0060] Step 4: Construction of feature consistency loss:

[0061] In addition to performing knowledge distillation at the output level of the teacher and student models, the present invention also uses multi-scale feature consistency to encourage the student to imitate the teacher's feature extraction ability. Let D represent the number of scales based on upsampling in the decoder, and the feature maps of the teacher and student models at the dth scale are represented by and To represent, d = 0 represents the last layer of the encoder, that is, the bottleneck layer, and the larger d is, the farther away from the bottleneck layer. The feature consistency loss function is defined as:

[0062]

[0063] Where sim represents cosine similarity, γ d is the weight of the d-th level scale.

[0064] Step 5: Construction of partial supervision loss function:

[0065] For the labeled pixels given in the graffiti annotation, a partial cross entropy loss function is used to learn from the labeled pixels, which is defined as follows:

[0066]

[0067] Where p is the probability map predicted by the student model, s is the one-hot encoding of the graffiti annotation, and ω s is a set of labeled pixels, and C represents the number of categories to be segmented. and are the probabilities that pixel i belongs to category c (the cth channel represents the cth category) in network prediction and graffiti annotation, respectively.

[0068] Since applying only the partial cross entropy loss function on sparse annotations cannot provide sufficient boundary supervision, the present invention further uses the conditional random field loss function to regularize the network output, which can encourage the prediction consistency between adjacent pixels with similar features:

[0069] Where Φ represents the set of pixel pairs in the image, and G(i, j) represents the Gaussian kernel bandwidth filter, as shown below:

[0070]

[0071] Among them, H i and I i The table represents the spatial position and spectral value of pixel i, ω is the normalized weight, σ H and σ I is a hyperparameter that controls the size of the Gaussian kernel.

[0072] Step 6: Model training based on mixed loss function:

[0073] Based on the partial supervision loss function and knowledge distillation loss of the student model in the above steps, the complete loss function is defined as:

[0074] L=L pCE +α·L CRF +β t ·(L PD +δL FD );

[0075] Where α is the weight of the conditional random field loss function, β t is the weight of the distillation loss function, and δ is the weight of the feature-level distillation loss function. t The sigmoid ramp function increases to a maximum value as shown below:

[0076]

[0077] Where t represents the current iteration number, t max represents the maximum number of iterations of the ramp-up function, β *is the control β t The maximum value of the hyperparameter is β. Since the exponential moving average weight of the continuous student model is used as the teacher model, the teacher model has poor performance in the early stage of training. As the training progresses, the performance of the teacher gradually improves. In this process, β t Gradually increasing , is conducive to better utilizing the pseudo labels generated by the teacher model to supervise the training of the student model.

[0078] When testing, you only need to send the test data to the student model to get the final prediction result. Figure 3 This is a comparison of the results of placenta and fetal brain segmentation in fetal magnetic resonance images by the present invention and other methods (including MLoss and pCE methods). It can be seen that the segmentation results of the present invention method are relatively close to the gold standard, while the segmentation results of other methods have a large deviation compared with the gold standard.

[0079] The present invention is described with reference to flowcharts and / or block diagrams of methods, devices (systems), and computer program products according to embodiments of the present invention. It should be understood that each process and / or block in the flowchart and / or block diagram, as well as the combination of processes and / or blocks in the flowchart and / or block diagram, can be implemented by computer program instructions. These computer program instructions can be provided to a processor of a general-purpose computer, a special-purpose computer, an embedded processor, or other programmable data processing device to produce a machine, so that the instructions executed by the processor of the computer or other programmable data processing device generate instructions for implementing the processes in the flowchart and / or block diagram. Figure 1 A process or multiple processes and / or boxes Figure 1 A device that provides the functions specified in a block or multiple blocks.

[0080] In the description of the present invention, it should be understood that the terms "first" and "second" are used for descriptive purposes only and should not be understood as indicating or implying relative importance or implicitly indicating the number of technical features indicated. Thus, a feature defined as "first" or "second" may explicitly or implicitly include one or more of the features. In the description of the present invention, the meaning of "plurality" is two or more, unless otherwise clearly and specifically defined.

[0081] Obviously, those skilled in the art can make various changes and modifications to the present invention without departing from the spirit and scope of the present invention. Thus, if these modifications and variations of the present invention fall within the scope of the claims of the present invention and their equivalents, the present invention is also intended to include these modifications and variations.

Claims

1. A self-distillation image segmentation method based on graffiti annotation, characterized in that: The following steps are involved: Obtain a dataset that requires medical image segmentation, provide graffiti annotations for each category to be segmented in each image in the training set, and fully annotate each image in the test set as the gold standard for testing; A neural network structure based on the teacher-student framework is constructed, trained with annotated data sets, and the knowledge of the teacher model is transferred to the student model by knowledge distillation. After obtaining the loss function of the student model, the trained student model is used as an image segmentation model for medical image segmentation.

2. The self-distillation image segmentation method based on graffiti annotation according to claim 1, characterized in that: In the process of constructing a neural network structure based on a teacher-student framework, the teacher model and the student model have the same encoder-decoder structure, implemented by UNet, but have different model weights, wherein the exponential moving average of the student model weights is used as the weight of the teacher model.

3. The self-distillation image segmentation method based on graffiti annotation according to claim 2, characterized in that: In the process of obtaining the weight of the teacher model, the weight of the teacher model is expressed as: In the formula, and are the weights of the student and teacher at the t-th and t-1-th iterations, respectively, and λ is a weight coefficient, which is obtained by the following formula: Among them, t represents the current iteration number, and 0.99 means that the weight λ of the exponential moving average gradually increases from 0 to 0.

99.

4. The self-distillation image segmentation method based on graffiti annotation according to claim 3, characterized in that: In the process of obtaining the loss function of the student model, the loss function of the student model is expressed as: L=L pCE +α·L CRF +b t ·(L PD +L FD ) Where, L pCE represents the partial supervision loss function, L CRF represents the conditional random field loss function, L PD represents the knowledge distillation loss function, L FD represents the feature consistency loss, α represents the weight of the conditional random field loss function, and β t Represents the weight of the distillation loss function.

5. The self-distillation image segmentation method based on graffiti annotation according to claim 4, characterized in that: In the process of obtaining the partial supervision loss function, the partial supervision loss function is expressed as: Where p is the probability map predicted by the student model, s is the one-hot encoding of the graffiti annotation, and ω s is a set of labeled pixels, and C represents the number of categories to be segmented; and are the probabilities that pixel i belongs to category c in network prediction and graffiti annotation, respectively.

6. The self-distillation image segmentation method based on graffiti annotation according to claim 5, characterized in that: In the process of obtaining the conditional random field loss function, the conditional random field loss function is expressed as: In the formula, Φ represents the set of pixel pairs in the image, G(i, j) represents the Gaussian kernel bandwidth filter, and H i and I i The table represents the spatial position and spectral value of pixel i, ω is the normalized weight, σ H and σ I is a hyperparameter that controls the size of the Gaussian kernel.

7. The self-distillation image segmentation method based on graffiti annotation according to claim 6, characterized in that: In the process of obtaining the knowledge distillation loss function, the knowledge distillation loss function is expressed as: In the formula, is the predicted score logical value output of the teacher model on the cth channel at the nth pixel, is the probability value of the soft label on the cth channel; τ is a hyperparameter that controls the smoothness, represents the soft prediction obtained by the student model through τ-Sofmax, N represents the number of pixels in the image, and C represents the number of target categories in the segmentation task.

8. The self-distillation image segmentation method based on graffiti annotation according to claim 7, characterized in that: In the process of obtaining the feature consistency loss, the feature consistency loss is expressed as: Where D represents the number of scales based on upsampling in the decoder, and the feature maps of the teacher and student models at the dth scale are represented by and To represent, sim represents the cosine similarity, and γd is the weight of the d-th level scale.

9. The self-distillation image segmentation method based on graffiti annotation according to claim 8, characterized in that: In the process of obtaining the weight of the distillation loss function, the weight of the distillation loss function is expressed as: In the formula, t represents the current iteration number, t max represents the maximum number of iterations of the ramp-up function, and β* is the control β t The maximum value of the hyperparameter.

10. A self-distillation image segmentation system based on graffiti annotation, the system is used to implement the self-distillation image segmentation method based on graffiti annotation as claimed in any one of claims 1 to 9, characterized in that: The system includes: The data collection and annotation module is used to obtain the data set that needs medical image segmentation, provide graffiti annotation for each category to be segmented in each image in the training set, and fully annotate each image in the test set as the test gold standard; The image segmentation module is used to construct a neural network structure based on the teacher-student framework, train it through annotated data sets, transfer the knowledge of the teacher model to the student model by using knowledge distillation, and after obtaining the loss function of the student model, use the trained student model as an image segmentation model for medical image segmentation.

Images