Three-dimensional tumor image semi-supervised segmentation method based on uncertainty guidance

Abstract

The invention provides a three-dimensional tumor image semi-supervised segmentation method based on uncertainty guidance. The method comprises the following steps: preprocessing an unlabeled three-dimensional tumor image, and inputting the unlabeled three-dimensional tumor image into a three-dimensional medical image segmentation basic model in combination with prompt information; a plurality of tumor segmentation masks are obtained through enhancement during testing, case-level uncertainty scores are calculated according to the tumor segmentation masks, and a voxel-level uncertainty thermodynamic diagram is generated; dividing the samples into a high-uncertainty sample set and a low-uncertainty sample set according to the case-level uncertainty score, performing manual labeling on the high-uncertainty samples, and taking a prediction result of the basic model as a pseudo label for the low-uncertainty samples; and combining the two types of samples to construct a mixed training set, and training a target model by adopting a weighted sampling strategy. According to the method, the samples can be effectively screened based on uncertainty guidance, the general segmentation capability of the basic model is fully utilized while the marking cost is remarkably reduced, and the precision of the model on a specific tumor segmentation task is improved.

Description

Technical Field

[0001] This invention belongs to the field of medical image segmentation technology, specifically relating to a semi-supervised segmentation method for three-dimensional tumor images based on uncertainty guidance. Background Technology

[0002] Medical imaging plays a crucial role in cancer diagnosis, treatment planning, and prognostic assessment. Accurate and reliable tumor segmentation not only helps clinicians precisely identify tumor boundaries and assess tumor volume but also significantly improves patient outcomes and drives the development of precision medicine. Despite continuous advancements in imaging technologies such as computed tomography (CT) and magnetic resonance imaging (MRI), tumors exhibit high heterogeneity in shape, size, and location, and often have similar densities and indistinct boundaries with surrounding soft tissues (such as blood vessels, muscles, or adjacent organs), posing significant challenges to accurate tumor segmentation.

[0003] In recent years, deep learning technology has made significant progress in the field of medical image segmentation. The introduction of uncertainty quantification has provided a new approach to solving these problems. Uncertainty reflects the model's confidence in its own predictions, characterizing the ambiguity of whether the model will output a correct segmentation given a given input. By evaluating the uncertainty of the segmentation results, the generated uncertainty heatmap can reveal regions with significant segmentation ambiguity, thus providing a more reliable reference for clinical diagnosis and treatment planning, and enhancing the scientific rigor and safety of decision-making.

[0004] Currently, most advanced automated tumor segmentation models (such as nnU-Net), while achieving high segmentation accuracy on publicly available datasets, still suffer from the following limitations: First, these models typically only output deterministic segmentation results, lacking quantification of prediction confidence and making it difficult to identify and highlight potentially mis-segmented regions. Second, model performance is highly dependent on large amounts of precise pixel-level labeled data, which requires experienced radiologists to invest significant time in manual delineation, resulting in high costs and low efficiency. Furthermore, existing models lack efficient interaction and correction mechanisms; physicians must either accept potentially erroneous results entirely or undergo complete manual modification, limiting the widespread adoption of such models in clinical practice.

[0005] To overcome data bottlenecks and interaction challenges, foundational medical image segmentation models (such as SegVol) have emerged. These models are pre-trained on massive amounts of medical data, learning common anatomical features and possessing powerful feature extraction and generalization capabilities. More importantly, they support multimodal interaction methods such as textual and spatial cues, effectively reducing semantic ambiguity in segmentation tasks such as "liver" and "liver tumor," thus better meeting actual clinical needs. Although these foundational models exhibit excellent general performance, their direct inference accuracy may be insufficient to meet clinical requirements when dealing with certain types of tumors (such as soft tissue tumors with blurred boundaries and low contrast), often necessitating further transfer training targeting the specific lesion.

[0006] Therefore, in the absence of large-scale labeled data, how to efficiently utilize the general capabilities of basic medical image segmentation models, achieve high-precision tumor segmentation through low-cost labeling strategies, and at the same time provide reliable uncertainty metrics to assist clinical decision-making is a technical challenge that urgently needs to be solved in the field of medical image segmentation. Summary of the Invention

[0007] In view of this, the present invention provides a semi-supervised segmentation method for three-dimensional tumor images based on uncertainty guidance, which can effectively screen samples based on uncertainty guidance, significantly reduce annotation costs, and make full use of the general segmentation capabilities of the basic model to improve the accuracy of the model on specific tumor segmentation tasks.

[0008] To achieve the above objectives, the present invention adopts the following technical solution:

[0009] A semi-supervised segmentation method for three-dimensional tumor images based on uncertainty guidance, comprising:

[0010] Step S1: Preprocess the unlabeled 3D tumor images and input them into the 3D medical image segmentation basic model in conjunction with the prompt information;

[0011] Step S2: Use time-of-test enhancement to obtain multiple tumor segmentation masks, calculate case-level uncertainty scores based on them, and generate voxel-level uncertainty heatmaps.

[0012] Step S3: Divide the samples into high-uncertainty and low-uncertainty sample sets according to the case-level uncertainty score. Manually label the high-uncertainty samples and use the prediction results of the basic model as pseudo-labels for the low-uncertainty samples.

[0013] Step S4: Merge the two types of samples to construct a mixed training set, and use a weighted sampling strategy to train the target model.

[0014] Furthermore, in step S1, the preprocessing includes: resampling, standardization, and normalization of the tumor images; the prompting information includes text prompts and spatial prompts, where the text prompts are text descriptions of the tumor category, and the spatial prompts can be either points or boxes.

[0015] Further, in step S1, the basic model for three-dimensional medical image segmentation adopts an encoder-decoder architecture, including: an image encoder, a spatial encoder, a text encoder, and a mask decoder; the image encoder is used to extract global and local features of the three-dimensional medical image; the spatial encoder is used to receive and encode spatial cue information, which includes points and boxes; the text encoder is used to encode the text description or feature vector of the tumor category; and the mask decoder is used to fuse image features, spatial encoded features, and text encoded features to output the final segmentation mask.

[0016] Furthermore, in step S2, the test-time enhancement includes one or more combinations of geometric transformation, intensity transformation, noise injection, and blurring.

[0017] Further, in step S2, the calculation process of the case-level uncertainty score specifically includes: for an unlabeled tumor image sample X, obtaining N tumor segmentation masks {P1, P2, ..., P...} based on test-time enhancement. N}, calculate all possible combinations (P) i P j The average pairwise Dice coefficient Then the case-level uncertainty score U of this sample case (X) = 1 - MPD(X).

[0018] Further, in step S2, the voxel-level uncertainty heatmap visualizes the uncertainty score of each voxel in the three-dimensional tumor image using Jet color mapping. For the voxel at spatial location v, the enhanced average prediction probability is calculated after N tests. Among them, P n (v) represents the predicted probability value of voxel v under the augmentation policy at the nth test, then the voxel-level uncertainty score

[0019] Further, in step S3, the process of dividing the high and low uncertainty samples is as follows: all samples are sorted from high to low according to the case-level uncertainty score, and divided into high uncertainty samples and low uncertainty samples according to a pre-set fixed threshold; the fixed threshold can be set to any value among 5%, 10%, 15% or 20%; the high uncertainty samples are manually labeled to obtain manually labeled samples, and the prediction results of the basic model are used as pseudo-labels to obtain pseudo-labeled samples for the low uncertainty samples.

[0020] Further, in step S4, the weighted sampling strategy assigns sampling weights based on the case-level uncertainty score of the samples. Specifically, this includes: assigning a fixed weight α (α≥1) to manually labeled samples; calculating the weight of pseudo-labeled samples based on their case-level uncertainty score u using an exponential decay function w=exp(-u / τ), where τ is a temperature coefficient adjusting the decay rate; and calculating the sum W of the sampling weights for all samples. sum =∑w, which gives the normalized sampling probability of each sample. Based on the sampling probability, samples are randomly selected to form training batches in each iteration step of model training.

[0021] Furthermore, in step S4, the target model is a deep neural network based on a three-dimensional U-Net structure.

[0022] The present invention also provides a computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements the steps of the method described in any of the above claims.

[0023] The present invention has the following beneficial effects:

[0024] 1. By using uncertainty-driven sample screening, manual annotation is performed only on cases with high uncertainty, which significantly reduces annotation costs and alleviates the dependence on large-scale labeled data.

[0025] 2. Make full use of the general segmentation capabilities of the pre-trained basic model and improve the accuracy and generalization performance of the model on specific tumor tasks through semi-supervised transfer learning.

[0026] 3. Introducing enhanced quantitative model prediction uncertainty during testing provides clinicians with intuitive visual aids to help identify potential segmentation ambiguities, thereby supporting more reliable clinical diagnosis and treatment decisions.

[0027] 4. A weighted sampling training strategy was designed to balance the contributions of manually labeled data and pseudo-labeled data in the model training process, effectively suppressing the adverse effects of noise and pseudo-labels on model performance and improving model robustness and final segmentation accuracy. Attached Figure Description

[0028] To more clearly illustrate the technical solutions in the embodiments of the present invention or the prior art, the drawings used in the description of the embodiments or the prior art will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0029] Figure 1This is a flowchart of a semi-supervised segmentation method for three-dimensional tumor images based on uncertainty guidance provided in an embodiment of the present invention;

[0030] Figure 2 This is a schematic diagram of a semi-supervised segmentation method for three-dimensional tumor images based on uncertainty guidance provided in an embodiment of the present invention;

[0031] Figure 3 This is a schematic diagram of the uncertainty quantification module provided in an embodiment of the present invention. Detailed Implementation

[0032] The technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. The described embodiments are only some embodiments of the present invention, and not all embodiments. All other embodiments obtained by those skilled in the art based on the embodiments of the present invention without creative effort are within the scope of protection of the present invention.

[0033] Example 1

[0034] like Figure 1 As shown, this embodiment provides a basic implementation of the uncertainty-guided semi-supervised segmentation method for 3D tumor images, including the following steps: preprocessing unlabeled 3D tumor images and inputting them into a 3D medical image segmentation base model in conjunction with prompting information; obtaining multiple tumor segmentation masks using test-time enhancement, calculating case-level uncertainty scores, and generating voxel-level uncertainty heatmaps; dividing the samples into high-uncertainty and low-uncertainty sample sets according to the case-level uncertainty scores, manually labeling high-uncertainty samples, and using the prediction results of the base model as pseudo-labels for low-uncertainty samples; merging the two types of samples to construct a hybrid training set, and training the target model using a weighted sampling strategy.

[0035] Example 2

[0036] like Figure 2 As shown in the figure, this invention discloses a semi-supervised segmentation method for three-dimensional tumor images based on uncertainty guidance. The specific implementation steps are as follows:

[0037] In this embodiment, pelvic CT images of cervical cancer patients are selected as unlabeled three-dimensional tumor images to be processed. In order to eliminate the differences caused by acquisition by different devices, the CT images are preprocessed, specifically including: resampling the image data to the median voxel spacing of all data; extracting the foreground with the global mean as the threshold, removing extreme outliers outside the range of 0.05% to 99.95%, and then performing Z-score standardization; and using the Min-Max normalization method to map the data values ​​to the range of [0, 1].

[0038] Preferably, in this embodiment, the SegVol model is selected as the basic model for three-dimensional medical image segmentation, combined with spatial and textual prompts. The spatial prompt is a three-dimensional bounding box, and the text prompt is "cervical cancer".

[0039] Using test-time enhancement techniques, case-level uncertainty scores are calculated and voxel-level uncertainty heatmaps are generated. Preferably, the test-time enhancement includes five basic enhancement methods: flipping, intensity scaling, intensity shifting, Gaussian noise, and Gaussian smoothing. These are combined in pairs to generate N (preferably 10 in this embodiment) fixed test-time enhancement strategies. For each unlabeled image X, these N test-time enhancement strategies are applied and input into the SegVol model for inference. The results are inversely transformed back to the original space to obtain the corresponding prediction mask set {P1, P2, ..., P...}. N}

[0040] To evaluate the model's case-level uncertainty for the current sample, firstly, the pairwise Dice coefficients between the N prediction masks are calculated. For any two prediction masks P... i and P j Its Dice coefficient is defined as:

[0041]

[0042] Among them, |P i ∩P j | represents the total number of voxels in the overlapping region between the i-th and j-th prediction masks; |P i |+|P j | represents the sum of the total number of voxels contained in the i-th prediction mask and the j-th prediction mask.

[0043] Next, calculate the average pairwise Dice coefficient (MPD) for all possible combinations:

[0044]

[0045] Finally, the case-level uncertainty score U for this sample is defined. case (X) is:

[0046] U case (X)=1-MPD(X)

[0047] U case The higher the (X) value, the greater the difference between the N prediction results, and the lower the confidence of the model in the segmentation result of the sample. This usually corresponds to samples with blurred tumor boundaries or complex morphology.

[0048] For a voxel at spatial location v, its voxel-level uncertainty U voxel (v) is:

[0049]

[0050] Among them, P n (v) represents the predicted probability value of voxel v under the enhancement strategy at the nth test. The result is processed by Min-Max normalization and mapped to the [0, 1] interval. Using Jet color mapping, a voxel-level uncertainty heatmap is generated. The redder the color, the higher the uncertainty of the voxel, and the bluer the color, the lower the uncertainty of the voxel.

[0051] All unlabeled samples D unlabeled According to the case-level uncertainty score U case (X) Sorting samples from highest to lowest uncertainty, and dividing them into high-uncertainty samples and low-uncertainty samples based on a pre-set fixed threshold; wherein, the fixed threshold can be set to any value among 5%, 10%, 15%, or 20%. In this embodiment, the fixed threshold is preferably 10%, that is, selecting the top 10% of samples by uncertainty score as high-uncertainty samples, and the remaining 90% as low-uncertainty samples. The high-uncertainty samples are manually labeled to form a real-world labeled dataset D. labeled For samples with low uncertainty, the original prediction results enhanced by the SegVol model without testing are used as pseudo-labels to form the pseudo-label dataset Dpseudo. labeled With the D labeled Merge and construct a hybrid training set.

[0052] The target model is trained based on the hybrid training set, and the target model is preferably nnU-Net.

[0053] To mitigate the noise impact caused by the fact that the amount of pseudo-labeled data is much larger than the amount of real-labeled data, this embodiment uses a weighted sampling strategy. A fixed weight α (preferably 1.5 in this embodiment) is assigned to manually labeled samples; the weight of pseudo-labeled samples is calculated based on their case-level uncertainty score u using an exponential decay function w = exp(-u / τ), where τ is a temperature coefficient adjusting the decay rate; and the sum W of the sampling weights of all samples is calculated. sum =∑w, which gives the normalized sampling probability of each sample. Based on the sampling probability, samples are randomly selected to form training batches in each iteration step of model training.

[0054] Preferably, this embodiment employs a combined loss function during training, combining Dice loss and cross-entropy loss. Wherein, the Dice loss L... Dice for:

[0055]

[0056] Where yi represents the true label value, which can be 0 or 1, pi represents the probability value predicted by the model, N represents the total number of primes, and ε is the smoothing term.

[0057] Cross-entropy loss L CE for:

[0058]

[0059] The total loss function is: L Total =L Dice +L CE .

[0060] Preferably, this embodiment uses an SGD optimizer and a multinomial decaying learning rate scheduler, with an initial learning rate of 1e-2, a weight decay of 1e-4, 1000 training epochs, and a batch size of 2.

[0061] The above description is merely a preferred embodiment of the present invention, but the scope of protection of the present invention is not limited thereto. Any equivalent substitutions or modifications made by those skilled in the art within the scope of the technology disclosed in the present invention, based on the technical solution and inventive concept of the present invention, should be covered within the scope of protection of the present invention.

Claims

1. A semi-supervised segmentation method for three-dimensional tumor images based on uncertainty guidance, characterized in that, include: Step S1: Preprocess the unlabeled 3D tumor images and input them into the 3D medical image segmentation basic model in conjunction with the prompt information; Step S2: Use time-of-test enhancement to obtain multiple tumor segmentation masks, calculate case-level uncertainty scores based on them, and generate voxel-level uncertainty heatmaps. Step S3: Divide the samples into high-uncertainty and low-uncertainty sample sets according to the case-level uncertainty score. Manually label the high-uncertainty samples and use the prediction results of the basic model as pseudo-labels for the low-uncertainty samples. Step S4: Merge the two types of samples to construct a mixed training set, and use a weighted sampling strategy to train the target model.

2. The uncertainty-guided semi-supervised segmentation method for three-dimensional tumor images as described in claim 1, characterized in that, In step S1, the preprocessing includes: resampling, standardization, and normalization of the tumor images; the prompting information includes text prompts and spatial prompts; the text prompts are text descriptions of the tumor category, and the spatial prompts can be either points or boxes.

3. The uncertainty-guided semi-supervised segmentation method for three-dimensional tumor images as described in claim 1, characterized in that, In step S1, the basic model for 3D medical image segmentation adopts an encoder-decoder architecture, including: an image encoder, a spatial encoder, a text encoder, and a mask decoder; the image encoder is used to extract global and local features of 3D medical images; the spatial encoder is used to receive and encode spatial cue information, which includes points and boxes; the text encoder is used to encode the text description or feature vector of the tumor category; and the mask decoder is used to fuse image features, spatial encoded features, and text encoded features to output the final segmentation mask.

4. The uncertainty-guided semi-supervised segmentation method for three-dimensional tumor images as described in claim 1, characterized in that, In step S2, the test-time enhancement includes one or more combinations of geometric transformation, intensity transformation, noise injection, and blurring.

5. The uncertainty-guided semi-supervised segmentation method for three-dimensional tumor images as described in claim 1, characterized in that, In step S2, the calculation process of the case-level uncertainty score specifically includes: for an unlabeled tumor image sample X, obtaining N tumor segmentation masks {P1, P2, ..., P...} based on test-time enhancement. N }, calculate all possible combinations (P) i P j The average pairwise Dice coefficient Then the case-level uncertainty score U of this sample case (X) = 1 - MPD(X).

6. The uncertainty-guided semi-supervised segmentation method for three-dimensional tumor images as described in claim 1, characterized in that, In step S2, the voxel-level uncertainty heatmap visualizes the uncertainty score of each voxel in the 3D tumor image using Jet color mapping; for the voxel at spatial location v, the enhanced average prediction probability is calculated over N tests. Among them, P n (v) represents the predicted probability value of voxel v under the augmentation policy at the nth test, then the voxel-level uncertainty score 7. The uncertainty-guided semi-supervised segmentation method for three-dimensional tumor images as described in claim 1, characterized in that, In step S3, the process of dividing the high and low uncertainty samples is as follows: all samples are sorted from high to low according to the case-level uncertainty score, and divided into high uncertainty samples and low uncertainty samples according to a pre-set fixed threshold; the fixed threshold can be set to any value among 5%, 10%, 15% or 20%; the high uncertainty samples are manually labeled to obtain manually labeled samples, and the prediction results of the basic model are used as pseudo-labels to obtain pseudo-label samples for the low uncertainty samples.

8. The uncertainty-guided semi-supervised segmentation method for three-dimensional tumor images as described in claims 1 and 7, characterized in that, In step S4, the weighted sampling strategy assigns sampling weights based on the case-level uncertainty score of the samples. Specifically, this includes: assigning a fixed weight α (α≥1) to manually labeled samples; calculating the weight of pseudo-labeled samples based on their case-level uncertainty score u using an exponential decay function w = exp(-u / τ), where τ is a temperature coefficient adjusting the decay rate; and calculating the sum W of the sampling weights for all samples. sum =∑w, which gives the normalized sampling probability of each sample. Based on the sampling probability, samples are randomly selected to form training batches in each iteration step of model training.

9. The uncertainty-guided semi-supervised segmentation method for three-dimensional tumor images as described in claim 1, characterized in that, In step S4, the target model is a deep neural network based on a three-dimensional U-Net structure.

10. A computer-readable storage medium having a computer program stored thereon, characterized in that, When the computer program is executed by a processor, it is able to implement the method steps as described in any one of claims 1 to 9.

Images