Adaptive method and system of medical image segmentation model in test stage

Abstract

The invention discloses a self-adaption method and system for a medical image segmentation model in a test stage, and belongs to the technical field of image segmentation, and the method comprises the steps: training an image segmentation model through employing a training image and a segmentation label, taking the trained model as a source domain model, deploying the source domain model to a target domain, and obtaining a target image segmentation model; according to the method, images with good prediction are automatically screened out, a dynamic feature library is constructed, primary prediction is refined by using a method based on feature affinity, semantic supervision and boundary supervision are performed, and an adaptive teacher-student strategy is provided to enhance the adaptability of the model. According to the invention, the model can be adjusted and optimized in real time in a dynamic segmentation task test scene, and the adaptability of the model to different feature distribution changes is enhanced, so that the image segmentation quality in the test scene with complex changes is improved.

Description

Technical Field

[0001] The present invention relates to the technical field of image segmentation, and in particular to an adaptive method and system for a medical image segmentation model in a testing phase. Background Art

[0002] With the rapid development of computer technology and artificial intelligence, medical image analysis, especially image segmentation technology, plays an increasingly important role in clinical diagnosis and treatment planning. Medical image segmentation, especially in a variety of medical imaging technologies such as magnetic resonance imaging (MRI), computed tomography (CT), ultrasound and X-ray, is used to identify and quantify various structures in the body, such as organs and lesion areas. Accurate medical image segmentation is essential for early diagnosis of diseases, evaluation of treatment effects, and radiotherapy and surgical planning.

[0003] However, due to the complexity of medical images and the variability in their acquisition process, such as imaging protocols of different devices and physiological differences between patients, automatic medical image segmentation faces many challenges. Traditional image segmentation methods, such as threshold segmentation, region growing, boundary tracking, and model-based methods, usually rely on manually set parameters or rules, which is not only time-consuming but also difficult to adapt to different clinical application scenarios.

[0004] In recent years, deep learning technology has made significant progress in the field of medical image segmentation due to its powerful feature learning and representation capabilities. With a large amount of training data, deep learning models can automatically learn complex functions to extract useful features from image data. However, the performance of these models often depends heavily on the quality and quantity of training data. In practical applications, especially in new clinical environments or when using new equipment, there are often significant distribution differences between training data and actual test data, resulting in poor model performance.

[0005] To solve this problem, domain adaptation techniques are expected to overcome the distribution differences between training and test data and improve the performance of the model on test data. However, such methods usually require access to data in both the source domain (the domain where the training data is located) and the target domain (the domain where the test data is located), which may be difficult to achieve in practice due to factors such as privacy protection, data transmission costs, and time constraints. Especially in the medical field, the sensitivity of patient data makes it more difficult to directly access source domain data.

[0006] Test-time adaptation (TTA) technology has become a key technology to solve the above problems because it does not require modifying the source model training process or accessing the source domain data. TTA aims to quickly adapt to new test data by making small adjustments to the model during the model deployment phase. However, most existing TTA methods focus on the adaptation of overall image-level recognition tasks (such as image classification) without addressing the fine-grained adaptation required for more complex dense prediction tasks (such as segmentation). In addition, these methods often ignore the problem of effective supervision in the absence of sufficient annotated data. Summary of the invention

[0007] In order to solve the problems of domain shift and insufficient dense supervision faced by existing test-time adaptation methods in medical image segmentation, especially the domain adaptation problem of relying only on unlabeled data in the target domain when there is no source domain data, the purpose of the present invention is to provide a test-time adaptation technology based on feature affinity refinement, which aims to adjust and optimize in real time in dynamic segmentation task test scenarios, enhance the adaptability to changes in different feature distributions, and thus improve the image segmentation quality in complex changing test scenarios.

[0008] In order to achieve the above technical objectives, the present application provides an adaptive method for a medical image segmentation model in a testing phase, comprising the following steps:

[0009] Constructing a trained medical image segmentation model as a source domain model and setting target domain data, wherein the medical image segmentation model is used to segment cardiac muscle, left ventricle and right ventricle in cardiac magnetic resonance images;

[0010] Based on the source domain model, the target domain data is segmented, sorted according to the soft neighborhood density index of the segmentation prediction results, and a feature library is constructed to correct the original prediction results of the source domain model;

[0011] The source domain model is used as a student model, and a teacher model is constructed based on the student model. The auxiliary pseudo-label loss is obtained according to the prediction result of the teacher model on the cardiac magnetic resonance image.

[0012] Based on the auxiliary pseudo-label loss, by constructing the pseudo-label semantic supervision loss and entropy consistency loss, the loss function of the student model is generated, which is used to update the loss gradient calculated for the student model. According to the updated results of the student model, the teacher model is updated, and the prediction results of the updated teacher model are used as the segmentation prediction results of the cardiac magnetic resonance image.

[0013] Preferably, in the process of obtaining the soft neighborhood density index, the soft neighborhood density index is expressed as:

[0014]

[0015] In the formula, Yes b The transpose of is C×HW, The shape is C×C, and the soft neighbor density Sb is The sum of the entropy, p b Represents the segmentation prediction result.

[0016] Preferably, in the process of constructing the feature library, the soft neighborhood density index of the segmentation prediction results of the target domain data is sorted, and the data with the lowest soft neighborhood density index is found and stored in the feature library. When the number of features in the feature library exceeds the library capacity, the features are updated according to the first-in-first-out strategy.

[0017] Preferably, in the process of correcting the original prediction result of the source domain model, the middle layer features of the cardiac magnetic resonance image are obtained, and the features most similar to the middle layer features in the feature library are searched according to the cosine similarity to correct the original prediction result.

[0018] Preferably, in the process of searching for the feature in the feature library that is most similar to the middle-layer feature, the feature Z′ that is most similar to the middle-layer feature zb b It is expressed as:

[0019]

[0020] Among them, <·> represents the inner product of the vector, |·| represents the modulus of the vector, represents the feature library, and z represents the feature of the feature library.

[0021] Preferably, in the process of updating the classroom model, the teacher model is updated in the following manner:

[0022]

[0023] In the formula, and f t are the weight parameters of the teacher and student models at time t, is a loss based on auxiliary pseudo-label The adaptive weight coefficient of .

[0024] Preferably, when obtaining the auxiliary pseudo-label loss In the process of It is expressed as:

[0025]

[0026] Where C is the total number of categories of the segmentation task, and p b，c Represent the probability of the cth category in the teacher's prediction results and the student's prediction, respectively.

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

[0028]

[0029]

[0030]

[0031] Where p′ b，c and p b，c Represent the probability of the cth category in the pseudo label and the original prediction, respectively, b The entropy graph is denoted as E(p b ), p′ b，c The entropy graph is denoted as E(p′ b ), N represents the total number of pixels in the image, n is a pixel index, and β and γ are hyperparameters that balance different loss terms.

[0032] Preferably, when obtaining p b The entropy graph is denoted as E(p b ) process, the entropy graph is recorded as E(p b ) is expressed as:

[0033]

[0034] The present invention discloses an adaptive system for a medical image segmentation model in a testing phase, comprising:

[0035] A model building and data setting module, used to build a trained medical image segmentation model as a source domain model and set target domain data, wherein the medical image segmentation model is used to segment cardiac muscle, left ventricle and right ventricle in cardiac magnetic resonance images;

[0036] The feature library construction module is used to segment the target domain data based on the source domain model, sort the data according to the soft neighborhood density index of the segmentation prediction results, and build a feature library to correct the original prediction results of the source domain model;

[0037] A teacher model building module is used to use the source domain model as a student model, and to build a teacher model based on the student model, and to obtain an auxiliary pseudo-label loss based on the prediction result of the teacher model on the cardiac magnetic resonance image;

[0038] The adaptive module in the testing phase is used to generate the loss function of the student model based on the auxiliary pseudo-label loss by constructing the pseudo-label semantic supervision loss and the entropy consistency loss, and is used to update the loss gradient calculated for the student model. According to the updated result of the student model, the teacher model is updated, and the prediction result of the updated teacher model is used as the segmentation prediction result of the cardiac magnetic resonance image.

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

[0040] (1) The present invention can adapt to the distribution changes of test images in real time by establishing a dynamic feature library and implementing a feature affinity refinement strategy, effectively improving the adaptability and accuracy of medical image segmentation. This method is particularly suitable for continuously arriving test batches, and can accurately capture and utilize high-quality predicted features, thereby improving the prediction performance of the overall model.

[0041] (2) This paper introduces a mean teacher framework based on adaptive exponential moving average, which guides the training of the student model by updating the teacher model in real time without accessing the source data, so as to quickly adapt to new test data. This strategy enhances the stability and robustness of the model in the face of rapid changes in data distribution.

[0042] (3) By combining semantic supervision and boundary supervision mechanisms, the present invention not only improves the accuracy of the segmentation results, but also enhances the model's ability to recognize complex boundaries. In particular, when processing small structures or fuzzy boundaries commonly seen in medical images, the region can be delineated more accurately. BRIEF DESCRIPTION OF THE DRAWINGS

[0043] 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.

[0044] Figure 1 Schematic diagram of the adaptation process during model testing in the present invention;

[0045] Figure 2 Schematic diagram of the prediction results of the primary prediction, refined prediction and teacher model in the present invention, and a schematic diagram of the entropy diagram;

[0046] Figure 3 The segmentation results of the present invention on the cardiac MRI dataset; (a) is the test image; (b) is the gold standard of the corresponding image. DETAILED DESCRIPTION

[0047] 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.

[0048] like Figure 1-3 As shown, the present invention provides a test-time adaptation method for a medical image segmentation model, which automatically screens out images with good predictions and constructs a dynamic feature library, refines primary predictions using a feature affinity-based method, performs semantic supervision and boundary supervision, and proposes an adaptive teacher-student strategy to enhance the adaptability of the model; specifically, the following steps are included:

[0049] Step 1: Acquisition of source domain model:

[0050] For an image segmentation task, a batch of images are collected in a data domain as training images. An image segmentation model is trained using this batch of training images and their segmentation annotations, and the trained model is used as the source domain model. A batch of cardiac magnetic resonance images from multiple centers (multiple centers means multiple medical centers) are collected. For the segmentation tasks of cardiac muscle, left ventricle, and right ventricle, a UNet model is trained on the data from one center (source domain) as the source domain model, and the data from other centers are used as the target domain data.

[0051] Step 2: Construction of feature library on target domain:

[0052] After the source domain model is deployed to the target domain, a dynamic feature library is established based on the test data of the target domain to store and update the features of test images with good prediction performance. The feature library measures the prediction quality of each test image based on the soft neighborhood density, selects the image with the highest prediction quality, and stores the features encoded by its encoder in the feature library. The feature library is dynamically updated in a first-in-first-out manner.

[0053] Step 3: Correction of prediction results based on feature library:

[0054] For a test image, based on the prediction results of the source domain model, by comparing the similarity between its features and the features stored in the feature library, the most similar feature in the feature library is selected for feature replacement, and this feature is decoded to generate a revised prediction result.

[0055] Step 4: Construction of adaptive teacher model:

[0056] The source model is expanded into a teacher-student model. The optimization of the teacher model adopts an adaptive exponential moving average strategy. By adjusting the speed of model weight updates, it quickly responds to changes in test data, thereby maintaining the stability and robustness of the model in continuous test batches.

[0057] Step 5: Establishment of pseudo-label supervision and entropy consistency loss:

[0058] The improved prediction results after feature replacement are used as pseudo labels, and the original prediction results are supervised at the semantic level to construct pseudo label supervision loss; at the same time, the prediction results of the teacher model are used as auxiliary supervision signals to construct auxiliary supervision loss. In addition, the entropy value of each pixel is calculated on the pseudo label map and the original prediction result map, and the two entropy maps are constrained to be consistent as segmentation boundary supervision information.

[0059] Step 6: Update of segmentation model:

[0060] Based on the loss function in step 5, the gradient of the parameters in the student model is calculated and the parameters are updated in the direction of gradient descent.

[0061] The method for constructing the feature library in step 2 is as follows:

[0062] The image segmentation model is expressed as Where h represents the feature extractor and g represents the segmentation prediction head. represents a batch of test images at the tth moment (a total of B images), where the bth image is represented by x b , and b=1,2,…,b=1,2,…,B. x b The intermediate feature is represented as z b =h(x b ), the segmentation prediction result is expressed as p b =g(z b ). b The size of p is C×HW, where C is the number of classes to be segmented and HW is the image size. b The quality of the soft neighborhood density index S b To measure:

[0063]

[0064] in, Yes b The transpose of , whose shape is C×HW. The shape is C×C, and the soft neighbor density S b yes The sum of the entropy of .

[0065] For all test images in a batch, their soft neighbor density S b Sort and find S b The lowest image, its feature z b Store in feature library In. Feature Library The capacity is fixed at L. When the number of features in the feature library exceeds this capacity, the features follow the first-in-first-out strategy. to update.

[0066] In step 3, the prediction result correction method based on the feature library is:

[0067] For a test image x b The intermediate layer features z b , according to the cosine similarity, find the most similar feature z′ in the feature library b :

[0068]

[0069] Where <·> represents the inner product of the vector, and |·| represents the modulus of the vector. b Replace the original feature z b , and input it into the segmentation head g to obtain the segmentation prediction result p′ b =g(z′ b ), and take it as x b The original prediction result p on b The correction result.

[0070] The construction method of the adaptive teacher model in step 4 is:

[0071] In the above segmentation model Building a teacher model based on in represents the feature extractor of the teacher, Represents the teacher's segmentation prediction head. Its update method is:

[0072]

[0073] in and f t are the weight parameters of the teacher and student models at time t respectively. is a loss based on auxiliary pseudo-label The adaptive weight coefficient of x b Input the teacher model and get the prediction results As auxiliary pseudo-label. Auxiliary pseudo-label loss for:

[0074]

[0075] Where C is the total number of categories of the segmentation task, and p b，c Represent the probability of the cth category in the teacher's prediction results and the student's prediction, respectively.

[0076] The pseudo-label supervision and entropy consistency loss in step 5 are established as follows:

[0077] First, p′ b As x b Pseudo-labels, construct pseudo-label semantic supervision loss:

[0078]

[0079] where p′ b，c and p b，c represent the probability of the cth category in the pseudo label and the original prediction, respectively.

[0080] Then, the entropy consistency loss is constructed. b The entropy graph is denoted as E(p b ), the calculation process is:

[0081]

[0082] Similarly, p′ b，c The entropy graph is denoted as E(p′ b ), in E(p b ) and E(p′ b ) establish a consistency constraint loss between them:

[0083]

[0084] Where N represents the total number of pixels in the image and n is a pixel index. The total loss function of the student model is:

[0085]

[0086] Where β and γ are hyperparameters that balance different loss terms.

[0087] Embodiment: Based on the above-mentioned adaptive method for testing a medical image segmentation model designed by the present invention, the present invention provides the following embodiment of adaptive segmentation during testing of multimodal heart and prostate MRI images. In the above-mentioned step 2, B=10 frames, C=4, L=40 are set; in step 5, β is taken as 5, and γ is taken as 1. Figure 2 is the primary prediction p in this embodiment b , refine the prediction p′ b And the prediction results of the teacher model By comparing the entropy graph, we can see that p b There are many wrong prediction areas in p′ b and The segmentation result in p is more accurate, and p b and p′ b The areas with larger entropy are more concentrated in the edge areas of the target. Figure 3 This is a comparison of the segmentation results of cardiac MRI in the target domain of the proposed method with the gold standard. It can be seen that the segmentation results are closer to the gold annotations.

[0088] 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.

[0089] 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.

[0090] 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. An adaptive method for a medical image segmentation model in the testing phase, characterized in that: The following steps are involved: Constructing a trained medical image segmentation model as a source domain model and setting target domain data, wherein the medical image segmentation model is used to segment cardiac muscle, left ventricle and right ventricle in a cardiac magnetic resonance image; Based on the source domain model, the target domain data is segmented, sorted according to the soft neighborhood density index of the segmentation prediction results, and a feature library is constructed to correct the original prediction results of the source domain model; The source domain model is used as a student model, and a teacher model is constructed according to the student model, and an auxiliary pseudo-label loss is obtained according to the prediction result of the teacher model on the cardiac magnetic resonance image; Based on the auxiliary pseudo-label loss, by constructing the pseudo-label semantic supervision loss and the entropy consistency loss, the loss function of the student model is generated, which is used to update the loss gradient calculated for the student model. According to the updated result of the student model, the teacher model is updated, and the prediction result of the updated teacher model is used as the segmentation prediction result of the cardiac magnetic resonance image.

2. The method for self-adapting a medical image segmentation model in a testing phase according to claim 1, characterized in that: In the process of obtaining the soft neighborhood density index, the soft neighborhood density index is expressed as: In the formula, Yes b The transpose of is C×HW, The shape is C×C, and the soft neighbor density S b yes The sum of the entropy, p b Represents the segmentation prediction result.

3. The method for self-adapting a medical image segmentation model in a testing phase according to claim 2, characterized in that: In the process of constructing the feature library, the soft neighborhood density index of the segmentation prediction results of the target domain data is sorted, and the data with the lowest soft neighborhood density index is found and stored in the feature library. When the number of features in the feature library exceeds the library capacity, the features are updated according to the first-in-first-out strategy.

4. The method for self-adapting a medical image segmentation model in a testing phase according to claim 3, characterized in that: In the process of correcting the original prediction result of the source domain model, the middle layer features of the cardiac magnetic resonance image are obtained, and the features most similar to the middle layer features in the feature library are searched according to the cosine similarity to correct the original prediction result.

5. The method for self-adapting a medical image segmentation model in a testing phase according to claim 4, characterized in that: In the process of searching for the most similar feature in the feature library to the middle-layer feature, b The most similar feature Z′ b It is expressed as: Among them, <·> represents the inner product of the vector, |·| represents the modulus of the vector, represents the feature library, and z represents the feature of the feature library.

6. The method for self-adapting a medical image segmentation model in a testing phase according to claim 5, characterized in that: In the process of updating the classroom model, the teacher model is updated in the following manner: In the formula, and f t are the weight parameters of the teacher and student models at time t, is a loss based on auxiliary pseudo-label The adaptive weight coefficient of .

7. The method for self-adapting a medical image segmentation model in a testing phase according to claim 6, characterized in that: In obtaining the auxiliary pseudo-label loss In the process of It is expressed as: Where C is the total number of categories of the segmentation task, and p b，c Represent the probability of the cth category in the teacher's prediction results and the student's prediction, respectively.

8. The method for self-adapting a medical image segmentation model in a testing phase according to claim 7, 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: Where p′ b，c and p b，c Represent the probability of the cth category in the pseudo label and the original prediction, respectively, b The entropy graph is denoted as E(p b ), p′ b，c The entropy graph is denoted as E(p′ b ), N represents the total number of pixels in the image, n is a pixel index, and β and γ are hyperparameters that balance different loss terms.

9. The method for self-adapting a medical image segmentation model in a testing phase according to claim 8, characterized in that: In getting p b The entropy graph is denoted as E(p b ) process, the entropy graph is recorded as E(p b ) is expressed as:

10. An adaptive system for a medical image segmentation model during a testing phase, characterized in that: include: A model building and data setting module, used to build a trained medical image segmentation model as a source domain model and set target domain data, wherein the medical image segmentation model is used to segment cardiac muscle, left ventricle and right ventricle in cardiac magnetic resonance images; A feature library construction module, used to segment the target domain data based on the source domain model, sort the target domain data according to the soft neighborhood density index of the segmentation prediction results, and construct a feature library for correcting the original prediction results of the source domain model; A teacher model construction module, used to use the source domain model as a student model, and to construct a teacher model based on the student model, and to obtain an auxiliary pseudo-label loss based on the prediction result of the teacher model on the cardiac magnetic resonance image; The adaptive module in the test phase is used to generate the loss function of the student model based on the auxiliary pseudo-label loss by constructing the pseudo-label semantic supervision loss and the entropy consistency loss, and is used to update the loss gradient calculated for the student model. According to the updated result of the student model, the teacher model is updated, and the prediction result of the updated teacher model is used as the segmentation prediction result of the cardiac magnetic resonance image.

Images