Weak supervision three-dimensional medical image segmentation method based on MedSAM

Abstract

The invention discloses a weak supervision medical image segmentation method based on MedSAM, and aims to solve the problems that traditional medical image segmentation depends on a large amount of annotation data, and training is time-consuming and labor-consuming, and improve the segmentation performance under the condition that the annotation cost is limited. A deep learning model feature extraction capability is utilized, and a category activation graph is generated through a backbone network to preliminarily position a target area. A CSD bidirectional common sense defense module and an AMS class activation graph superposition module are provided, the former is combined with global information to improve the robustness of the model, and the latter is used for determining the maximum target range to guarantee the segmentation effect. A dynamic threshold segmentation strategy and a MedSAM model are combined to enhance a thermodynamic diagram, semantic features are integrated to avoid dependence on prototype learning, and weak supervision high-precision segmentation is realized. Experiments show that the multi-data-set performance of the method is superior to that of the existing method, the application prospect is wide, and the method can be applied to various medical image tasks and other fields which are high in labeling cost and need efficient segmentation.

Description

Technical Field

[0001] The invention belongs to the technical field of medical image segmentation, and in particular relates to a weakly supervised three-dimensional medical image segmentation method based on MedSAM. Background Art

[0002] Most existing weakly supervised semantic segmentation methods train classification networks and extract localization maps from convolutional neural networks based on class activation maps. However, the initial CAM is often incomplete or redundant. To address this challenge, researchers have proposed a variety of improvement strategies at different stages to improve the quality of CAM and the final segmentation effect.

[0003] In order to solve the incompleteness problem of CAM, some studies have trained the classification network through auxiliary tasks and designed additional loss functions to guide the model to discover more target areas. The "erasing" strategy is one of the common methods, which forces the network to explore more undiscovered areas by erasing the salient parts of the image or feature map. In addition, there are studies that accumulate multiple activation results during the training process to improve model performance; other methods have proposed innovative improvement solutions from the perspectives of cross-image mining, self-supervision mechanisms, and anti-adversarial attacks. To alleviate the redundancy problem of CAM, previous studies have used softmax cross entropy as an additional loss to reactivate the model, or introduced out-of-distribution data to improve generalization ability. Recently, some methods based on the Transformer architecture have been introduced into the WSSS task and have shown competitive performance.

[0004] In order to optimize the initial CAM, improving the quality of the localization map by learning semantic affinity has become one of the mainstream methods. For example, PSA trains the network to learn the semantic affinity between pixels and propagates the semantics of the strong response area in the attention map to pixels with similar semantics. IRNet and BES significantly improve the segmentation effect by synthesizing category boundaries and expanding the coverage of the target area to the boundary. In addition, some studies further optimize CAM by combining saliency maps to obtain more accurate background information or distinguish co-occurring target objects.

[0005] In the stage of training segmentation models, existing methods usually generate pseudo labels by applying a global threshold to CAM, but ignore the confidence information of each pixel and fail to fully tap the potential of CAM. To address this shortcoming, the PMM model proposes a disguised underfitting strategy to improve segmentation performance by reweighting the loss of potential noise pixels; URN estimates uncertainty by scaling the prediction map multiple times and combining conditional random fields. However, the improvement of PMM is limited to the loss level and does not directly use confidence information, while URN has low computational efficiency due to the need for multiple CRF processing.

[0006] On the other hand, in recent years, with the rapid development of artificial intelligence technology, the application of large-scale basic models in the field of weakly supervised semantic segmentation has received increasing attention, and various segmentation tasks and methods using large-scale basic models have sprung up. In this development process, researchers have focused on two key aspects. On the one hand, they are committed to the architectural design of the basic model, striving to create a model framework with a more reasonable structure and better performance, so that it can better adapt to different segmentation task requirements; on the other hand, they focus on developing flexible and user-friendly prompts to facilitate users to interact with the model more conveniently and efficiently, and guide the model to output segmentation results that meet expectations.

[0007] Many scholars are also actively exploring how to use existing visual basic models to assist weakly supervised semantic segmentation. For example, Liang et al., Qin et al., and Ghiasi et al. cleverly used powerful visual basic models such as CLIP or DINOv2, and trained them (usually by fine-tuning) to work with specific segmentation heads or adapters. After such processing, these basic models can be better integrated into the semantic segmentation task, dig out more potential segmentation clues, and thus improve the accuracy and efficiency of segmentation.

[0008] Particularly worth mentioning is the arbitrary segmentation model, which has demonstrated remarkable category-independent segmentation capabilities as soon as it came out, and can perform segmentation more accurately regardless of the type of target object. This outstanding performance has undoubtedly achieved great success, causing strong repercussions in academia and industry, and has inspired many researchers to actively think about and try to design innovative methods that can achieve effective segmentation without a lot of training. On this basis, the introduction of MedSAM further expanded the application boundaries of SAM and focused on medical image segmentation tasks. MedSAM is optimized for the characteristics of medical images and can achieve efficient and accurate target area segmentation on complex medical data such as MRI and CT. This improvement not only reduces the dependence on large-scale annotated data, but also provides strong support for clinical auxiliary diagnosis and medical image analysis, becoming a highlight in the field of medical image segmentation. The present invention is based on the basic model of medical image segmentation MedSAM, and realizes a weakly supervised three-dimensional medical image segmentation method.

[0009] The purpose of 3D medical image segmentation is to segment 3D medical images such as CT and MRI into dense predicted pixels corresponding to lesions or organs. Volumetric medical image segmentation has been explored from different perspectives, including U-net, 3D convolutional networks, Transformers, diffusion models, domain adaptation, etc. Its core goal is to perform detailed segmentation of 3D medical images such as CT and MRI, and finally form dense predicted pixels corresponding to lesions or organs, providing powerful image analysis support for applications such as medical diagnosis. The success of large-scale basic models has brought new perspectives and challenges to this field. Summary of the invention

[0010] The present invention belongs to the technical field of medical image segmentation, and specifically relates to a weakly supervised three-dimensional medical image segmentation method based on MedSAM. The method is based on deep learning, uses a backbone network to extract image features and generate a CAM class activation heat map, combines CAM feature information, and further enhances the segmentation effect through MedSAM. After inputting the class activation heat map, different from the traditional weakly supervised segmentation method, the present invention innovatively introduces a segmentation framework that does not require training labels or rely on prototype learning, and designs a CSD bidirectional common sense defense module and an AMS class activation map overlay module. Through this new weakly supervised segmentation framework, the present invention significantly improves the segmentation quality when only image-level labels are input.

[0011] The present invention provides a weakly supervised three-dimensional medical image segmentation method based on MedSAM, the main steps of the method are as follows: Step 1: Obtain n target MRI images. We slice the 3D image data along a specific axis to generate 2D slices, which are divided into two categories: target area and non-target area.

[0012] Step 2: Perform classification network training to obtain a CAM class activation heat map, or use other methods to provide a CAM class activation heat map; perform dynamic threshold segmentation on the class activation heat map and stack it to obtain a superimposed segmentation mask.

[0013] Step 3: Use the superimposed segmentation mask and class activation heat map as input to the AMS module to obtain the preliminary segmentation result.

[0014] Step 4: The segmentation result obtained in step 3, the class activation heat map, and the superimposed segmentation mask are used as the input of the CSD module to further enhance the segmentation effect.

[0015] Step 5: Use the segmentation result obtained in step 4 as the input of the MedSAM-based segmentation Module to obtain the segmentation result enhanced by the large model.

[0016] Step 6: superimpose the segmentation results obtained in step 5 according to the operation of step 2, replacing the superimposed segmentation mask in step 2, and repeat steps 3, 4, and 5 to obtain the final segmentation result.

[0017] Preferably, the classification model described in step S2 is specifically operated as follows: Step S2.1: For each input image, use the data augmentation method to preprocess the image by cropping and flipping; Step S2.2: Train the classification network with a learning rate of 0.001 for 10 epochs; crop and flip the image using data augmentation methods; extract the class activation heat map M for the target area in each MRI image * ; Set a hyperparameter τ, set the position greater than the product of the maximum value of the heat map and τ to 1, otherwise set it to 0, and get the initial segmentation mask for each frame

[0018] Preferably, in step S3, the initial segmentation mask is superimposed, normalized to the maximum and minimum, and segmented by dynamic threshold along the axial direction specified in S1 to obtain a superimposed segmentation mask. The specific steps are as follows: Step S3.1: The superposition operation is: Step S3.2: Maximum and minimum normalization is performed to obtain the superimposed heat map S, and the superimposed segmentation mask is obtained by dynamic threshold segmentation. A slice of 3D image data can be regarded as a frame of video data. Compared with video data, it also has continuous visual feature transformation. Then, the high response area is segmented by dynamic thresholding: Step S3.3: Reduce the random noise of individual slices by taking the intersection and strengthen the consistency signal between slices. As the initial segmentation mask after preliminary quality enhancement; in

[0019] As a preferred embodiment, the two-way common sense defense module has the following specific steps: Step S4.1: Determine whether the current slice has poor class activation performance, using bidirectional extreme judgment conditions: represents the highest response area of ​​the class activation heat map, Indicates the highest corresponding area of ​​the stacked graph; Represents the dynamic threshold segmentation area of ​​the class activation map, Represents the dynamic threshold segmentation region of the stacked image; Step S4.2: If the current slice class activation performs well and does not meet the bidirectional extreme judgment condition, then no change is made; If the performance is poor and meets the bidirectional extreme judgment condition, then enhance them by utilizing the continuous visual feature changes in temporally adjacent frames; The previous frame or the next frame may provide better segmentation; bidirectional optimization is used to make changes using adjacent frames, including reverse optimization: And positive direction optimization: Where X t represents the original input image at frame t, Represents an operation that extracts the average brightness value of the image within the mask, or the average brightness value of the entire image if no input mask is present.

[0020] Preferably, in step S5, a mask bounding box extractor is used to extract For the bounding box B t As the input of MedSAM, we get P t is the segmentation mask: P t =MedSAM(X t ,B t ) Assume the segmentation mask P t Contains multiple connected regions C t ; If for all connected regions C t，k (k∈i,…,K), satisfying the following conditions L(X t )>L(X t ,C t，k ), that is, the brightness of the connected area is less than the average brightness of the entire image, it is considered to be an error, and the target area may not be included in any connected area; in this case, we replace the original segmentation mask with the frame that meets the following conditions by comparing the brightness of the t-1 frame and the t+1 frame: Among them, L(X t ) represents the overall average brightness of the image, L(X t ,C t，k ) indicates that the image is in C t，k The average brightness of the area.

[0021] Compared with the prior art, the present invention has the following beneficial effects: First, we propose a prototype-free weakly supervised medical image segmentation framework, which frees users from training or providing prototypes.

[0022] Second, to reduce the uncertainty of the initial seeds in weakly supervised image segmentation, we adopt an activation map superposition (AMS) strategy to ensure cross-frame consistency and thus improve reliability.

[0023] Finally, to mitigate unexpected failures generated by large base models, a bidirectional common sense defense (CSD) mechanism based on extreme value checking is proposed to prevent the segmentation model from generating unreasonable outputs. BRIEF DESCRIPTION OF THE DRAWINGS

[0024] Figure 1 It is a flow chart of the model structure of the present invention.

[0025] Figure 2 It is a schematic diagram of the extreme inspection conditions in the bidirectional common sense defense mechanism of the present invention. DETAILED DESCRIPTION

[0026] refer to Figure 1 , Figure 1 The model structure flow chart of the present invention includes: S1. Obtain n target MRI images. We slice the 3D image data along a specific axis to generate 2D slices, which are divided into two categories: target area and non-target area.

[0027] S2, use the Vi'TB-16 model pre-trained based on ImageNet-21k for classification training; input the 2D slices of the target area into S2 to obtain the trained model, and use Grad-CAM and Class-aware Attention-based Affinity modules to generate the class activation heat map of the target area; perform dynamic threshold segmentation on the obtained class activation heat map, and obtain the high response area as the initial segmentation mask.

[0028] S3. We define a class activation map superposition module, which can simply and effectively integrate the global semantic feature relationship for all initial segmentation masks from the same patient, superimpose the initial segmentation masks along the axis specified in S1, perform maximum and minimum normalization, and dynamic threshold segmentation to obtain a superimposed segmentation mask. The activation information from a single slice is usually not enough to cover the entire target area, and combined with global information, a global representation that enhances spatial consistency and reduces uncertainty can be formed. This can continuously activate the target area throughout the 3D space, help reduce the impact of random noise in a single slice, and strengthen consistent signals between slices to highlight any potential target areas. The initial segmentation mask and the obtained superimposed segmentation mask are intersected to reduce the information that is incorrectly activated in the initial segmentation mask.

[0029] S4,We define a bidirectional common sense defense module, which can perform extreme checks in both directions, such as Figure 2 As shown in . If the extreme conditions are met, the class activation map of the frame is considered to be poorly performed; otherwise, it is considered to be well-performed. For the poorly performing slices, its initial segmentation mask is replaced with the initial segmentation mask of the previous or next frame.

[0030] S5. For each frame, the initial segmentation mask is extracted by the mask bounding box extractor to obtain the bounding box of the target area. The bounding box and the original image are used as the input of the MedSAM model to obtain the enhanced segmentation mask. At the same time, a post-processing strategy based on average brightness is designed to reduce the error of the segmentation result.

[0031] S6. The enhanced segmentation mask obtained in S5 is superimposed according to the operation in S3 and the superimposed segmentation mask is replaced. The initial segmentation mask and the newly obtained superimposed segmentation mask are intersected, and the final segmentation mask is obtained by going through steps S4 and S5 again.

[0032] The specific operation of the classification model in step S2 is as follows: For each input image, Vi'TB-16 is used for training with a small learning rate for 10 epochs. The image is preprocessed by cropping and flipping using data augmentation methods. The class activation heat map is extracted for the target area in each MRI image using Grad-CAM and CAA modules.

[0033] First, we extract feature embeddings from the last layer of ViT, where N is the number of patches and d is the embedding dimension. During backpropagation, we compute the output y for category c c The gradient is: Next, these gradients are globally averaged pooled to obtain patch weights Then, we calculate each embedding F i The patch and its corresponding weight The weighted sum of c .

[0034] Finally, the previously generated class activation maps are enhanced by utilizing the attention weights from the last multi-head self-attention layer: M * =B c ⊙A t ·vec(M c ) Where ⊙ is the Hadamard product, t is the number of iterations, vec(·) is the vectorized operation, and A is the symmetric affinity matrix derived from the attention weights. The specific method of dynamic threshold segmentation is: set a hyperparameter τ, set the position greater than the product of the maximum value of the heat map and τ to 1, and vice versa, to 0, and obtain the initial segmentation mask for each frame

[0035] In step S3, the initial segmentation mask is superimposed, normalized to the maximum and minimum, and segmented by dynamic threshold along the axis specified in S1 to obtain a superimposed segmentation mask. The superimposition operation is: After maximum and minimum normalization, the superimposed heat map S is obtained, and the superimposed segmentation mask is obtained by dynamic threshold segmentation. A slice of 3D image data can be regarded as a frame of video data. Compared with video data, it also has continuous visual feature transformation. The sequence connection between continuous slices can be used to solve many challenges in single 2D image processing. Then, the high response area is segmented by dynamic thresholding: Then, the random noise of individual slices is reduced by taking the intersection, and the consistency signal between slices is strengthened. As the segmentation mask after preliminary quality enhancement.

[0036] In step S4, the two-way common sense defense module has the following specific steps: Step 1: Determine whether the current slice has poor class activation performance, using bidirectional extreme judgment conditions: Indicates the highest response area of ​​the class activation heat map; Indicates the highest corresponding area of ​​the stacked graph; Represents the dynamic threshold segmentation area of ​​the class activation map; Represents the dynamic threshold segmentation area of ​​the stacked image; Step 2: If the current slice class activation performs well, that is, it does not meet the two-way extreme judgment condition, then no change is made. If it performs poorly, that is, it meets the two-way extreme judgment condition, then it is enhanced by utilizing the continuous visual feature changes in the temporally adjacent frames, where the previous frame or the next frame may provide better segmentation. Adopt bidirectional optimization to make changes using adjacent frames, including reverse optimization: And positive direction optimization: Where X t represents the original input image at frame t, Represents an operation that extracts the average brightness value of the image within the mask, or the average brightness value of the entire image if no input mask is present.

[0037] In step S5, a mask bounding box extractor is used to extract For the bounding box B t As the input of MedSAM, we get P t is the segmentation mask: P t =MedSAM(X t ,B t ) Assume the segmentation mask P t Contains multiple connected regions C t If for all connected regions C t，k (k∈i,…,K), satisfying the following conditions L(X t )>L(X t ,C t，k ), that is, the brightness of the connected area is less than the average brightness of the entire image, it is considered to be an error, and the target area may not be included in any connected area; in this case, we replace the original segmentation mask with the frame that meets the following conditions by comparing the brightness of the t-1 frame and the t+1 frame: Among them, L(X t ) represents the overall average brightness of the image, L(X t ,C t，k ) indicates that the image is in C t，k The average brightness of the area.

[0038] The above is only a basic implementation example of the present invention, and the above examples are only to help understand the technical solution and core ideas of the present invention. It should be pointed out that for ordinary technicians in this technical field, without departing from the principle of the present invention, several improvements and modifications can be made to the present invention, and these improvements and modifications also fall within the scope of protection of the claims of the present invention.

Claims

1. A weakly supervised 3D medical image segmentation method based on MedSAM, characterized in that: The steps include: S1. Obtain n target MRI images. We slice the 3D image data along a specific axis to generate 2D slices, which are divided into two categories: target area and target area. S2. Perform classification network training to obtain a class activation heat map, or use other methods to provide a CAM class activation heat map; perform dynamic threshold segmentation on the obtained class activation heat map to obtain a high response area as an initial segmentation mask; S3, defines a class activation map superposition module, for all initial segmentation masks from the same patient, integrates the global semantic feature relationship, superimposes the initial segmentation masks along the axis specified in S1, performs maximum and minimum normalization, and performs dynamic threshold segmentation to obtain a superimposed segmentation mask; performs an intersection operation on the initial segmentation mask and the obtained superimposed segmentation mask to reduce the information of erroneous activation in the initial segmentation mask; S4. A bidirectional common sense defense module is defined, which can perform extreme checks in both directions; if the extreme conditions are met, the class activation map of the frame is considered to be poorly performed; otherwise, the class activation map of the frame is considered to be well performed; for the poorly performing slices, its initial segmentation mask is replaced with the initial segmentation mask of the previous frame or the next frame; S5. For each frame, the initial segmentation mask is extracted by the mask bounding box extractor to obtain the bounding box of the target area. The bounding box and the original image are used as the input of the MedSAM model to obtain the enhanced segmentation mask. At the same time, a post-processing strategy based on average brightness is designed to reduce the error of the segmentation result. S6. The enhanced segmentation mask obtained in S5 is superimposed according to the operation in S3 and the superimposed segmentation mask is replaced. The initial segmentation mask and the newly obtained superimposed segmentation mask are intersected, and the final segmentation mask is obtained by going through steps S4 and S5 again.

2. The weakly supervised three-dimensional medical image segmentation method based on MedSAM according to claim 1, characterized in that: The specific operation of the classification network in step S2 is as follows: Step S2.1: For each input image, preprocess the image by cropping and flipping it using a data augmentation method; Step S2.2: Train the classification network with a learning rate of 0.001 for 10 epochs; crop and flip the image using data augmentation methods; extract the class activation heat map M for the target area in each MRI image * ; Set a hyperparameter τ, set the position greater than the product of the maximum value of the heat map and τ to 1, otherwise set it to 0, and get the initial segmentation mask for each frame 3. The weakly supervised three-dimensional medical image segmentation method based on MedSAM according to claim 2, characterized in that: In step S3, the initial segmentation mask is superimposed, normalized to the maximum and minimum, and segmented by dynamic threshold along the axis specified in S1 to obtain a superimposed segmentation mask. The specific steps are as follows: Step S3.1: The superposition operation is: Step S3.2: Maximum and minimum normalization is performed to obtain the superimposed heat map S, and the superimposed segmentation mask is obtained by dynamic threshold segmentation. A slice of 3D image data can be regarded as a frame of video data. Compared with video data, it also has continuous visual feature transformation. Then, the high response area is segmented by dynamic thresholding: Step S3.3: Reduce the random noise of individual slices by taking the intersection and strengthen the consistency signal between slices. As the initial segmentation mask after preliminary quality enhancement; in 4. The weakly supervised three-dimensional medical image segmentation method based on MedSAM according to claim 3, characterized in that: In step S4, the two-way common sense defense module has the following specific steps: Step S4.1: Determine whether the current slice has poor class activation performance, using bidirectional extreme judgment conditions: represents the highest response area of ​​the class activation heat map, Indicates the highest corresponding area of ​​the stacked graph; represents the dynamic threshold segmentation area of ​​the class activation map, Dynamic threshold segmentation area representing the stacked image; Step S4.2: If the current slice class activation performs well and does not meet the bidirectional extreme judgment condition, then no change is made; if the performance is poor and meets the bidirectional extreme judgment condition, then it is enhanced by utilizing the continuous visual feature changes in temporally adjacent frames; The previous frame or the next frame may provide better segmentation; bidirectional optimization is used to make changes using adjacent frames, including reverse optimization: And positive direction optimization: Where X t represents the original input image at frame t, Represents an operation that extracts the average brightness value of the image within the mask, or the average brightness value of the entire image if no input mask is present.

5. The weakly supervised three-dimensional medical image segmentation method based on MedSAM according to claim 4, characterized in that: In step S5, a mask bounding box extractor is used to extract For the bounding box B t As the input of MedSAM, we get P t is the segmentation mask: P t =MedSAM(X t ,B t ) Assume the segmentation mask P t Contains multiple connected regions C t ; If for all connected regions C t，k (k∈i,…,K), satisfying the following conditions L(X t )>L(X t ,C t，k ), that is, the brightness of the connected area is less than the average brightness of the entire image, it is considered to be an error, and the target area may not be included in any connected area; in this case, we replace the original segmentation mask with the frame that meets the following conditions by comparing the brightness of the t-1 frame and the t+1 frame: Among them, L(X t ) represents the overall average brightness of the image, L(X t ,C t，k ) indicates that the image is in C t，k The average brightness of the area.

Images