SAM-based image segmentation method capable of prompting multiple organs

Abstract

The invention belongs to the field of computer vision, and particularly relates to an SAM-based promptable multi-organ image segmentation method, which comprises the following steps: acquiring a to-be-processed medical image and a region-of-interest specific prompt, and preprocessing the to-be-processed medical image and the region-of-interest specific prompt to obtain a corresponding superficial layer image block, a superficial layer embedding vector and a prompt vector; the shallow image blocks and the shallow embedded vectors are sent to a double-branch multi-stage feature encoder and processed, and mixed features are obtained; processing the mixed feature and the prompt vector by adopting a feature decoder to obtain a corresponding prediction result; and calculating the total loss of the model, adjusting model parameters according to a loss value, and finally obtaining a trained visceral organ image segmentation model capable of being prompted. According to the invention, the accuracy of the medical image segmentation result is improved.

Description

Technical Field

[0001] The invention belongs to the field of computer vision, and in particular relates to a SAM-based multi-organ image segmentation method. Background Art

[0002] Medical image segmentation aims to accurately divide structures or tissues in medical images into different regions or labels using computer-aided techniques. This enables clinicians to diagnose lesions, locate organs, or plan surgery more accurately. Medical image segmentation is of great significance to medical image analysis, computer-aided diagnosis, and personalized medicine.

[0003] In 2023, META AI published a paper titled Segment Anything, in which the proposed "Segment Anything Model" (SAM) caused a huge sensation. SAM is a basic model in the field of image segmentation, which has the characteristics of "promptable, zero-sample generalization, and downstream application-friendly". However, studies have shown that SAM is not suitable for the field of medical images, and researchers attribute this problem to the "field gap". Although SAM uses a large amount of natural image data and image masks in the training stage, images such as CT and MRI in the medical field do not fall into the category of natural images.

[0004] Visual converters use embedded vectors to represent the relative position of each image block in the entire image, so that the model can learn this information, which can ensure that the model is good at capturing global information when processing images. However, for medical image segmentation, the boundary accuracy of the segmented area is also a major factor in measuring the segmentation effect. Therefore, networks based on visual converters often cannot complete the task of medical image segmentation well. In addition, the training of visual converters has high hardware requirements and a long training time, which greatly increases the training cost of the model. Summary of the invention

[0005] In order to solve the problems existing in the background technology, improve the accuracy of medical image segmentation results and reduce the training cost of the model, the present invention provides a SAM-based multi-organ image segmentation method, comprising: inputting a medical image to be segmented into a trained medical image segmentation model to obtain a segmented target area map; wherein the medical image segmentation model comprises: a shallow feature extractor based on a convolutional neural network, a shallow feature extractor based on a visual converter, a dual-branch multi-stage feature encoder and a feature decoder;

[0006] The shallow feature extractor based on convolutional neural network is used to preprocess the medical image to be segmented to obtain shallow image blocks;

[0007] The shallow feature extractor based on visual converter is used to preprocess the medical image to be segmented to obtain a shallow embedding vector;

[0008] The dual-branch multi-stage feature encoder is used to perform feature fusion of shallow image blocks and shallow embedding vectors to obtain mixed features;

[0009] The feature decoder is used to decode the prompt vector and the mixed features of the target area to obtain a segmented target area map.

[0010] Preferably, the shallow feature extractor based on convolutional neural network uses 4 convolution layers with a convolution kernel size of 2 and a step size of 2 to preprocess the medical image to be segmented. Except for the last convolution layer, each convolution layer will stretch the channel dimension of the input feature map to 4 times, and the last convolution layer will stretch the channel dimension of the input feature map to a preset dimension to obtain a shallow image block.

[0011] Preferably, the shallow feature extractor based on visual transformer segments the medical image to be segmented into image blocks of fixed size, and then flattens each image block into a one-dimensional image block vector; and maps the flattened image block vector into a high-dimensional space through a linear layer to obtain a shallow embedding vector R H×W×C , where H stands for height, W stands for width, and C stands for channel.

[0012] Preferably, the dual-branch multi-stage feature encoder comprises: N cascaded single-stage dual-branch feature extractors; each single-stage dual-branch feature extractor comprises: a feature extraction branch based on a convolutional neural network, a feature extraction branch based on a visual converter, and an adaptive feature fusion device;

[0013] In the feature extraction branch based on the convolutional neural network, the input shallow image block is input into the detail extraction convolution layer and passes through the convolution layers of sizes 1×1, 3×3, and 1×1 in sequence to fully extract the local features and boundary information in the input shallow image block, and then passes through the redistribution convolution module, which includes a 3×3 convolution layer and a deconvolution layer, to remap the input feature map to a space of size 16×16 to obtain the first intermediate image feature; the first intermediate image feature and the input shallow image block are feature added to obtain the second intermediate image feature;

[0014] The input shallow embedding vector in the feature extraction branch based on the visual converter will be linearly projected through the self-attention layer of the visual converter to obtain the Q, K, and V vectors of the feature; then the self-attention operation of the visual converter is performed to extract the global information of the image and obtain the attention feature vector; the attention feature vector is then input into the fine-tuning module, and feature extraction is performed from the channel dimension and the spatial dimension to obtain a hybrid fine-tuning embedding feature, and the hybrid fine-tuning embedding feature and the input shallow embedding vector are added to obtain the first intermediate embedding feature; the first intermediate embedding feature is input into the multi-layer perceptron, and different abstract levels of the vector are learned layer by layer to obtain the second intermediate embedding feature; the first intermediate embedding feature and the second intermediate embedding feature are added to obtain the third intermediate embedding feature;

[0015] In the adaptive feature fuser, the second intermediate image feature and the third intermediate embedding feature are fused to obtain the input shallow image block and input shallow embedding vector of the next single-stage dual-branch feature extractor:

[0016]

[0017] in, represents the input shallow image patch of the i-th single-stage two-branch feature extractor, represents the input shallow embedding vector of the i-th single-stage dual-branch feature extractor; when i = 1, and Represents the shallow image patch and shallow embedding vector of the input dual-branch multi-stage feature encoder;

[0018] The shallow image block and shallow embedding vector output by the last single-stage dual-branch feature extractor are added to obtain the mixed feature.

[0019] Preferably, extracting features from the channel dimension and the spatial dimension to obtain the first intermediate embedding feature includes: passing the attention feature vector through an average pooling layer, and then performing weight extraction on the channel dimension. The specific process is as follows:

[0020] X channel =X in ⊙W channel

[0021] W channel =Sigmoid(Linear(AP(X in )))

[0022] Where X channel represents the feature map processed by the channel dimension, X in represents the attention feature vector, W channelRepresents the weight of the channel dimension. AP represents the average pooling layer, and Linear represents the linear layer. By obtaining the weight information of each channel, the model will determine which channel's information is more important.

[0023] The feature map processed by the channel dimension is used to extract the features of the spatial dimension. First, a 4×4 convolution is performed to compress the spatial dimension, and then a deconvolution is performed to restore the spatial dimension to obtain the first intermediate image feature.

[0024] Preferably, the feature decoder employs the decoding part of SAM.

[0025] Preferably, the hint vector of the target area includes: boundary box information of the target area or information of a plurality of pixels in the target area.

[0026] Preferably, when training the medical image segmentation model, for each training sample, the model will perform 5 iterative training. After each iterative training, the model will randomly select 3 pixel point information in the error area between the predicted mask and the actual label, and splice them into the prompt vector of the target area as the prompt vector of the target area for the next round of training.

[0027] Preferably, the loss function used when training the medical image segmentation model is as follows:

[0028]

[0029] Among them, λ F , D and λ I Represents the weight of the loss function, L F represents focal loss, L D represents dice loss, L I represents iou loss; represents the mask predicted by the model and Y represents the actual label.

[0030] The present invention has at least the following beneficial effects

[0031] In view of the fact that the existing visual transformer network has poor segmentation effect and high training cost, the present invention improves the segmentation accuracy of the network in medical organ images by extracting the global information and local features of the image respectively and adaptively fusing the features; and reduces the training cost of the model by freezing some parameters of the model and introducing a fine-tuning module at the same time. BRIEF DESCRIPTION OF THE DRAWINGS

[0032] Figure 1 It is a schematic diagram of the framework of the method process in the present invention;

[0033] Figure 2Schematic diagram of the framework of the dual-branch multi-stage feature encoder in the present invention;

[0034] Figure 3 It is a schematic diagram of the framework of the single-stage dual-branch feature extractor in the present invention;

[0035] Figure 4 It is a schematic diagram of the framework of the fine-tuning module in the present invention. DETAILED DESCRIPTION

[0036] The following describes the embodiments of the present invention by specific examples, and those skilled in the art can easily understand other advantages and effects of the present invention from the contents disclosed in this specification. The present invention can also be implemented or applied through other different specific embodiments, and the details in this specification can also be modified or changed in various ways based on different viewpoints and applications without departing from the spirit of the present invention. It should be noted that the illustrations provided in the following embodiments only illustrate the basic concept of the present invention in a schematic manner, and the following embodiments and features in the embodiments can be combined with each other without conflict.

[0037] See also Figure 1 to Figure 4 The present invention provides a SAM-based method for segmenting multiple organ images, comprising: inputting a medical image to be segmented into a trained medical image segmentation model to obtain a segmented target region map; wherein the medical image segmentation model comprises: a shallow feature extractor based on a convolutional neural network, a shallow feature extractor based on a visual converter, a dual-branch multi-stage feature encoder and a feature decoder;

[0038] The shallow feature extractor based on convolutional neural network is used to preprocess the medical image to be segmented to obtain shallow image blocks;

[0039] The shallow feature extractor based on visual converter is used to preprocess the medical image to be segmented to obtain a shallow embedding vector;

[0040] The dual-branch multi-stage feature encoder is used to perform feature fusion of shallow image blocks and shallow embedding vectors to obtain mixed features;

[0041] The feature decoder is used to decode the prompt vector and the mixed features of the target area to obtain a segmented target area map.

[0042] In this embodiment, the present invention allows SAM to better adapt to the field of medical images while reducing computational overhead on the basis of retaining the parameters of the pre-trained SAM; most of the existing segmentation methods are based on convolutional neural networks and visual converters, and convolutional neural networks are good at capturing local information of images, while visual converters are good at extracting global information of images. Therefore, the present invention combines the advantages of the two and carefully designs an efficient dual-branch multi-stage feature encoder, so that the model can take into account the mixed features of global information and local information; to improve the effect of the target segmentation result.

[0043] Preferably, the shallow feature extractor based on convolutional neural network uses 4 convolution layers with a convolution kernel size of 2 and a step size of 2 to preprocess the medical image to be segmented. Except for the last convolution layer, each convolution layer will stretch the channel dimension of the input feature map to 4 times, and the last convolution layer will stretch the channel dimension of the input feature map to a preset dimension to obtain a shallow image block.

[0044] In this embodiment, the shallow feature extractor based on the convolutional neural network uses 4 convolution layers with a convolution kernel size of 2 and a stride of 2 to preprocess the medical image to be segmented. Except for the last convolution layer, each convolution layer stretches the channel dimension of the input feature map to 4 times. The last convolution layer stretches the channel dimension of the input feature map to 768, and divides the medical image to be segmented into 16×16×768 shallow image blocks.

[0045] Preferably, the shallow feature extractor based on visual transformer segments the medical image to be segmented into image blocks of fixed size, and then flattens each image block into a one-dimensional image block vector; and maps the flattened image block vector into a high-dimensional space through a linear layer to obtain a shallow embedding vector R H×W×C , where H stands for height, W stands for width, and C stands for channel.

[0046] In this embodiment, the shallow feature extractor based on the visual transformer divides the medical image to be segmented into 16×16 image blocks, and then flattens each 16×16 image block into a one-dimensional vector. Next, a linear layer is used to map these flattened image block vectors into a high-dimensional space to obtain a shallow embedding vector R H×W×C .

[0047] See also Figure 2 , Figure 3 and Figure 4 Preferably, the dual-branch multi-stage feature encoder comprises: N cascaded single-stage dual-branch feature extractors; each single-stage dual-branch feature extractor comprises: a feature extraction branch based on a convolutional neural network, a feature extraction branch based on a visual converter, and an adaptive feature fusion device;

[0048] In the feature extraction branch based on the convolutional neural network, the input shallow image block is input into the detail extraction convolution layer and passes through the convolution layers of sizes 1×1, 3×3, and 1×1 in sequence to fully extract the local features and boundary information in the input shallow image block, and then passes through the redistribution convolution module, which includes a 3×3 convolution layer and a deconvolution layer, to remap the input feature map to a space of size 16×16 to obtain the first intermediate image feature; the first intermediate image feature and the input shallow image block are feature added to obtain the second intermediate image feature;

[0049] The input shallow embedding vector in the feature extraction branch based on the visual converter will be linearly projected through the self-attention layer of the visual converter to obtain the Q, K, and V vectors of the feature; then the self-attention operation of the visual converter is performed to extract the global information of the image and obtain the attention feature vector; the attention feature vector is then input into the fine-tuning module, and feature extraction is performed from the channel dimension and the spatial dimension to obtain a hybrid fine-tuning embedding feature, and the hybrid fine-tuning embedding feature and the input shallow embedding vector are added to obtain the first intermediate embedding feature; the first intermediate embedding feature is input into the multi-layer perceptron, and different abstract levels of the vector are learned layer by layer to obtain the second intermediate embedding feature; the first intermediate embedding feature and the second intermediate embedding feature are added to obtain the third intermediate embedding feature;

[0050] In the adaptive feature fuser, the second intermediate image feature and the third intermediate embedding feature are fused to obtain the input shallow image block and input shallow embedding vector of the next single-stage dual-branch feature extractor:

[0051]

[0052] in, represents the input shallow image patch of the i-th single-stage two-branch feature extractor, represents the input shallow embedding vector of the i-th single-stage dual-branch feature extractor; when i = 1, and Represents the shallow image patch and shallow embedding vector of the input dual-branch multi-stage feature encoder; parameter α∈(0.5,1);

[0053] The shallow image block and shallow embedding vector output by the last single-stage dual-branch feature extractor are added to obtain the mixed feature.

[0054] In this embodiment, in the feature extraction branch based on the visual converter, the present invention introduces a fine-tuning module to extract features from the channel dimension and the spatial dimension to obtain the first intermediate embedding feature. When training the model, the self-attention layer and the multi-layer perceptron with the highest computational cost can be frozen, and only the parameters of the fine-tuning module can be adjusted to obtain excellent segmentation capabilities, which greatly reduces the computational overhead of the model.

[0055] Preferably, extracting features from the channel dimension and the spatial dimension to obtain the first intermediate embedding feature includes: passing the attention feature vector through an average pooling layer, and then performing weight extraction on the channel dimension. The specific process is as follows:

[0056] X channel =X in ⊙W channel

[0057] W channel =Sigmoid(Linear(AP(X in )))

[0058] Where X channel represents the feature map processed by the channel dimension, X in represents the attention feature vector, W channel Represents the weight of the channel dimension. AP represents the average pooling layer, and Linear represents the linear layer. By obtaining the weight information of each channel, the model will determine which channel's information is more important.

[0059] The feature map processed by the channel dimension is used to extract the features of the spatial dimension. First, a 4×4 convolution is performed to compress the spatial dimension, and then a deconvolution is performed to restore the spatial dimension to obtain the first intermediate image feature.

[0060] In this embodiment, the feature decoder uses the decoding part of SAM, inputs the mixed feature and the hint vector of the target area, and obtains the segmentation result of the image.

[0061] Preferably, the hint vector of the target area includes: bounding box information of the target area or multiple pixel point information within the target area; the bounding box information of the target area is based on the user's box selection information on the area to be segmented, and the multiple pixel point information of the target area includes multiple pixel point information selected by the user in the target area.

[0062] In this embodiment, the interactivity between the user and the model is improved by this selection method. The user can select the area of ​​interest and pass it to the model in the form of pixel points or frame selection information, providing a customized segmentation function.

[0063] Preferably, when training the medical image segmentation model, for each training sample, the model will perform 5 iterative training. After each iterative training, the model will randomly select 3 pixel point information in the error area between the predicted mask and the actual label, and splice them into the prompt vector of the target area as the prompt vector of the target area for the next round of training.

[0064] In this embodiment, the segmentation performance of the pre-trained SAM is retained while reducing the computational cost. When training the model, the parameters of the self-attention layer and the multi-layer perceptron in the dual-branch multi-stage feature encoder are frozen, and only the parameters of the fine-tuning module are adjusted.

[0065] In this embodiment, the iterative training method adopted by the model can effectively improve the model's fit to the training data, and ultimately improve the segmentation accuracy of the model in actual use.

[0066] Preferably, the loss function used when training the medical image segmentation model is as follows:

[0067]

[0068] Among them, λ F , D and λ I Represents the weight of the loss function, L F represents focal loss, L D represents dice loss, L I represents iou loss; represents the mask predicted by the model and Y represents the actual label.

[0069] In this embodiment, through the above loss function design, there is no need to design different loss functions for different prompt modes during the training process. This unified training structure allows the model to focus more on the overall target rather than biased towards category loss or regression loss. The loss function used in this method is of great help to the optimization of the model and can obtain better results.

[0070] In summary, the present invention mainly improves SAM and adopts a dual-branch multi-stage feature encoder to optimize the feature extraction process, effectively improving the segmentation accuracy of organ images. At the same time, a fine-tuning module is added to reduce the computational overhead of model training, so that the model retains the parameters of the pre-trained SAM and can be well adapted to the medical image dataset.

[0071] Finally, it should be noted that the above embodiments are only used to illustrate the technical solution of the present invention rather than to limit it. Although the present invention has been described in detail with reference to the preferred embodiments, those skilled in the art should understand that the technical solution of the present invention can be modified or replaced by equivalents without departing from the purpose and scope of the technical solution, which should be included in the scope of the claims of the present invention.

Claims

1. A SAM-based multi-organ image segmentation method, characterized in that: include: Inputting the medical image to be segmented into a trained medical image segmentation model to obtain a segmented target area map; wherein the medical image segmentation model includes: a shallow feature extractor based on a convolutional neural network, a shallow feature extractor based on a visual converter, a dual-branch multi-stage feature encoder and a feature decoder; The shallow feature extractor based on convolutional neural network is used to preprocess the medical image to be segmented to obtain shallow image blocks; The shallow feature extractor based on visual converter is used to preprocess the medical image to be segmented to obtain a shallow embedding vector; The dual-branch multi-stage feature encoder is used to perform feature fusion of shallow image blocks and shallow embedding vectors to obtain mixed features; The feature decoder is used to decode according to the prompt vector and the mixed features of the target area to obtain a segmented target area map.

2. The SAM-based multi-organ image segmentation method according to claim 1, characterized in that: The shallow feature extractor based on convolutional neural network uses 4 convolution layers with a convolution kernel size of 2 and a stride of 2 to preprocess the medical image to be segmented. Except for the last convolution layer, each convolution layer will stretch the channel dimension of the input feature map to 4 times. The last convolution layer will stretch the channel dimension of the input feature map to a preset dimension to obtain a shallow image block.

3. The SAM-based multi-organ image segmentation method according to claim 1, characterized in that: The shallow feature extractor based on visual transformer divides the medical image to be segmented into image blocks of fixed size, and then flattens each image block into a one-dimensional image block vector; a linear layer is used to map the flattened image block vector to a high-dimensional space to obtain a shallow embedding vector R H×W×C , where H stands for height, W stands for width, and C stands for channel.

4. The SAM-based multi-organ image segmentation method according to claim 1, characterized in that: The dual-branch multi-stage feature encoder comprises: N cascaded single-stage dual-branch feature extractors; each single-stage dual-branch feature extractor comprises: a feature extraction branch based on a convolutional neural network, a feature extraction branch based on a visual converter, and an adaptive feature fusion device; In the feature extraction branch based on the convolutional neural network, the input shallow image block is input into the detail extraction convolution layer and passes through the convolution layers of sizes 1×1, 3×3, and 1×1 in sequence to fully extract the local features and boundary information in the input shallow image block, and then passes through the redistribution convolution module, which includes a 3×3 convolution layer and a deconvolution layer, to remap the input feature map to a space of size 16×16 to obtain the first intermediate image feature; the first intermediate image feature and the input shallow image block are feature added to obtain the second intermediate image feature; In the feature extraction branch based on the visual converter, the input shallow embedding vector will be linearly projected through the self-attention layer of the visual converter to obtain the Q, K, and V vectors of the feature; then the self-attention operation of the visual converter is performed to extract the global information of the image and obtain the attention feature vector; the attention feature vector is then input into the fine-tuning module, and feature extraction is performed from the channel dimension and the spatial dimension to obtain a hybrid fine-tuning embedding feature, and the hybrid fine-tuning embedding feature and the input shallow embedding vector are added to obtain the first intermediate embedding feature; the first intermediate embedding feature is input into the multi-layer perceptron, and different abstract levels of the vector are learned layer by layer to obtain the second intermediate embedding feature; the first intermediate embedding feature and the second intermediate embedding feature are added to obtain the third intermediate embedding feature; In the adaptive feature fuser, the second intermediate image feature and the third intermediate embedding feature are fused to obtain the input shallow image block and input shallow embedding vector of the next single-stage dual-branch feature extractor: in, represents the input shallow image patch of the i-th single-stage two-branch feature extractor, represents the input shallow embedding vector of the i-th single-stage two-branch feature extractor; when i = 1, and Represents the shallow image patch and shallow embedding vector of the input dual-branch multi-stage feature encoder; The shallow image block and shallow embedding vector output by the last single-stage dual-branch feature extractor are added to obtain the mixed feature.

5. The SAM-based multi-organ image segmentation method according to claim 4, characterized in that: The step of extracting features from the channel dimension and the spatial dimension to obtain the first intermediate embedding feature includes: passing the attention feature vector through an average pooling layer, and then performing weight extraction on the channel dimension. The specific process is as follows: X channel =X in ⊙W channel W channel =Sigmoid(Linear(AP(X in ))) Where X channel represents the feature map processed by the channel dimension, X in represents the attention feature vector, W channel Represents the weight of the channel dimension. AP represents the average pooling layer, and Linear represents the linear layer. By obtaining the weight information of each channel, the model will determine which channel's information is more important. The feature map processed by the channel dimension is used to extract the features of the spatial dimension. First, a 4×4 convolution is performed to compress the spatial dimension, and then a deconvolution is performed to restore the spatial dimension to obtain the first intermediate image feature.

6. The SAM-based multi-organ image segmentation method according to claim 1, characterized in that: The feature decoder uses the decoding part of SAM.

7. The SAM-based multi-organ image segmentation method according to claim 1, characterized in that: The hint vector of the target area includes: boundary box information of the target area or information of a plurality of pixels in the target area.

8. The SAM-based multi-organ image segmentation method according to claim 1, characterized in that: When training the medical image segmentation model, for each training sample, the model will perform 5 iterative training. After each iterative training, the model will randomly select 3 pixel information in the error area between the predicted mask and the actual label, and splice it into the prompt vector of the target area as the prompt vector of the target area for the next round of training.

9. The SAM-based multi-organ image segmentation method according to claim 1, characterized in that: The loss function used when training the medical image segmentation model is as follows: Among them, λ F , D and λ I Represents the weight of the loss function, L F represents focal loss, L D represents dice loss, L I represents iouloss; represents the mask predicted by the model and Y represents the actual label.

Images