A SAM-guided continuous testing adaptive approach for category-balanced teachers

Abstract

A continuous test-time adaptation method based on a SAM-guided category-balanced teacher, comprising the following steps: pre-training a source model with source domain data; predicting a target domain image through the source model to generate preliminary pseudo-labels; building and initializing a target model; predicting the target domain image through the SAM model; using a pseudo-label generation algorithm prompted by difficulty points to identify difficult areas of the target domain image and generate SAM pseudo-labels; using a pixel-level pseudo-label optimization algorithm to assign weights to pixels using prediction confidence and optimize preliminary pseudo-labels and SAM pseudo-labels; predicting the target domain image through the target model to generate target predictions; calculating the consistency loss between the target predictions and the SAM pseudo-labels; and then calculating the consistency loss between the target predictions and the preliminary pseudo-labels; adaptively training the target model through the consistency loss to reduce the error accumulation caused by noisy pseudo-labels; and calculating the category-balanced soft distillation loss through a category-balanced soft distillation algorithm guided by difficulty knowledge.

Description

Technical Field

[0001] The present invention relates to the field of medical image segmentation technology, and in particular to a continuous testing adaptation method based on a SAM-guided category-balanced teacher. Background Art

[0002] Medical image segmentation plays a key role in clinical applications such as computer-aided diagnosis and anatomical structure delineation. Recent advances in deep learning (DL) have significantly improved segmentation performance. However, due to domain shift caused by differences in imaging equipment, protocols, and operators, the generalization ability of models trained on source domain data is significantly reduced when applied to target domains unseen in real clinical scenarios.

[0003] Although domain adaptation (DA) technology can alleviate domain shift to a certain extent, it often requires a large amount of target domain data and generates a fixed model; test-time adaptation (TTA) only uses test data to adapt to the unknown target domain during the inference phase. However, TTA usually assumes that the target domain is a static environment, which is difficult to hold true in the constantly evolving clinical reality. To this end, continuous test-time adaptation (CTTA) came into being, which dynamically adapts to distribution shift during the inference phase. Existing CTTA methods are mainly divided into two categories: model-based methods and parameter-efficient methods. Model-based methods reduce error accumulation and catastrophic forgetting by optimizing self-supervised loss or resetting model parameters; while parameter-efficient methods use adapters or visual cues to update a small number of parameters.

[0004] The defects of existing technologies are: first, due to the strong dependence on the performance of the source model, it is easy to generate noisy pseudo-labels under significant domain offset, which in turn causes error accumulation and negative transfer; second, it fails to effectively solve the problem of category imbalance in medical images, that is, the area of ​​the foreground area where the key anatomical structure is located is much smaller than the background tissue, making the learning process easily dominated by the background.

[0005] Therefore, it is necessary to design a new CTTA method to solve the above technical difficulties. Summary of the Invention

[0006] In view of the above-mentioned defects of the prior art, the present invention provides a continuous test-time adaptation method based on a SAM-guided category-balanced teacher, which solves the problems of model performance degradation caused by the accumulation of pseudo-label errors in the field of medical image segmentation and the imbalance of foreground and background categories in medical images.

[0007] In order to achieve the above object, the technical solution adopted by the present invention is:

[0008] In the first aspect, a continuous test-time adaptation method for a category-balanced teacher based on SAM guidance includes the following steps:

[0009] S1. Build a source model and pre-train the source model using source domain data; predict target domain images using the source model to generate preliminary pseudo labels; build and initialize the target model;

[0010] S2. Predicting the target domain image using the SAM model;

[0011] The step S2 comprises:

[0012] S21, using a difficult point prompt pseudo label generation algorithm to identify difficult areas of the target domain image and generate a SAM pseudo label;

[0013] S22, optimizing the preliminary pseudo-label and the SAM pseudo-label by using a pixel-level pseudo-label optimization algorithm and assigning weights to pixels using prediction confidence;

[0014] S3, generating a target prediction by predicting the target domain image using the target model; calculating the consistency loss between the target prediction and the SAM pseudo-label; and then calculating the consistency loss between the target prediction and the preliminary pseudo-label; and adaptively training the target model using the consistency loss to reduce error accumulation caused by noisy pseudo-labels.

[0015] S4. Calculate the category-balanced soft distillation loss through the category-balanced soft distillation algorithm guided by difficult knowledge, which is used to guide the target model to perform foreground segmentation.

[0016] Preferably, step S21 includes the following steps:

[0017] According to the preliminary pseudo-label and the green channel image of the target domain image, a foreground area category prototype is calculated; the pixel farthest from the foreground area category prototype is selected as a difficulty point prompt; the geometric center of the largest connected foreground area in the preliminary pseudo-label is selected as a center point prompt; based on the difficulty point prompt and the center point prompt, a SAM model generates multiple candidate SAM pseudo-labels; the similarity between each candidate SAM pseudo-label and the preliminary pseudo-label is calculated and compared through the Dice coefficient to generate the SAM pseudo-label.

[0018] Preferably, the class balance soft distillation algorithm guided by difficult knowledge in step S4 comprises the following steps:

[0019] Calculate the foreground category weight according to the target pseudo label and the SAM pseudo label;

[0020] Constructing an uncertain region label based on a source model prediction probability map and an entropy map generated by predicting the target domain image by the source model;

[0021] The class-balanced soft distillation loss function is calculated using the foreground class weight and the uncertain region label.

[0022] Preferably, in step S22, the pixel-level pseudo-label optimization algorithm includes: receiving the prediction output of the target domain image; and mapping the probability value of each pixel in the prediction output using a Gaussian function to generate a weight map.

[0023] In a second aspect, an electronic device comprises a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein when the processor executes the computer program, the steps of the method provided in the first aspect or any possible implementation of the first aspect are implemented.

[0024] In a third aspect, a computer-readable storage medium stores a computer program thereon, wherein the computer-readable storage medium stores instructions, which, when the instructions are executed on a computer or a processor, cause the computer or processor to execute the method provided in the first aspect or any possible implementation of the first aspect.

[0025] Compared with the prior art, the beneficial effects of the present invention are embodied in:

[0026] 1. Unlike existing techniques that rely solely on the source model to generate pseudo-labels, which can easily generate a large number of noisy labels and cause error accumulation in scenarios with significant domain shift, this invention improves pseudo-label generation and optimization by leveraging the powerful prior knowledge of the SAM model, effectively alleviating the problems of pseudo-label noise and error accumulation.

[0027] 2. Unlike existing techniques that ignore class imbalance and cause the model to be easily dominated by background, this paper overcomes the impact of foreground-background imbalance through a class-balanced soft distillation mechanism, significantly improving the robustness and accuracy of medical image segmentation in adaptive scenarios during continuous testing. BRIEF DESCRIPTION OF THE DRAWINGS

[0028] Figure 1 is a schematic flow chart of the method of Example 1;

[0029] Figure 2 Schematic diagram of the method framework of Example 1. DETAILED DESCRIPTION

[0030] In order to make the technical means, creative features, objectives and effects of the invention easier to understand, the present invention is further described with reference to specific figures. However, the present invention is not limited to the following implementation cases.

[0031] It should be noted that the structures, proportions, sizes, etc. illustrated in the drawings in this specification are only used to match the contents disclosed in the specification so that people familiar with this technology can understand and read them. They are not used to limit the conditions under which the present invention can be implemented. Therefore, they have no substantive technical significance. Any structural modification, change in proportional relationship or adjustment of size should still fall within the scope of the technical content disclosed in the present invention without affecting the efficacy and purpose that can be achieved by the present invention.

[0032] Example 1:

[0033] like Figure 1 、 Figure 2 A SAM-guided category-balanced teacher framework (Sam-CBT) is shown, which aims to dynamically adjust the pre-trained source model. (Based on labeled source domain data Training, including is the image data, The label data corresponding to the image, the model pre-trained with labeled source domain data is collectively referred to as the source model) to cope with the unlabeled target domain data during deployment The adaptation process is performed online, the target domain data is only available once, and the source domain data is no longer accessible.

[0034] Step S1: Source model (teacher model) pre-training and preliminary pseudo-label generation

[0035] First, in the labeled source domain data Train the source model After training, the source model will be used to train the unlabeled target domain images. Make predictions and generate preliminary pseudo labels.

[0036] For each medical image in the target domain , first use the pre-trained source model Make a prediction and the output result is a soft prediction probability map The soft prediction refers to the model's output of the predicted probability that each pixel belongs to the foreground category, with a value range between 0 and 1. Then, after processing through the Sigmoid function, the soft prediction is converted into a hard pseudo label for supervised training. The formula is as follows:

[0037]

[0038] Among them, the threshold , is the indicator function. If the probability value is greater than 0.5, the corresponding pseudo label is 1 (indicating that the pixel belongs to the target area), otherwise it is 0 (indicating that the pixel does not belong to the target area).

[0039] Step S2: Pseudo-label optimization strategy (Sam-PR) guided by the SAM model (teacher model):

[0040] The Sam-PR strategy integration source model proposed in this invention , target model (Student Model) and SAM Model The source model generates preliminary pseudo labels , SAM generates SAM pseudo labels , the two jointly train the target model To improve pseudo-label quality and reduce error accumulation, we propose two strategies: the Hard Point Prompted Pseudo-Label Generation (HPPG) algorithm and the Pixel-Level Pseudo-Label Optimization (PPR) algorithm. The optimized pseudo-labels are further guided by hard distillation (directly using the optimized pseudo-labels as a supervisory signal for training).

[0041] Step S21: Difficult Point Prompt Pseudo-label Generation (HPPG):

[0042] The input of the HPPG algorithm is the category prototype and difficult area extracted by preliminary pseudo-labels, and the output is the optimal prediction (SAM generates ).

[0043] The HPPG algorithm includes the following steps:

[0044] 1. Category prototype calculation: Based on the source model pseudo-label and green channel image, calculate the prototype feature value of the foreground area of ​​each category ;

[0045] 2. Hint point extraction: Obtain difficult point hints by calculating the difference with the prototype ; and calculate the center point of the largest connected foreground area ;

[0046] 3. Input the cue point to the SAM model to generate candidate prediction results, and perform Dice similarity evaluation with the source model pseudo-label to select the optimal output. The details are as follows:

[0047] The category prototype extracted by preliminary pseudo-labels identifies the difficult areas in the target image (areas that are difficult to segment in the image). Green Channel Pixel value calculation foreground area category prototype Specifically, the category prototype is defined as Medium Category Corresponding foreground pixels exist Mean over the channel:

[0048]

[0049]

[0050] distance The farthest pixel is judged as the most difficult sample and is selected as the hard point prompt At the same time, the geometric center of the largest connected foreground area in the preliminary pseudo-label is selected as the center point hint . Combined with the center point prompt Tips on difficult points , SAM will generate multiple candidate prediction results. Finally, the Dice coefficient is used to measure the relationship between the candidate prediction and the source model pseudo label. similarity, select the best prediction SAM pseudo-labels as output.

[0051] Step S22: Pixel-level pseudo-label optimization (PPR):

[0052] The input of the PPR algorithm is the predicted probability map generated by the SAM model or the source model , the output is based on the pixel confidence weight , used for weighted distillation loss to suppress the influence of pseudo-label noise. The PPR algorithm is an existing technology.

[0053] The PPR algorithm includes the following steps:

[0054] 1. Get the predicted probability map: receive the predicted output of the target image , that is, the probability value of each pixel belonging to the foreground;

[0055] 2. Calculate pixel confidence weight: Use Gaussian function to calculate the probability value of each pixel Mapping and generating weight graph The details are as follows:

[0056] The weights of pixels are assigned by prediction confidence to reduce the influence of noisy pseudo labels. Specifically, pixels with higher confidence will be given higher weights. The confidence level of is calculated as follows:

[0057]

[0058] in 、 and are hyperparameters. These weights It is then integrated into the binary cross entropy loss and Dice loss functions to effectively suppress the instability caused by noisy pseudo labels during the model adaptation process.

[0059] Step S3: Consistency training with hard distillation loss

[0060] The input to the hard distillation loss algorithm is the target model prediction , preliminary pseudo-labels , SAM pseudo-label , and the corresponding confidence weight 、 , the output is the consistency loss used for target model training

[0061] Hard distillation losses include the following steps:

[0062] 1. Target model prediction generation: Input enhanced target image , through the target model Get the predicted output .

[0063] 2. Loss Calculation: What constitutes the eventual consistency loss , used for back-propagation to update the target model parameters and achieve cross-domain robust adaptive training.

[0064] The essence of this method is to calculate the consistency loss between the target model prediction results and the pseudo labels generated by Sam-PR optimization and the preliminary labels after PPR optimization. :

[0065]

[0066] in Represents the target model for the enhanced image The predicted output of is the weighted segmentation loss function after fusing pixel weights.

[0067] Step S4: Hard Knowledge Guided Class Balanced Soft Distillation (HKC-SD) Algorithm

[0068] The foreground-background imbalance in medical image segmentation is a core issue that needs to be addressed. Background classes usually dominate the pixel space, leading to the poor performance of traditional loss functions (e.g. ) focuses too much on the background area. To this end, the present invention proposes the HKC-SD method to address this imbalance problem.

[0069] Although it is a common strategy to solve the category imbalance by increasing the foreground loss weight, the presence of noisy pseudo labels in the CTTA scenario will lead to unstable weight distribution. This method innovatively combines SAM prediction with difficult category information to calculate category balance weights: by quantizing the target model prediction With SAM prediction The inconsistency between the two images guides the model to focus on the difficult-to-classify foreground areas.

[0070] The input of the HKC-SD algorithm is: target model prediction , SAM pseudo-label , Source model prediction probability map and its entropy graph ; The output is the class-balanced soft distillation loss , used to improve the foreground area segmentation performance.

[0071] HKC-SD includes the following steps:

[0072] 1. Dynamic estimation of class weights: Calculate the number of pixels that appear in the SAM prediction for each class and the number of pixels that differ from the target model prediction to form the class weight. ;

[0073] 2. Uncertain region identification: Fusing the predicted probability and entropy information of the source model, identifying pixel areas where the model lacks confidence, and constructing the uncertainty region ;

[0074] 3. Class-balanced loss calculation: Using dynamic class weights and uncertain region labels, construct and calculate the class-balanced soft distillation loss function , used to guide the model to focus on the foreground and improve the recognition ability of difficult areas. The details are as follows:

[0075] Foreground class weight The calculation formula is:

[0076]

[0077] in and are the height and width of the image respectively. The formula is calculated by counting the foreground pixels confirmed by SAM ( ) and the model prediction divergence area ( ), effectively capturing difficult samples such as false positives and false negatives. This dual statistical mechanism suppresses the interference of noisy pseudo-labels on weight distribution, ultimately improving foreground segmentation accuracy.

[0078] By introducing the HKC-SD loss , the present invention can learn richer probability distributions and apply category balance weights, thereby effectively solving the foreground-background imbalance problem:

[0079]

[0080] in Predicting entropy by fusion source model Identify the uncertain region. Here e(·) is the formula for calculating entropy, It's a prediction. is the entropy of the prediction, expressed as the uncertainty of the prediction. The design applies class weights to encourage knowledge transfer from uncertain distributions. HKC-SD improves foreground segmentation and addresses the class imbalance problem by focusing on high uncertainty regions.

[0081] By combining these modules, the Sam-CBT method overcomes the limitations of relying solely on pre-trained source models and unlabeled target domain data, enabling online adaptive updates to the continuously evolving distribution of the target domain. This overall approach not only improves pseudo-label generation and optimization but also mitigates the problems of pseudo-label noise accumulation and class imbalance through hard and soft distillation mechanisms, respectively, significantly enhancing the model's robustness and accuracy in medical image segmentation tasks.

[0082] Those skilled in the art will readily appreciate other embodiments of the present invention after considering the specific embodiments of this specification. The embodiments disclosed in this application are intended only to illustrate the principles of the present invention and do not limit the scope of protection of the present invention. Any modification, variation, or improvement of the contents described in this specification that does not depart from the spirit and principles of the present invention shall be included within the scope of protection of the present invention.

Claims

1. A continuous test-time adaptation method based on a SAM-guided category-balanced teacher, characterized by: The following steps are involved: S1. Build a source model and pre-train the source model using source domain data; predict the target domain image using the source model to generate preliminary pseudo labels; Build and initialize the target model; S2. Predicting the target domain image using the SAM model; The step S2 comprises: S21, using a difficult point prompt pseudo label generation algorithm to identify difficult areas of the target domain image and generate a SAM pseudo label; S22, optimizing the preliminary pseudo-label and the SAM pseudo-label by using a pixel-level pseudo-label optimization algorithm and assigning weights to pixels using prediction confidence; S3. Predicting the target domain image through the target model to generate a target prediction; calculating the consistency loss between the target prediction and the optimized SAM pseudo-label; and then calculating the consistency loss between the target prediction and the optimized preliminary pseudo-label; adaptively training the target model through the consistency loss between the target prediction and the optimized SAM pseudo-label and the consistency loss between the target prediction and the optimized preliminary pseudo-label to reduce the error accumulation caused by noisy pseudo-labels; S4. Calculate the category-balanced soft distillation loss through the category-balanced soft distillation algorithm guided by difficult knowledge, which is used to guide the target model to perform foreground segmentation.

2. The continuous test-time adaptation method of the SAM-guided category-balanced teacher according to claim 1, characterized in that The step S21 includes the following steps: Calculating a foreground region category prototype based on the preliminary pseudo-label and the green channel image of the target domain image; selecting a pixel farthest from the foreground region category prototype as a difficulty point prompt; Selecting the geometric center of the largest connected foreground area in the preliminary pseudo-label as the center point hint; According to the difficulty point prompt and the center point prompt, the SAM model generates a plurality of candidate SAM pseudo labels; The SAM pseudo-label is generated by calculating and comparing the similarity between each candidate SAM pseudo-label and the preliminary pseudo-label through the Dice coefficient.

3. The continuous test-time adaptation method of the SAM-guided category-balanced teacher according to claim 1, characterized in that The class balance soft distillation algorithm guided by difficult knowledge in step S4 includes the following steps: Calculating foreground category weights based on the target prediction and the SAM pseudo-label; Constructing an uncertain region label based on a source model prediction probability map and an entropy map generated by predicting the target domain image by the source model; The class-balanced soft distillation loss function is calculated using the foreground class weight and the uncertain region label.

4. The method for continuous test-time adaptation of a class-balanced teacher based on SAM guidance according to claim 2, characterized in that In step S22, the pixel-level pseudo-label optimization algorithm includes: receiving the prediction output of the target domain image; and mapping the probability value of each pixel in the prediction output using a Gaussian function to generate a weight map.

5. An electronic device comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein: When the processor executes the computer program, the steps of the method according to any one of claims 1 to 4 are implemented.

6. A computer-readable storage medium having a computer program stored thereon, wherein the computer-readable storage medium stores instructions, characterized in that: When the instructions are executed on a computer or a processor, the computer or the processor is caused to perform the steps of the method according to any one of claims 1 to 4.

Images