A weakly supervised 3D medical image segmentation method based on MedSAM

Abstract

The present invention discloses a weakly supervised medical image segmentation method based on MedSAM, which aims to solve the problem that traditional medical image segmentation relies on a large amount of labeled data and is time-consuming and labor-intensive to train, and to improve the segmentation performance under the condition of limited labeling cost. It utilizes the feature extraction capability of the deep learning model and generates a category activation map through the backbone network to preliminarily locate the target area. A CSD bidirectional common sense defense module and an AMS class activation map superposition module are proposed. The former combines global information to improve the robustness of the model, and the latter determines the maximum range of the target to ensure the segmentation effect. Combining the dynamic threshold segmentation strategy and the MedSAM model enhanced heat map, the semantic features are integrated to avoid relying on prototype learning, and weakly supervised high-precision segmentation is achieved. Experiments show that its performance in multiple data sets is better than existing methods, and it has broad application prospects. It can be used for a variety of medical imaging tasks and other fields with high labeling costs and the need for efficient segmentation.

Description

Technical Field

[0001] The present 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 a classification network and extract localization maps from a convolutional neural network based on class activation maps. However, the initial CAM is often incomplete or redundant. To address this challenge, researchers have proposed various improvement strategies at different stages to enhance the quality of CAM and the final segmentation results.

[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 commonly used methods. By erasing the salient parts of the image or feature map, the network is forced to explore more undiscovered areas. In addition, some studies have accumulated 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] To optimize the initial CAM, learning semantic affinity to improve localization map quality has become a mainstream approach. For example, PSA trains the network to learn semantic affinity between pixels, propagating the semantics of strongly responsive areas in the attention map to semantically similar pixels. IRNet and BES significantly improve segmentation by synthesizing category boundaries and extending the coverage of the target region to the boundary. Additionally, some studies have further optimized CAM by incorporating saliency maps to obtain more precise background information or distinguish co-occurring target objects.

[0005] When training segmentation models, existing methods typically generate pseudo-labels by applying a global threshold to the CAM, but ignore the confidence information of each pixel and fail to fully exploit the potential of the CAM. To address this shortcoming, the PMM model proposes a disguised underfitting strategy that improves segmentation performance by reweighting the loss of potentially noisy pixels. The URN model, on the other hand, uses conditional random fields (CRFs) to estimate uncertainty by scaling the prediction map multiple times. However, the PMM's improvements are limited to the loss level and do not directly utilize confidence information. The URN, on the other hand, requires multiple CRFs, resulting in low computational efficiency.

[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 emerged like mushrooms after rain. 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 the needs of different segmentation tasks; 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, guiding the model to output the desired segmentation results.

[0007] Many researchers are also actively exploring how to leverage existing visual foundational models to aid weakly supervised semantic segmentation. For example, Liang et al., Qin et al., and Ghiasi et al. cleverly leverage powerful visual foundational models such as CLIP or DINOv2, training them (usually through fine-tuning) to work with specific segmentation heads or adapters. This allows these foundational models to be better integrated into semantic segmentation tasks, uncovering more potential segmentation cues and improving segmentation accuracy and efficiency.

[0008] Particularly worth mentioning is the arbitrary segmentation model, which has demonstrated remarkable category-independent segmentation capabilities as soon as it came out. No matter what type of target object it is, it can perform segmentation more accurately. This outstanding performance has undoubtedly achieved great success, causing a strong response 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 labeled 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 medical image segmentation basic model MedSAM to realize a weakly supervised three-dimensional medical image segmentation method.

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

[0010] The present invention belongs to the field of medical image segmentation technology, 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 the 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 and does not 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:

[0012] 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 target area.

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

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

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

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

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

[0018] Preferably, the classification model in step S2 is specifically operated as follows:

[0019] Step S2.1: For each input image, use the data augmentation method to preprocess the image by cropping and flipping;

[0020] 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, and vice versa to 0, to obtain the initial segmentation mask for each frame

[0021]

[0022] Preferably, in step S3, the initial segmentation masks are 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:

[0023] Step S3.1: The superposition operation is:

[0024]

[0025] 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:

[0026]

[0027] 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;

[0028] in

[0029] As a preferred option, the two-way common sense defense module has the following specific steps:

[0030] Step S4.1: Determine whether the current slice has poor class activation performance, using bidirectional extreme judgment conditions:

[0031]

[0032] Indicates the highest response area of ​​the class activation heat map, Indicates the highest corresponding area of ​​the stacked graph;

[0033] Represents the dynamic threshold segmentation area of ​​the class activation map, Representing 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;

[0034] The previous or next frame may provide better segmentation; bidirectional optimization is used to make changes using adjacent frames, including optimization in the opposite direction:

[0035]

[0036] And positive direction optimization:

[0037]

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

[0039] 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 compare the brightness of the t-1 frame and the t+1 frame and replace the original segmentation mask with the frame that meets the following conditions:

[0040]

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

[0042] Compared with the prior art, the present invention has the following beneficial effects:

[0043] First, we propose a prototype-free weakly supervised medical image segmentation framework, which frees users from training or providing prototypes.

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

[0045] 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

[0046] Figure 1 This is a flow chart of the model structure of the present invention.

[0047] Figure 2 Schematic diagram of the extreme inspection conditions in the bidirectional common sense defense mechanism of the present invention. DETAILED DESCRIPTION

[0048] refer to Figure 1 , Figure 1 This is a flow chart of the model structure of the present invention, which includes:

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

[0050] S2 uses the Vi'TB-16 model pre-trained on ImageNet-21k for classification training. 2D slices of the target area are input into S2 to obtain the trained model. Grad-CAM and the Class-aware Attention-based Affinity module are used to generate the class activation heat map of the target area. Dynamic threshold segmentation is performed on the obtained class activation heat map to obtain the high-response area as the initial segmentation mask.

[0051] S3. We define a class activation map superposition module that can simply and effectively integrate global semantic feature relationships for all initial segmentation masks from the same patient. The initial segmentation masks are superimposed, normalized, and dynamically thresholded along the axis specified in S1 to obtain a superimposed segmentation mask. The activation information from a single slice is usually not enough to cover the entire target area, but 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 area. 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.

[0052] 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 a frame is considered to be poorly performing; otherwise, it is considered to be well performing. For poorly performing slices, its initial segmentation mask is replaced with the initial segmentation mask of the previous or next frame.

[0053] S5. For each frame, the initial segmentation mask is extracted through the mask bounding box extractor to obtain the bounding box of the target region. The bounding box and the original image are used as input to the MedSAM model to obtain the enhanced segmentation mask. A post-processing strategy based on average brightness is also designed to reduce the error of the segmentation results.

[0054] 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 performing steps S4 and S5 again.

[0055] The specific operations of the classification model in step S2 are as follows:

[0056] Each input image is trained using Vi'TB-16 with a small learning rate for 10 epochs. Data augmentation methods are used to preprocess the image by cropping and flipping. Grad-CAM and CAA modules are used to extract class activation heatmaps for the target region in each MRI image.

[0057] 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 Gradient:

[0058]

[0059] Next, these gradients are globally averaged pooled to obtain patch weights Then, calculate each embedding F i The patch and its corresponding weight The weighted sum of c .

[0060] Finally, the previously generated class activation maps are enhanced by utilizing the attention weights from the last multi-head self-attention layer:

[0061] M * =B c ⊙A t ·vec(M c )

[0062] 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 set it to 0 otherwise, to obtain the initial segmentation mask for each frame

[0063]

[0064] In step S3, the initial segmentation mask is superimposed, normalized, and segmented by dynamic threshold along the axis specified in S1 to obtain a superimposed segmentation mask. The superimposition operation is:

[0065]

[0066] 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 considered a frame of video data. Compared to video data, it also has continuous visual feature transformations. The sequential connections between consecutive slices can be exploited to address many challenges in single 2D image processing. Then, high-response regions are segmented using dynamic thresholding:

[0067]

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

[0069] In step S4, the bidirectional common sense defense module has the following specific steps:

[0070] Step 1: Determine whether the current slice has poor class activation performance, using bidirectional extreme judgment conditions:

[0071] 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. Bidirectional optimization is used to utilize the changes in the adjacent frames, including reverse optimization:

[0072]

[0073] And positive direction optimization:

[0074]

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

[0076] 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 compare the brightness of the t-1 frame and the t+1 frame and replace the original segmentation mask with the frame that meets the following conditions:

[0077]

[0078] Among them, L(X t ) represents the overall average brightness of the image, L(X t ,Ct，k ) indicates that the image is in C t，k The average brightness of the area.

[0079] The above is merely a basic example of the present invention, and is intended only to help understand the technical solution and core concept of the present invention. It should be noted that those skilled in the art may make various improvements and modifications to the present invention without departing from the principles of the present invention, and such 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 are as follows: 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, the global semantic feature relationship is integrated, and the initial segmentation masks are superimposed along the axis specified in S1, and the maximum and minimum normalization and dynamic threshold segmentation are performed to obtain a superimposed segmentation mask. The initial segmentation mask and the obtained superimposed segmentation mask are intersected to reduce the information of false 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 slice, its initial segmentation mask is replaced with the initial segmentation mask of the previous or next frame. S5. For each frame’s initial segmentation mask, a bounding box of the target region is obtained through a mask bounding box extractor. The bounding box and the original image are used as inputs to the MedSAM model to obtain an 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 performing 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 operations of the classification network in step S2 are 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, and vice versa to 0, to obtain 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 masks are superimposed, normalized, 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 bidirectional 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: 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, 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 or next frame may provide better segmentation; bidirectional optimization is used to make changes using adjacent frames, including optimization in the opposite direction: 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 compare the brightness of the t-1 frame and the t+1 frame and replace the original segmentation mask with the frame that meets the following conditions: 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