Efficient three-dimensional medical image segmentation large model construction method based on single slice labeling

Abstract

The invention provides an efficient three-dimensional medical image segmentation large model construction method based on single slice labeling, and belongs to the technical field of medical image segmentation. According to the method, prompt segmentation is carried out on a single slice through a 2D segmentation large model, features of a segmentation result are propagated to other slices of a 3D medical image through a feature propagation model FPM, the 3D segmentation large model is constructed cooperatively, and three-stage progressive training is carried out on the 3D segmentation large model by using various 3D medical image data sets. Under the extreme sparse marking condition that only one slice is marked on each 3D image, the large 3D segmentation model of which the segmentation performance is comparable with or even exceeds that of a fully supervised model can still be trained. According to the method, the manual annotation cost is greatly reduced, the precision and generalization ability of the model in multiple medical image segmentation tasks are remarkably improved, and the method has high efficiency and practical value in actual medical image segmentation application and can be flexibly applied to different medical segmentation tasks.

Description

Technical Field

[0001] This invention belongs to the field of medical image segmentation technology, specifically involving an efficient training method for constructing a high-precision 3D segmentation model by combining a basic segmentation model and a feature propagation model when only a single slice is labeled in each 3D medical image. Background Technology

[0002] 3D medical image segmentation is a core task in medical image analysis, widely used in clinical diagnosis, surgical planning, and disease monitoring. Traditional 3D medical image segmentation methods rely on voxel-level dense annotation, which is extremely costly. For example, a 3D MRI or CT image typically contains hundreds of slices, and detailed annotation of each slice requires a significant investment of time and manpower from professional physicians.

[0003] In recent years, semi-supervised learning methods have alleviated the annotation cost problem to some extent by combining a small amount of labeled data with a large amount of unlabeled data. However, existing semi-supervised methods generally still require complete 3D annotations for each labeled image, which remains difficult to achieve in real clinical scenarios. In contrast, in real medical environments, doctors tend to annotate only key slices, resulting in sparse and discontinuous annotations.

[0004] With the emergence of large-scale segmentation models (such as the Segment Anything Model, SAM), efficient segmentation through prompt-driven methods has become possible. However, directly applying these models to 3D medical images still presents challenges, such as insufficient robustness to low-contrast structures and inadequate transfer of cross-slice information. Therefore, how to fully leverage the capabilities of basic models to achieve high-quality 3D segmentation with minimal annotation is a key research challenge. Summary of the Invention

[0005] To address the problems existing in the prior art, this invention proposes a method for constructing a high-efficiency 3D medical image segmentation large model based on single-slice annotation. This method uses a 2D segmentation large model to perform high-quality segmentation of a single slice, and then uses a Feature Propagation Model (FPM) to propagate the features of the segmentation results to other slices of the 3D medical image, thereby collaboratively constructing a 3D segmentation large model. The 3D segmentation large model is trained in three stages using multiple 3D medical image datasets, so that even with only one slice annotated for each 3D image, it can still achieve segmentation performance comparable to that of traditional fully supervised deep learning models.

[0006] The technical solution of the present invention is as follows: A method for constructing a large-scale 3D medical image segmentation model based on single-slice annotation is characterized by: using a 2D segmentation model to provide cue segmentation for a single slice; using a Feature Propagation Model (FPM) to propagate the features of the segmentation result to other slices of the 3D medical image to collaboratively construct the large-scale 3D segmentation model; and using multiple 3D medical image datasets to perform a three-stage progressive training on the large-scale 3D segmentation model. In each 3D medical image dataset, only one slice is annotated, with the rest being unannotated slices. The three-stage progressive training specifically includes: Phase 1: Semi-supervised pre-training of the Feature Propagation Model (FPM); First, a large 2D segmentation model is used to perform cue segmentation on a single labeled slice to generate pseudo-labels. Then, the pseudo-labels are used as supervision signals to perform semi-supervised pre-training of FPM. During the training process, each training session takes three frames of slices as input, where the first frame is a labeled reference slice, and the second and third frames are unlabeled or labeled slices. FPM first predicts the result of the second frame based on the reference slice, and then combines the first two frames to predict the third frame. Backpropagation is performed through cross-entropy and Dice loss to optimize the model parameters. Phase Two: Construction of the Pseudo-Label Library and Self-Training of the Feature Propagation Model (FPM); After initial pre-training, the pre-trained FPM from Phase 1 is used to perform bidirectional inference on unlabeled slices, i.e., inference is performed from the labeled slices to both the front and back directions of the 3D image to obtain the prediction result for each unlabeled slice. For each unlabeled slice, the prediction result of the nearest reference slice is selected as its initial pseudo-label. Initial pseudo-labels are generated for all unlabeled slices in each 3D medical image, and a pseudo-label library is constructed and stored. Then, the pseudo-label library is used as a supervision signal to self-train the FPM. During training, labeled slices are used as reference frames or pseudo-labels of unlabeled slices are used as new reference frames. At the same time, the pseudo-label library is dynamically updated during self-training. Phase 3: Fine-tuning of high-confidence pseudo-labels for large 2D segmentation models; The FPM trained in Phase 2 is used to perform bidirectional inference on the unlabeled slices again to obtain a new set of pseudo-labels. This set is compared with the final pseudo-label library in Phase 2, and high-confidence samples are selected from it to fine-tune the 2D segmentation model to adapt it to the feature distribution of medical images.

[0007] Furthermore, the three-frame slices used as input to the Feature Propagation Model (FPM) employ a three-frame triplet sampling rule, with each training iteration using three-frame triplets. Enter FPM, where This represents the number of labeled frames in the last two frames. This represents the number of unlabeled frames in the last two frames, with the following constraints: (1) The reference frame is always the annotation frame; (2) During training, there is a 50% probability that all frames are unlabeled and a 50% probability that one frame is labeled and one frame is unlabeled. (3) Skip sampling is used for unlabeled frames, and the skip step size is gradually increased linearly or segmentally as training progresses.

[0008] Furthermore, in the training of the Phase 1 and Phase 2 Feature Propagation Model (FPM), FPM performs backpropagation using cross-entropy and Dice loss. The loss function is specifically calculated as follows: Order No. The prediction results for labeled / unlabeled frames are: Marked as The pseudo-tag is The overall training loss of FPM is defined as: Among them, for frames with real annotations, the supervised loss for: In the formula This represents pixel-level cross-entropy loss. Indicates Dice loss: Using pseudo-labels to supervise unlabeled frames, the unsupervised loss is defined as: in This is an indicator function; it returns 1 if a condition is met, and 0 otherwise. (Threshold) Used to eliminate pre-existing confidence levels that are too low. Pixel measurement; pseudo-label In Phase 1, the loss is generated by a large 2D segmentation model, and in Phase 2, it is generated by the FPM prediction itself. The labels for the supervised loss in Phase 2 come from a pseudo-label library or real labels.

[0009] Furthermore, in the first stage, before generating pseudo-labels using the large 2D segmentation model, the large 2D segmentation model is fine-tuned. A bounding box selection enhancement strategy is introduced to randomly perturb the ground truth bounding boxes of the reference slices. During fine-tuning, labeled slices are used, and the original bounding boxes and several perturbed bounding boxes are used to provide hints to the large 2D segmentation model to generate prediction results. The loss function used is supervised loss.

[0010] Furthermore, during the self-training process of the Phase 2 Feature Propagation Model (FPM), the pseudo-label library is dynamically updated. The update rule is as follows: when the prediction confidence of the Feature Propagation Model (FPM) for a certain pixel exceeds the threshold and is inconsistent with the current pseudo-label, the pseudo-label of that pixel is updated.

[0011] Furthermore, the 2D segmentation large model fine-tuning in stage three uses the cross-entropy and Dice loss employed during the training of the feature propagation model FPM in stage one.

[0012] The technical effects of this invention are as follows: This invention proposes a highly efficient method for constructing a large-scale 3D medical image segmentation model based on single-slice annotation. It performs high-quality segmentation on a single slice and uses a Feature Propagation Model (FPM) to propagate the segmentation features to other slices of the 3D medical image, thereby collaboratively constructing a large-scale 3D segmentation model. The model is trained in a three-stage progressive manner using multiple 3D medical image datasets. Even under extremely sparse annotation conditions where only one slice is annotated per 3D image, it can still train a large-scale 3D segmentation model with segmentation performance comparable to or even exceeding that of a fully supervised model. This method not only significantly reduces the cost of manual annotation but also significantly improves the accuracy and generalization ability of the model in multiple medical image segmentation tasks, such as LMS (leg muscles), LA (heart), and KiTS19 (kidneys). Dice scores reach 91.43%, 91.38%, and 93.82%, respectively, and the HD95 metric is also superior to existing mainstream methods, fully validating the efficiency and practical value of this invention in real-world medical image segmentation applications. This invention can be flexibly applied to different medical segmentation tasks. Attached Figure Description

[0013] Figure 1 This is an overall flowchart of the efficient 3D medical image segmentation large model construction method based on single slice annotation of the present invention. Detailed Implementation

[0014] The present invention will be further clearly and completely described below with reference to the accompanying drawings and specific embodiments.

[0015] This invention presents an efficient method for constructing a large-scale 3D medical image segmentation model based on single-slice annotation, applicable to various 3D medical image datasets, such as cardiac MRI, abdominal CT, and muscle MRI data. Before training, for each 3D medical image, one slice is randomly selected as the labeled slice, and the corresponding annotation is retained. The remaining slices are considered unlabeled slices, meaning their corresponding annotation information is not used during training. This invention combines a 2D segmentation model with a spatial-temporal feature propagation model (FPM) to construct a 3D segmentation model. The 3D segmentation model is then trained in a three-stage progressive manner using various 3D medical image datasets. The specific steps include: Phase 1: Semi-supervised pre-training of the Feature Propagation Model (FPM), such as Figure 1As shown in Phase 1, a large 2D segmentation model, such as SAM (Segment Anything Model), is first used to perform cue segmentation on a single labeled slice, generating high-quality pseudo-labels. Then, these pseudo-labels are used as supervisory signals to perform semi-supervised pre-training of FPM. During training, each training iteration uses three frames of slices as input: the first frame is a labeled reference slice, and the second and third frames are unlabeled or labeled slices. FPM first predicts the result of the second frame based on the reference slice, then combines the results of the first two frames to predict the third frame. Backpropagation is performed using cross-entropy and Dice loss to optimize the model parameters.

[0016] The three-frame slices used as input employ a three-frame triplet sampling rule, with each training iteration using three-frame triplets. Enter FPM, where This represents the number of labeled frames in the last two frames. This represents the number of unlabeled frames in the last two frames, with the following constraints: (1) The reference frame is always the labeled frame (i.e., sampled from the selected single-label slice set); (2) During training, there is a 50% probability that all frames are unlabeled, and a 50% probability that one frame is labeled and one frame is unlabeled (i.e., =1, =1 or =0, =2); (3) Skip sampling is used for unlabeled frames. The skip step size is set in the early stage of training. As training gradually Linear or piecewise increases to (To reduce the risk of propagation of false label adjacency errors).

[0017] The Feature Propagation Model (FPM) performs backpropagation using cross-entropy and Dice loss. The loss function is calculated as follows: Order to the first The prediction results for labeled / unlabeled frames are: Marked as The pseudo-tag is The overall training loss of FPM is defined as: Among them, monitoring losses (For frames with actual annotations) the following: In the formula This represents pixel-level cross-entropy loss. Indicates Dice loss: Unsupervised loss (using pseudo-labels to supervise unlabeled frames) is defined as: in The indicator function (1 if the condition is met, 0 otherwise), threshold Used to exclude those with low confidence levels. Predicted pixels (taken here) Pseudo-tags In this stage, SAM (which generates prompts based on FPM predictions) is used. In the later stage two, pseudo-labels are generated. It is generated by FPM prediction itself.

[0018] To improve the pseudo-label generation quality of SAM under imperfect cues, this stage introduces a box-prompt augmentation strategy for fine-tuning. This strategy randomly perturbs the ground truth boxes of the reference slices to enhance SAM's generalization ability under cue noise, thereby generating more reliable pseudo-labels for FPM training.

[0019] To generate more robust SAM pseudo-labels from FPM prediction results, the bounding box suggestion enhancement strategy introduced in this invention employs the following steps: (1) For an unlabeled slice, the prediction result is obtained by inference using FPM. ; (2) Define pseudo-regions : That is, only those with a confidence level of 1 And non-background pixels are included in the pseudo region, among which Represents the row and column coordinates of pixels in an image ( For row index, (for column indexes) This indicates the pixel The predicted probability, This represents the maximum predicted probability value for that pixel across all categories. This indicates the category with the highest predicted probability for this pixel, denoted by [symbol]. It indicates the logical relation "AND"; (3) Calculate the tight wrapping Minimum bounding rectangle ;in Left of the rectangle Top corner coordinates The coordinates of the bottom right corner of the rectangle are used as the original bounding box. (4) To improve the robustness of SAM to inaccurate prompts, point jitter is used to perturb the frame corners. Several perturbed boxes are obtained, and the perturbation formula is: in The perturbation intensity hyperparameter (here) ), The direction of the disturbance. , .

[0020] SAM performs fine-tuning before generating pseudo-labels, using labeled slices and combining the original boxes with several sub-labels. The perturbed bounding boxes provide hints to the SAM to generate prediction results, and the loss function uses supervised loss. This improves the quality of pseudo-tags generated by SAM under noisy cues. During fine-tuning, only the mask decoder and point-prompt encoder of SAM are updated, while other parameters remain frozen.

[0021] Phase Two: Construction of the pseudo-label library and self-training of the feature propagation model (FPM), such as... Figure 1 As shown in Phase 2 of the diagram. After initial pre-training, the FPM pre-trained in Phase 1 is used to perform bidirectional inference on unlabeled slices, i.e., inference is performed from the labeled slices in both directions towards the front and back of the 3D image, obtaining the prediction result for each unlabeled slice. For each unlabeled slice, the prediction result of the nearest reference slice is selected as its initial pseudo-label; initial pseudo-labels are generated for all unlabeled slices in each 3D medical image, and a pseudo-label library is constructed and stored. Then, the pseudo-label library is used as a supervision signal to self-train the FPM. During training, it is no longer limited to using labeled slices as reference frames; the pseudo-labels of unlabeled slices can also be used as new reference inputs. At the same time, the pseudo-label library is dynamically updated during training: when the prediction confidence of the FPM model for a certain pixel exceeds a threshold and is inconsistent with the current pseudo-label, the pseudo-label of that pixel is updated to gradually improve the quality of the pseudo-labels.

[0022] Specifically, the pre-trained FPM from Phase 1 is used to perform bidirectional inference on the unlabeled slices of each training image, as follows: a. Using the labeled slices as a reference, reason forward and backward to obtain the prediction result for each unlabeled slice. Forward prediction results and backward prediction results: t represents any unlabeled slice; b. For any unlabeled slice The prediction result generated by the "most recent labeled reference frame" is selected as its initial pseudo-label. Sign (i.e., select the prediction result from the reference frame with the smallest distance); c. Store the set of pseudo-labels for all unlabeled slices into a pseudo-label bank (denoted as ). ), labeled slices Fill directly with the actual annotations.

[0023] Specifically, the FPM is self-trained using a pseudo-label library as a supervision signal, as follows: Use pseudo-tag library The pseudo-labels in the dataset serve as supervision signals, and the pre-trained FPM from Phase 1 is self-trained using the same training method as in Phase 1. The training loss form remains consistent with Phase 1, but the supervision loss... The annotations can come from a pseudo-label library instead of human annotations. If a frame is a real annotation (human annotation), then the real annotation is still used. This stage allows reference frames to come from any high-confidence slice in the pseudo-label library, instead of being limited to the original real-labeled frames; FPM self-trains until the expected number of training iterations is reached.

[0024] Furthermore, the update rules for the pseudo-label library are designed as follows: During the FPM self-training process in this stage, given the FPM prediction of a certain unlabeled slice in one iteration of training... (a certain pixel) If the confidence level condition is met: If the pseudo-label is inconsistent with the current pseudo-label, then the pseudo-label of the corresponding pixel in the pseudo-label library will be used. Updated to; in, The higher threshold is used to determine the category with the highest predicted probability for this pixel. To ensure that only in non- The pseudo-tags are only overwritten when there is sufficient confidence, thus achieving a "gradual cleanup" of pseudo-tags. The above formula is the core update rule for the pseudo-tag library, which enables the quality of the pseudo-tag library to automatically improve as FPM capabilities improve.

[0025] Phase 3: Fine-tuning of high-confidence pseudo-labels for large 2D segmentation models, such as... Figure 1 As shown in Stage 3 of the model. In this stage, after FPM has fully converged, high-confidence samples are selected from the pseudo-label library to fine-tune SAM, making it adapt to the feature distribution of medical images and further improving the overall segmentation accuracy.

[0026] Specifically: First, the FPM trained in Phase 2 is used to perform bidirectional inference on the unlabeled slices again, resulting in a new set of pseudo-labels. The current content of the pseudo-tag library is recorded as follows: ; For any unlabeled slice, if and The Intersectionover Union (IoU) ratio on the target object exceeds the threshold. The pseudo-label is then considered a high-confidence sample and included in the fine-tuning set of SAM. in .

[0027] Next, the SAM is fine-tuned using box-prompts on selected high-confidence samples. Fine-tuning loss. The supervised loss method of Phase 1 will still be used. (i.e., cross-entropy and Dice loss), but box-prompt augmentation is no longer used in stage three to ensure that the prompts are as accurate as possible during fine-tuning.

[0028] This invention proposes an efficient method for constructing a large-scale 3D medical image segmentation model based on single-slice annotation. To verify its effectiveness, this method is compared with traditional and existing mainstream models (methods), including SAM, SemiSAM, MedSAM, SAM-Med2D, SAM2, U-Net, V-Net, MT, UA-MT, FixMatch, and DyCON. Segmentation tests are conducted under different levels of labeled data. The meaning of the labeled percentage is as follows: 0%: completely unlabeled; 0.96%, 1.14%, and 1.23%: Only one slice is labeled for each 3D medical image, accounting for approximately 1% of the total slices, which is considered extremely sparse annotation; 100%: Full annotation; Testing was conducted on three public datasets: LMS (leg muscles), LA (heart), and KiTS19 (kidney). Evaluation metrics included: Dice: segmentation accuracy metric, with higher values ​​indicating a closer relationship between the segmentation result and the ground truth annotation; HD95: segmentation boundary precision metric, with lower values ​​indicating smaller boundary errors.

[0029] As shown in Table 1, using only about 1% of the annotations, the method of this invention achieves Dice scores of 91.43%, 91.38%, and 93.82% on the three datasets, respectively. The HD95 score also significantly outperforms traditional and existing mainstream methods, with segmentation boundary accuracy approaching or even exceeding that of fully supervised models. This demonstrates that the proposed method possesses superior segmentation performance and practical application value in extremely sparse annotation scenarios, greatly reducing the cost of manual annotation.

[0030] Table 1. Comparison of segmentation performance of each model under different annotation ratios. Finally, it should be noted that the purpose of disclosing the embodiments is to help further understand the present invention. However, those skilled in the art will understand that various substitutions and modifications are possible without departing from the spirit and scope of the present invention and the appended claims. Therefore, the present invention should not be limited to the content disclosed in the embodiments, and the scope of protection of the present invention is defined by the scope of the claims.

Claims

1. A method for constructing a large-scale 3D medical image segmentation model based on single-slice annotation, characterized in that, A 2D segmentation model is used to segment a single slice with prompts. A Feature Propagation Model (FPM) is then used to propagate the segmentation features to other slices in the 3D medical image, collaboratively constructing a large 3D segmentation model. This model undergoes a three-stage progressive training process using multiple 3D medical image datasets. In each 3D medical image dataset, only one slice is labeled, with the rest being unlabeled. The three-stage progressive training specifically includes: Phase 1: Semi-supervised pre-training of the Feature Propagation Model (FPM); First, a large 2D segmentation model is used to perform cue segmentation on a single labeled slice to generate pseudo-labels. Then, the pseudo-labels are used as supervision signals to perform semi-supervised pre-training of FPM. During the training process, each training session takes three frames of slices as input, where the first frame is a labeled reference slice, and the second and third frames are unlabeled or labeled slices. FPM first predicts the result of the second frame based on the reference slice, and then combines the first two frames to predict the third frame. Backpropagation is performed through cross-entropy and Dice loss to optimize the model parameters. Phase Two: Construction of the Pseudo-Label Library and Self-Training of the Feature Propagation Model (FPM); After initial pre-training, the pre-trained FPM from Phase 1 is used to perform bidirectional inference on unlabeled slices, i.e., inference is performed from the labeled slices to both the front and back directions of the 3D image to obtain the prediction result for each unlabeled slice. For each unlabeled slice, the prediction result of the nearest reference slice is selected as its initial pseudo-label. Initial pseudo-labels are generated for all unlabeled slices in each 3D medical image, and a pseudo-label library is constructed and stored. Then, the pseudo-label library is used as a supervision signal to self-train the FPM until the expected number of iterations is reached. During training, labeled slices are used as reference frames or pseudo-labels of unlabeled slices are used as new reference frames. At the same time, the pseudo-label library is dynamically updated during self-training. Phase 3: Fine-tuning of high-confidence pseudo-labels for large 2D segmentation models; The FPM trained in Phase 2 is used to perform bidirectional inference on the unlabeled slices again to obtain a new set of pseudo-labels. This set is compared with the final pseudo-label library in Phase 2, and high-confidence samples are selected from it to fine-tune the 2D segmentation model to adapt it to the feature distribution of medical images.

2. The method as described in claim 1, characterized in that, The three-frame slices used as input to the Feature Propagation Model (FPM) employ a three-frame triplet sampling rule, with each training iteration using three-frame triplets. Enter FPM, where This represents the number of labeled frames in the last two frames. This represents the number of unlabeled frames in the last two frames, with the following constraints: (1) The reference frame is always the annotation frame; (2) During training, there is a 50% probability that all frames are unlabeled and a 50% probability that one frame is labeled and one frame is unlabeled. (3) Skip sampling is used for unlabeled frames, and the skip step size is gradually increased linearly or segmentally as training progresses.

3. The method as described in claim 1, characterized in that, In the training of the Feature Propagation Model (FPM) for Phase 1 and Phase 2, FPM performs backpropagation using cross-entropy and Dice loss. The loss function is calculated as follows: Order No. The prediction results for labeled / unlabeled frames are: Marked as The pseudo-tag is The overall training loss of FPM is defined as: Among them, for frames with real annotations, the supervised loss for: In the formula This represents pixel-level cross-entropy loss. Indicates Dice loss: Using pseudo-labels to supervise unlabeled frames, the unsupervised loss is defined as: in This is an indicator function; it returns 1 if a condition is met, and 0 otherwise. (Threshold) Used to remove predicted pixels with excessively low confidence; pseudo-labels In Phase 1, the loss is generated by a large 2D segmentation model, and in Phase 2, it is generated by the FPM prediction itself. The labels for the supervised loss in Phase 2 come from a pseudo-label library or real labels.

4. The method as described in claim 1, characterized in that, In the first stage, before generating pseudo-labels using the large 2D segmentation model, the large 2D segmentation model is fine-tuned. A bounding box selection enhancement strategy is introduced to randomly perturb the ground truth bounding boxes of the reference slices. During fine-tuning, labeled slices are used, and the original bounding boxes and several perturbed bounding boxes are used to provide hints to the large 2D segmentation model to generate prediction results. The loss function used is supervised loss.

5. The method as described in claim 1, characterized in that, During the self-training process of the Phase 2 Feature Propagation Model (FPM), the pseudo-label library is dynamically updated. The update rule is: when the prediction confidence of the Feature Propagation Model (FPM) for a certain pixel exceeds the threshold and is inconsistent with the current pseudo-label, the pseudo-label of that pixel is updated.

6. The method as described in claim 1, characterized in that, The 2D segmentation large model fine-tuning in stage three uses the cross-entropy and Dice loss employed during the training of the feature propagation model FPM in stage one.

Images