Weakly supervised medical image segmentation method and device based on sam collaborative learning and cross-layer feature aggregation enhancement

Abstract

The application discloses a weakly supervised medical image segmentation method and device based on SAM collaborative learning and cross-layer feature aggregation enhancement, which improves segmentation accuracy by combining a weakly supervised segmentation network with SAM collaborative learning. The application proposes a cross-layer feature aggregation module to integrate adjacent layers across layers to enhance the representation ability of features at different resolutions. To enhance the performance of the segmenter, a SAM segmentation mask constraint based on collaborative learning is proposed to improve the accuracy of the segmentation result. In addition, to solve the uncertainty problem of the SAM pseudo label, the application proposes a filtering mechanism to discard the wrong pseudo label under the guidance of the scribble label combined with the cross-consistency constraint. The application effectively improves the performance of the weakly supervised network in segmenting medical images.

Description

Technical Field

[0001] This invention relates to the field of image processing, and in particular to a weakly supervised medical image segmentation method and apparatus based on SAM collaborative learning and cross-layer feature aggregation enhancement. Background Technology

[0002] Medical image segmentation is a crucial task, aiming to accurately extract and label specific structures or tissues, such as organs and tumors, in medical images, providing valuable information for fields such as medical diagnosis. In recent years, deep learning-based medical image segmentation has made significant progress, providing reliable support for clinical decision-making. However, most methods employ fully supervised training strategies that rely on large-scale pixel-level annotations, which typically require highly specialized experience, are time-consuming and costly. The powerful segmentation performance of SAM (Segment Anything Model) has attracted considerable attention. Therefore, weakly supervised learning (WSL) combined with SAM offers an effective solution to this problem and is gaining increasing popularity.

[0003] In recent years, various weakly supervised medical image segmentation methods have been explored. In these methods, different sparse labels are applied to segmentation tasks for various organs. For example, Yu et al. proposed a weakly supervised cell segmentation framework based on scribble-level annotation (Yu, Siyue, et al. "Structure-consistent weakly supervised salient object detection with local saliency coherence." In: AAAI conference on artificial intelligence.) to achieve automatic cell segmentation. Liu et al. proposed a COVID-19 infection CT slice segmentation method to segment the lung infection region (Liu, Xiaoming, et al. "Weakly supervised segmentation of COVID-19 infection with scribble annotation on CT images."). Jiang et al. proposed using SAM to generate pseudo-labels, and then using these pseudo-labels to train a weakly supervised semantic segmentation model (Jiang, Peng-Tao, and Yuqi Yang. "Segment Anything is A Good Pseudo-labelGenerator for Weakly Supervised Semantic Segmentation." arXiv preprint). Chen et al. introduced the CAM method to address the challenges of distinguishing the boundary between foreground and background and to handle obvious co-occurrence during training (Chen, Zhang, et al. "C-cam: Causal cam for weakly supervised semantic segmentation on medical image." In: Computer Vision and Pattern Recognition. 2022.). Additionally, Roth proposed using extreme point clicks as a minimal form of user interaction to train deep learning-based segmentation models (Roth, Holger R., et al. "Going to extremes: weakly supervised medical image segmentation." Machine Learning and Knowledge Extraction 3.2(2021):507-524.).The methods described above clearly demonstrate that weakly supervised medical image segmentation remains a challenging task. Therefore, it is crucial to propose an efficient and accurate weakly supervised segmentation strategy to address this issue. Summary of the Invention

[0004] The purpose of this invention is to provide a weakly supervised medical image segmentation method and apparatus based on SAM collaborative learning and cross-layer feature aggregation enhancement.

[0005] The technical solution for achieving the objective of this invention is as follows: Firstly, this invention provides a weakly supervised medical image segmentation method based on SAM collaborative learning and cross-layer feature aggregation enhancement, comprising:

[0006] Multiple different medical image datasets were collected to form training and test sets. The training set was processed with different labels using a weakly supervised method.

[0007] A neural network CEANet with Res2Net50 as the backbone network framework and including cross-layer feature fusion and multi-layer feature enhancement is constructed.

[0008] The neural network is trained using a prediction mask provided by SAM with various loss constraints, the network parameters are optimized, and the best model is selected, including partial loss constraints and structural consistency constraints.

[0009] The test set samples are input into the final model file to obtain the final segmented image.

[0010] Secondly, the present invention provides a weakly supervised medical image segmentation device based on SAM collaborative learning and cross-layer feature aggregation enhancement, comprising:

[0011] The data acquisition module collects multiple datasets of medical images from different datasets to form training and testing samples;

[0012] The cross-layer feature aggregation module is used to construct a neural network with Res2Net as the backbone network and fuse the features of every two adjacent layers.

[0013] The multi-layer feature enhancement module is used to aggregate image feature information at different scales, so that both high-level and low-level features at different scales are preserved.

[0014] The structural consistency and cross-consistency modules are used to constrain image predictions, utilizing graffiti annotations and masks generated from SAM to constrain the supervised segmentation model.

[0015] The segmentation image generation module is used to input test samples into the final segmentation model to obtain the final segmentation image.

[0016] Thirdly, the present invention provides an electronic device including a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor executes the program to implement the steps of the method described in the first aspect.

[0017] Fourthly, the present invention 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 the first aspect.

[0018] Compared with the prior art, the significant advantages of the present invention are:

[0019] (1) A cross-level feature aggregation module is proposed to enhance the representation capability of features at different resolutions by fusing cross-level features.

[0020] (2) This invention proposes a multi-layer feature enhancement module, which integrates high-level multi-scale features and then uses the integrated features to enhance the features of the current layer. It focuses on high-level semantics while taking into account low-level features, effectively improving the network's perception capability.

[0021] (3) This invention proposes a prompt box generation strategy to optimize the collaborative learning process. By fusing the boxes obtained from the graffiti labels and predicted features, it provides good prompts for SAM so as to generate higher quality mask images.

[0022] (4) Experimental results on five widely used polyp datasets show that SAM-CLNet outperforms existing weakly supervised segmentation methods. Attached Figure Description

[0023] Figure 1 This is a structural diagram of SAM-CLNet.

[0024] Figure 2 The diagram shows the framework of the module of this invention, where (a) is CEM and (b) is FAM.

[0025] Figure 3 This is a flowchart illustrating the generation process of the bounding box in this invention.

[0026] Figure 4 A visual comparison chart of the results from this invention and the SCW method. Detailed Implementation

[0027] like Figures 1-3 As shown, this invention proposes a weakly supervised medical image segmentation method based on SAM collaborative learning and cross-layer feature aggregation enhancement, comprising:

[0028] (1) Provide a cross-layer feature enhancement module to integrate the feature information of two adjacent layers and enhance the representation capability of features at different resolutions.

[0029] Cross-level Enhancement Module: This module adjusts the high-level features F by upsampling. i+1 As input to CEM, these two features are each passed through a 3×3 convolutional layer to obtain F. l and F h To obtain preliminary characteristics, F l and F h F represents the low-level and high-level features in two adjacent feature layers, respectively. l =Conv(F i ), F h =Conv(U(f) i+1 U represents performing an upsampling operation on the feature map; F l and F h Each feature is smoothed by a 3×3 convolutional layer and then fed into an activation function to learn each other's weights. Pixel-wise multiplication is used to enhance the F-value. l and F h Get F′ l and F′ h , where F′ h =S(Canv(F) l ))×Conv(F h ),F′ l =S(U(Conv(F) h )))×Conv(F l ), where S represents the feature activation sigmoid, U represents the upsampled features, and the original concatenated features are fused to obtain Where F is the fusion feature, represented as F = Conv(Cat(F) l ,F h Then The input is fed into the corresponding layer of the decoder.

[0030] (2) Provides a multi-level feature aggregation module: used for feature enhancement to capture richer features across multiple different levels. First, cascade features. i>2. Features are extracted using a 1×1 convolutional layer to reduce the number of channels. O cat This represents the concatenation operation. Feature F cas Feature F is generated after batch normalization and ReLU activation function using a 3×3 convolutional layer. con After that, in F conWe apply global average pooling to generate a global weight map using the sigmoid activation function. We then compare the resulting weight map with the corresponding feature F. con Element-wise multiplication produces a globally attention-enhanced feature F. g Stack F on the initial enhanced features con and F g Residual connections are applied, and finally, the fused features are batch normalized using 3×3 convolutions and activated by ReLU to obtain the final segmentation features. This process can be represented as: Where S represents the sigmoid operation, O avg This is global average pooling.

[0031] (3) Provides a collaborative learning strategy for structural consistency and cross-consistency constraints: fully utilizes the main segmentation prediction S1 and graffiti annotation G. s This generates a prompt box. The prompt box is then input into SAM to generate a mask, which serves as additional labels for training our segmenter. This constraint is denoted as L. bce (S1,G M In this process, our two branch networks can work together to optimize for more promising segmentation results by minimizing L. dom To optimize the segmentation process, defined as l dom =L ce (S1,G s )+L bce (S1, G) M )+L ss (S1, S) ↓ ), S1,G M These represent the prediction image and the mask provided by SAM, respectively. ↓ It is a predicted output image, a scaled-down version of the original input image. ce L represents the cross-entropy loss function. mse This represents the widely used mean squared error loss, and L ss This represents the structural consistency loss function.

[0032] Based on the same inventive concept, the present invention also provides a weakly supervised medical image segmentation device based on SAM collaborative learning and cross-layer feature aggregation enhancement, comprising:

[0033] The data acquisition module collects multiple datasets of medical images from different datasets to form training and testing samples;

[0034] The cross-layer feature aggregation module is used to construct a neural network with Res2Net as the backbone network and fuse the features of every two adjacent layers.

[0035] The multi-layer feature enhancement module is used to aggregate image feature information at different scales, so that both high-level and low-level features at different scales are preserved.

[0036] The structural consistency and cross-consistency modules are used to constrain image predictions, utilizing graffiti annotations and masks generated from SAM to constrain the supervised segmentation model.

[0037] The segmentation image generation module is used to input test samples into the final segmentation model to obtain the final segmentation image.

[0038] The implementation methods of the above modules are the same as those of the aforementioned weakly supervised medical image segmentation method, and will not be repeated here.

[0039] The detailed steps of the method of the present invention will be described below with reference to the accompanying drawings and embodiments.

[0040] Example

[0041] Figure 1 This is a structural diagram of SAM-CLNet.

[0042] 1.1 Collect the public datasets ETIS, CVC-ClinicDB, CVC-ColonDB, EndoScene and Kvasir, and take 1450 images from CVC-ClinicDB and Kvasir as the training set, and the rest as the test set.

[0043] 1.2 The image resolution was adjusted to 320×320 using bilinear interpolation.

[0044] 2.1 Input image I into the model, and after encoding, obtain two sets of features. F1 represents the lowest level of features, and F5 represents the highest level of features.

[0045] 2.2 Input F5 and F4 into the CEM module for fusion to obtain F4 and F3 are input into the CEM module for fusion to obtain... F3 and F2 are input into the CEM module for fusion to obtain... F2 and F1 are input into the CEM module and fused to obtain... We get

[0046] 2.3 The results obtained in 2.2 It is used as input to the FAM module for high-level feature fusion. This layer directly passes through convolution to output the predicted image S4. Then... As a high-level feature enhancement The corresponding prediction map S3 is obtained, and so on. and After fusion, it is used as a high-level fusion feature to enhance The corresponding prediction map S2 is obtained, and Used to enhance after fusion The corresponding prediction map S1 is obtained as the final model output.

[0047] 3.1 Combine S1 with the graffiti annotations for calculation and processing to generate a prompt box, and give SAM a prompt as the prompt information for its segmentation.

[0048] 3.2 Using the mask generated by SAM, structural consistency constraints are added to calculate the constraint surface L of the predicted graph S1. bce (S1, G) M ).

[0049] 3.3 Apply foreground-background consistency constraints to S1 for the graffiti annotations, and calculate L. ce (S1,G s ).

[0050] 3.4 The image after scaling the input image by 0.3 is input into the network to obtain a smaller-scale prediction map. Calculate the consistency loss L ss (S1,S ↓ ).

[0051] 4.1 Calculate the gradient and update the weights of the segmentation network.

[0052] 5.1 Test the trained network on the test dataset to obtain metrics and visualization results, and compare them with the SCW method.

[0053] Figure 4 A visualization comparison of the results of this invention and the SCW method is provided. It is evident that the SAM-CLNet of this invention outperforms existing weakly supervised segmentation methods.

[0054] The embodiments described above are merely illustrative of several implementation methods of this application, and while the descriptions are relatively specific and detailed, they should not be construed as limiting the scope of the invention patent. It should be noted that those skilled in the art can make various modifications and improvements without departing from the concept of this application, and these all fall within the protection scope of this application. Therefore, the protection scope of this patent application should be determined by the appended claims.

Claims

1. A weakly supervised medical image segmentation method based on SAM collaborative learning and cross-layer feature aggregation enhancement, characterized in that, include: Multiple different medical image datasets were collected to form training and test sets. The training set was processed with different labels using a weakly supervised method. The CEANet neural network, with Res2Net50 as its backbone and incorporating cross-layer feature fusion and multi-layer feature enhancement, is constructed as follows: After passing through an encoder with Res2Net50 as the backbone network, features are obtained. , and First pass The number of channels in the convolutional layers is unified; Adjusting high-level features through upsampling operations To make it have the same The same resolution; then, each of these two features is processed through a... The convolutional layer is obtained and To obtain preliminary characteristics, and These represent the low-level and high-level features in two adjacent feature layers, respectively. U represents performing an upsampling operation on the feature map; Will and Each of them passed through one The smooth features of the convolutional layers are fed into the activation function to learn each other's weights, and pixel-wise multiplication is used to enhance them. and get and ,in , , where S represents the sigmoid activation function; After fusing the original cascaded features, we obtain Where F is the fusion feature, denoted as Then The input is fed into the corresponding layer of the decoder, where Cat indicates that a concatenation operation is performed on the feature map; The neural network is trained using a prediction mask provided by SAM with various loss constraints, the network parameters are optimized, and the best model is selected, including partial loss constraints and structural consistency constraints. The test set samples are input into the final model file to obtain the final segmented image.

2. The weakly supervised medical image segmentation method based on SAM collaborative learning and cross-layer feature aggregation enhancement according to claim 1, characterized in that, In the first Layer characteristics When introducing the decoder, first aggregate all features from the previous CEM modules, and then combine the aggregated features with... Combined, and then fed into the decoder to generate the predicted segmentation map; for First, cascade features , ,.., All these features are cascaded and then passed through a Convolutional layers are used to reduce channels; therefore, cascaded features can be obtained. , described as ,in Represents the concatenation operation; characteristics After a The convolutional layers are followed by batch normalization and ReLU activation function to generate features. Then, in Global average pooling is applied, followed by a sigmoid activation function to generate a global weight map; the weight map and the corresponding features... Element-wise multiplication produces globally attention-enhanced features. To preserve the original feature information of each layer, stacking is performed on top of the initial enhanced features. and To apply residual connections; then, the fused features are processed through a... The final features are obtained by batch normalization of convolutions and ReLU activation function; This process can be represented as: Where S represents the sigmoid operation, This is global average pooling.

3. The weakly supervised medical image segmentation method based on SAM collaborative learning and cross-layer feature aggregation enhancement according to claim 1, characterized in that, The SAM-based collaborative learning method is as follows: In the weakly supervised segmenter, the main segmentation prediction S1 and graffiti annotations are used to generate the prompt "prompt"; These prompts are then fed into SAM to generate masks, which serve as additional labels for training the segmenter. In this process, the two branch networks work together to optimize in order to achieve a more promising segmentation result by minimizing... To optimize the segmentation process, defined as , These represent the prediction image and the mask provided by SAM, respectively. This is the predicted output image of the downsampled version of the original input image. The mask provided by the graffiti annotation. Represents the cross-entropy loss function. This represents the widely used mean squared error loss, and This represents the structural consistency loss function.

4. A weakly supervised medical image segmentation device based on SAM collaborative learning and cross-layer feature aggregation enhancement, characterized in that, include: The data acquisition module collects multiple datasets of medical images from different datasets to form training and testing samples; The cross-layer feature aggregation module is used to construct a neural network with Res2Net as the backbone, fusing features from every two adjacent layers. The specific implementation is as follows: After passing through an encoder with Res2Net50 as the backbone network, features are obtained. , and First pass The convolutional layers achieve channel uniformity, and high-level features are adjusted through upsampling operations. To make it have the same The same resolution; then, each of these two features is processed through a... The convolutional layer is obtained and To obtain preliminary characteristics; and Each of them passed through one The smooth features of the convolutional layers are fed into the activation function to learn each other's weights, and pixel-wise multiplication is used to enhance them. and get and ,in , After fusing the original cascaded features, the following is obtained: Where F is the fusion feature, represented as Then The input is fed into the corresponding layer of the decoder, where Cat indicates that a concatenation operation is performed on the feature map; The multi-layer feature enhancement module is used to aggregate image feature information at different scales, so that both high-level and low-level features at different scales are preserved. The structural consistency and cross-consistency modules are used to constrain image predictions, utilizing graffiti annotations and masks generated from SAM to constrain the supervised segmentation model. The segmentation image generation module is used to input test samples into the final segmentation model to obtain the final segmentation image.

5. The weakly supervised medical image segmentation device based on SAM collaborative learning and cross-layer feature aggregation enhancement according to claim 4, characterized in that, The specific implementation method of the multi-layer feature enhancement module is as follows: First, cascaded features , ,.., Then through a Convolutional layers are used to reduce channels and obtain cascaded features. , described as ,in Represents the concatenation operation; characteristics After a The convolutional layers are followed by batch normalization and ReLU activation function to generate features. ;exist Global average pooling is applied, followed by a sigmoid activation function to generate a global weight map; the weight map and the corresponding features... Element-wise multiplication produces globally attention-enhanced features. Stacking on the initial enhanced features and To apply residual connections, the fused features pass through a The final features are obtained by batch normalization of convolutions and ReLU activation function; This process can be represented as: Where S represents the sigmoid operation, This is global average pooling.

6. The weakly supervised medical image segmentation device based on SAM collaborative learning and cross-layer feature aggregation enhancement according to claim 4, characterized in that, The specific implementation methods of the structural consistency and cross-consistency constraints in collaborative learning are as follows: Using principal segmentation prediction and graffiti annotations To generate the prompt; These prompts are then fed into SAM to generate a mask, which serves as additional labels for training the segmenter. This part of the constraint is denoted as... During this process, the two branch networks can work together to optimize and achieve a more promising segmentation result by minimizing... To optimize the segmentation process, defined as , These represent the prediction image and the mask provided by SAM, respectively. This is the predicted output image of the downsampled version of the original input image. Represents the cross-entropy loss function. This represents the widely used mean squared error loss, and This represents the structural consistency loss function.

7. An electronic device comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, characterized in that, When the processor executes the program, it implements the steps of the method as described in any one of claims 1-3.

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

Images