Spinal image segmentation method based on multiple maps and SAM model

Abstract

The invention discloses a spine image segmentation method based on multiple maps and an SAM model, and the method comprises the steps: carrying out the preliminary segmentation of a to-be-segmented image through a simple network, carrying out the registration of a prompt image and a rough segmentation result, and taking the registration as multiple maps; a ViT model of SAM pre-training weight is used as an image encoder to extract input image features, and an adapter is inserted into the input image features for fine tuning; a CNN is used as a prompt encoder to extract multi-map features, and a feature map sampled at each time is reserved; and fusing the features extracted by the two encoders by using a decoding module of the mask decoder, obtaining a feature vector through a multi-layer MLP, and finally obtaining a prediction result through point multiplication. According to the method, the segmentation result is automatically generated through the multi-atlas generation sub-network and the segmentation sub-network without manual intervention, the low-efficiency operation of manually marking points and frames in the medical segmentation field of the SAM model is improved, and meanwhile, the segmentation result is more accurate in combination with the prior knowledge of the multi-atlas.

Description

Technical Field

[0001] The present invention belongs to the field of medical image technology, and in particular relates to a spinal image segmentation method based on multi-atlas and SAM model. Background Art

[0002] The spine is an indispensable and important part of the human body, providing vital support. Accurate spinal segmentation results are also crucial for the treatment of spinal diseases, helping doctors better understand the patient's spinal anatomy and diagnose and treat them. However, accurate manual labeling of the spine is a time-consuming and inefficient task. Furthermore, manual labeling is prone to errors due to the individual experience of the labeler. This is also a challenge for automatic labeling methods based on deep learning. Due to the complex anatomical structure of the spinal region, the noise in medical images, the low contrast of the nerve region in CT, and other factors, the segmentation results of the spinal region may not be very accurate.

[0003] Multi-atlas segmentation is stable and accurate. Traditional multi-atlas segmentation uses multiple sets of images and their labels as templates. The image to be segmented is then registered with the labeled template image using the SyN or Demons algorithm, resulting in multiple registered labeled images. Labels from these multiple labeled images are then fused using methods such as weighted voting. The resulting image undergoes further post-processing to produce the final segmentation result. With the advancement of deep learning, multi-atlas segmentation and deep learning have also been integrated to achieve a combined advantage.

[0004] The SAM model provides a foundational model for image segmentation. Its high accuracy and simplicity offer new insights into segmentation. It uses the ViT model pre-trained with MAE as an image feature extractor. Because ViT models vary in size, SAM offers three models with different parameter sizes. A larger parameter size results in a larger model and more accurate segmentation results. Furthermore, it provides three methods for indicating the image region to be segmented: points, boxes, and text. The operator simply clicks on the desired area, draws a box, or enters text to accurately segment a 2D image. However, the SAM model is trained on natural images, while medical images are often acquired using methods such as X-rays and magnetic resonance imaging. Consequently, properties such as pixel intensity and texture differ significantly between medical and natural images, leading to unsatisfactory performance of the SAM model in medical image segmentation. Furthermore, 3D images such as CT and MRI are composed of hundreds of stacked 2D images. Segmenting the entire 3D image requires segmenting each individual 2D image. Manually performing these operations, such as clicking and drawing boxes on each 2D image within the 3D image, remains a complex and labor-intensive task.

[0005] Therefore, how to combine multi-atlas segmentation with the SAM model, improve the shortcomings of SAM and use it and the high accuracy of multi-atlas to complete the precise segmentation of the spinal region is the key task of the present invention. Summary of the Invention

[0006] Purpose of the invention: The present invention proposes a spinal image segmentation method based on multi-atlas and SAM model, which improves the inefficient operation of SAM model in the field of medical segmentation, which requires manual annotation of points and boxes. At the same time, it combines the prior knowledge generated by multi-atlas to make the segmentation results more accurate.

[0007] Summary of the invention: The present invention provides a spinal image segmentation method based on multi-atlas and SAM model, comprising the following steps:

[0008] (1) Use a simple network to perform preliminary segmentation on the image to be segmented to obtain a coarse segmentation result, then align the prompt image with the coarse segmentation result and use it as a multi-atlas;

[0009] (2) Use the ViT model with SAM pre-trained weights as the image encoder to extract input image features and insert the adapter into it for fine-tuning;

[0010] (3) Use CNN as a prompt encoder to extract multi-atlas features and retain the feature maps of each downsampling to help subsequent images be upsampled to higher resolution;

[0011] (4) The decoding module of the mask decoder is used to fuse the features extracted by the two encoders, and the feature vector is obtained through a multi-layer MLP, and finally the prediction result is obtained through point multiplication.

[0012] Furthermore, the implementation process of step (1) is as follows:

[0013] Randomly reserved pictures in the dataset are used as prompt images. The input image generates a coarse segmentation result after passing through the UNet segmentation network; then the SyN algorithm is used to align the label of the prompt image to the coarse segmentation result, and the corresponding registration result is used as a multi-atlas.

[0014] Furthermore, the implementation process of using the ViT model with SAM pre-trained weights as an image encoder to extract input image features in step (2) is as follows:

[0015] The two-dimensional image is divided into multiple image blocks and position encoding is added. Image features are extracted after passing through multiple layers of Transformer modules. Each Transformer module is spliced ​​in the order of a multi-head attention module, a layer normalization, an MLP module, and a layer normalization. After the second layer normalization, an adapter module is arranged in parallel with the MLP for fine-tuning. Each layer normalization uses residual connections to improve accuracy.

[0016] Furthermore, the process of inserting the adapter for fine-tuning in step (2) is as follows:

[0017] After global pooling, the input feature map is compressed by channel, and then two fully connected layers compress and then restore the feature map. The sigmoid layer then obtains the weight of the feature map by channel. After obtaining the weighted feature map, it is compressed and restored through convolution and deconvolution. After residual connection, the final output feature map is obtained through layer normalization.

[0018] Furthermore, the implementation process of step (3) is as follows:

[0019] Multi-atlas is used as input, and four downsampling modules are used for feature extraction; each downsampling module contains a maximum pooling layer and a double convolution module; each double convolution module consists of a convolution layer, layer normalization, and a GeLU layer; in the convolution module, the convolution layer increases the number of channels of the original image while keeping the image length and width unchanged, and then uses layer normalization to stabilize the network, and uses a GeLU layer for activation; after the multi-atlas is input, it first passes through a convolution module to increase the number of channels to 16 while keeping the length and width of the multi-atlas unchanged, and then after 4 downsamplings, the number of channels is gradually increased to 256 and the length and width of the multi-atlas are reduced to 1 / 16 of the original; the features extracted by each downsampling module will be retained to help restore the predicted low-resolution segmentation results to the original size.

[0020] Furthermore, in the maximum pooling layer, the convolution kernel size is 2 and the stride is 2, which reduces the length and width of the image by half while keeping the original number of channels unchanged.

[0021] Furthermore, the convolution kernel size of the convolution layer is 3, the stride is 1, and the padding is 1.

[0022] Furthermore, the mask decoder in step (4) includes a decoding module, four upsampling modules, a cross attention layer, a multi-layer MLP and a dot product operation; the decoding module fuses features from the image encoder and features from the prompt encoder to generate updated image features and prompt features; in the upsampling module, after passing through the deconvolution layer, layer normalization, and GeLU layer, the updated image features obtained are concatenated with the corresponding down-sampled features retained by the prompt encoder, and then pass through two layers of convolution modules to obtain higher resolution results; the prompt features generated by the decoding module are updated again in the token-to-image cross attention layer, and then pass through the multi-layer MLP to generate a feature vector; finally, the feature vector generated by the MLP layer is dot-multiplied with the high-resolution image features generated by the upsampling module to obtain the final segmentation result.

[0023] Furthermore, the decoding module uses the decoding module of the SAM decoder, which has two layers, and each layer structure is: self-attention layer, token-to-image cross-attention layer, MLP layer, image-to-token cross-attention layer.

[0024] Furthermore, the convolution kernel size of the deconvolution layer is 2, the step size is 2, the number of channels remains unchanged after deconvolution, and the length and width of the image are doubled.

[0025] Beneficial effects: Compared with the existing technology, the present invention has the following beneficial effects: the present invention uses the anatomical prior knowledge contained in multiple atlases to make the segmentation results more accurate; the present invention improves the inefficient operation of the SAM model in the field of medical segmentation, which requires manual annotation of points and boxes, and adopts an image encoder that fine-tunes the original SAM model and automatically generates segmentation results through a hint encoder and a mask decoder; the present invention achieves the high precision of SAM, and is fully automatic without the need for manual annotation. BRIEF DESCRIPTION OF THE DRAWINGS

[0026] Figure 1 Schematic diagram of the spine image segmentation network structure proposed by the present invention; DETAILED DESCRIPTION

[0027] The present invention will be further described below with reference to the accompanying drawings.

[0028] The present invention proposes a spine image segmentation method based on multi-atlas and SAM model. Figure 1 The spinal image segmentation network shown in the figure takes as input the spinal image to be segmented and the corresponding prompt image, and outputs the segmentation result. The image to be segmented and the prompt image are input into the multi-atlas generation subnetwork to generate a multi-atlas image. The image to be segmented is input into the image encoder to extract high-dimensional image features. The generated multi-atlas is input as a prompt into the prompt encoder for feature extraction, and the features extracted by each downsampling module are saved to help restore the predicted low-resolution segmentation result to the original size. Finally, the image features and multi-atlas features are fused in the decoder and the original resolution is restored by upsampling. Specifically, the following steps are included:

[0029] Step 1: Use a simple network to perform preliminary segmentation on the image to be segmented to obtain a coarse segmentation result. Then, align the prompt image with the coarse segmentation result and use it as a multi-atlas.

[0030] In the multi-atlas generation subnetwork, 20 images reserved in the dataset are used as prompt images. The input image is passed through a simple segmentation network, such as a Unified Network (UNet), to generate a coarse segmentation result. The SyN algorithm is then used to register the original prompt image to the coarse segmentation result, and the corresponding registration result is used as the multi-atlas.

[0031] Step 2: Use the ViT model with SAM pre-trained weights as the image encoder to extract input image features and insert the adapter module into it for fine-tuning.

[0032] In the segmentation sub-network, the SAM architecture is used, consisting of two encoders (image encoder and prompt encoder) and one decoder (mask decoder) to extract and fuse image and multi-atlas features to automatically generate segmentation results.

[0033] The image is first tiled, splitting the two-dimensional image into multiple tiles and adding positional encodings. After passing through multiple layers of Transformer modules, image features are extracted. In the image encoder, each Transformer module is constructed in the order of a multi-head attention module, a layer normalization module, an MLP module, and a layer normalization module. After the second layer normalization, an adapter module is placed alongside the MLP module for fine-tuning. Each layer normalization uses residual connections to improve accuracy. When using the SAM model with the smallest number of parameters, the Transformer layer has 12 layers. During training, the Transformer weights of the original SAM weights are frozen, and the adapter weights are updated for fine-tuning.

[0034] In the adapter module, the input feature map is first weighted by channel, and a channel-weighted feature map is obtained through element-wise multiplication. This is achieved by performing global pooling to compress the input feature map by channel. Two fully connected layers then compress and restore the feature map, and then pass it through a sigmoid layer to obtain the channel-weighted feature map weights. After obtaining the weighted feature map, convolution and deconvolution are used for compression and restoration. After residual connection, layer normalization is performed to obtain the final output feature map.

[0035] Step 3: Use CNN as a hint encoder to extract multi-atlas features and retain the feature maps of each downsampling to help subsequent images upsample to higher resolution.

[0036] The multi-atlas generated in step 1, registered to the input image, is used as input and features are extracted using four downsampling modules. Each downsampling module contains a maximum pooling layer and a double convolution module. The double convolution module expands the receptive field to extract features without changing the size of the feature map. Each double convolution module consists of a convolution layer, layer normalization, and a GeLU layer. In the maximum pooling layer, the convolution kernel size is 2 and the stride is 2. This reduces the image width and height by half while maintaining the original number of channels. In the convolution module, the convolution kernel size is 3, the stride is 1, and the padding is 1. This increases the number of channels of the original image while maintaining the original width and height. Layer normalization is then used to stabilize the network, and GeLU layers are used for activation. After the multi-atlas input, it first passes through a convolution module to reduce the number of channels to 16 while maintaining the multi-atlas width and height. Then, four downsampling steps are performed to gradually increase the number of channels to 256 and reduce the multi-atlas width and height to 1 / 16 of the original. The features extracted by each downsampling module will be retained to help restore the predicted low-resolution segmentation results to the original size.

[0037] Step 4: Use the decoding module of the mask decoder to fuse the features extracted by the two encoders, obtain the feature vector through a multi-layer MLP, and finally obtain the prediction result through dot multiplication.

[0038] The mask decoder architecture consists of a decoding module, four upsampling modules, a crisscross attention layer, a three-layer MLP, and a dot product operation. In the mask decoder, the decoding module uses the decoding module of the SAM decoder, consisting of two layers: a self-attention layer, a token-to-image crisscross attention layer, an MLP layer, and an image-to-token crisscross attention layer. The decoding module fuses features from the image encoder with those from the hint encoder to produce updated image features and hint features. In the upsampling module, after deconvolution, layer normalization, and a GeLU layer, the updated image features are concatenated with the corresponding downsampled features retained from the hint encoder. The updated features are then passed through two convolutional modules to achieve a higher resolution. The deconvolution layer uses a kernel size of 2 and a stride of 2. While the number of channels remains unchanged after deconvolution, the image's width and height are doubled. The hint features generated by the decoding module are further updated in the token-to-image crisscross attention layer. After passing through multiple MLP layers, a feature vector is generated. Finally, the feature vector generated by the MLP layer is dot-multiplied with the high-resolution image features generated by the upsampling module to produce the final segmentation result.

[0039] The validation was conducted on a spinal dataset and a neural dataset. The spinal dataset had four categories: background, nerve, intervertebral disc, and cone, with each 2D image size of 512*512; the neural dataset had two categories: background and nerve, with each 2D image size of 160*320. The present invention was compared with the point and box generation of the SAM model. Both points and boxes were generated in advance based on real labels to simulate manual selection of points and boxes: the real labels were divided into a foreground point set and a background point set, and some points were randomly selected from each category of an image as prompt points. Within the foreground point set, a box that perfectly enclosed all foreground points was selected as a narrow bounding box, and a box that was randomly expanded by a certain number of voxels in both width and height was selected as a prompting box. As shown in Table 1, the segmentation accuracy of the present invention was better than that of the SAM using 1, 3, and 5 points, as well as the accuracy of using a tight bounding box, a 10-voxel expanded bounding box, and a 20-voxel expanded bounding box (i.e., DICE = 0.9165, MIOU = 0.8491, and ASD = 1.7607).

[0040] Table 1 Segmentation results on the spine dataset

[0041]

[0042] As shown in Table 2, the segmentation accuracy of the present invention's SAM using 1, 3, and 5 points is also better than the accuracy of using a tight bounding box, a 10-voxel enlarged bounding box, and a 20-voxel enlarged bounding box (i.e., DICE = 0.8525, MIOU = 0.7484, and ASD = 0.7596). From these results, it can be seen that the present invention integrates prior knowledge from multiple atlases and thus achieves high accuracy that exceeds the best performance of SAM. Furthermore, this method does not require any manual annotation of points or boxes as prompts, making the present invention fully automated and highly efficient.

[0043] Table 2 Segmentation results on neural datasets

[0044]

[0045] This method, based on multi-atlas and the SAM model, achieves precise spinal segmentation. It uses two encoders to extract features from the image to be segmented and the multi-atlas, and a mask decoder to fuse these features, ultimately producing high-resolution segmentation results. The use of multi-atlas prior knowledge and the SAM framework ensures high accuracy. Furthermore, this method requires no manual annotation, making it highly efficient.

[0046] The above is only a preferred embodiment of the present invention. It should be pointed out that for ordinary technicians in this technical field, several improvements and modifications can be made without departing from the principles of the present invention. These improvements and modifications should also be regarded as the scope of protection of the present invention.

Claims

1. A spine image segmentation method based on multi-atlas and SAM model, characterized in that: The following steps are involved: (1) Use a simple network to perform preliminary segmentation on the image to be segmented to obtain a rough segmentation result, then align the prompt image with the rough segmentation result and use it as a multi-atlas; (2) Use the ViT model with SAM pre-trained weights as the image encoder to extract input image features and insert the adapter into it for fine-tuning; (3) Use CNN as a prompt encoder to extract multi-atlas features and retain the feature maps of each downsampling to help subsequent images upsample to higher resolution; (4) The decoding module of the mask decoder is used to fuse the features extracted by the two encoders, and the feature vector is obtained through a multi-layer MLP. Finally, the prediction result is obtained through dot multiplication.

2. The spine image segmentation method based on multi-atlas and SAM model according to claim 1, characterized in that: The step (1) is implemented as follows: Use randomly reserved pictures in the dataset as prompt images, and generate a rough segmentation result after the input image passes through the UNet segmentation network; then use the SyN algorithm to align the label of the prompt image to the rough segmentation result, and use the corresponding registered result as a multi-atlas.

3. The spine image segmentation method based on multi-atlas and SAM model according to claim 1, characterized in that: The implementation process of using the ViT model with SAM pre-trained weights as an image encoder to extract input image features in step (2) is as follows: The two-dimensional image is divided into multiple image blocks and position encoding is added. The image features are extracted after passing through multiple layers of Transformer modules. Each Transformer module is spliced ​​in the order of a multi-head attention module, a layer normalization, an MLP module, and a layer normalization. After the second layer normalization, an adapter module is arranged in parallel with the MLP for fine-tuning. Each layer normalization uses residual connection to improve accuracy.

4. The spine image segmentation method based on multi-atlas and SAM model according to claim 1, characterized in that: The process of inserting the adapter for fine-tuning in step (2) is as follows: After global pooling, the input feature map is compressed by channel, and then two fully connected layers compress and then restore the feature map. The sigmoid layer is used to obtain the weight of the feature map by channel. After obtaining the weighted feature map, it is compressed and restored through convolution and deconvolution. After residual connection, the final output feature map is obtained through layer normalization.

5. The spine image segmentation method based on multi-atlas and SAM model according to claim 1, characterized in that: The implementation process of step (3) is as follows: Use multi-atlas as input and use four downsampling modules for feature extraction; each downsampling module contains a maximum pooling layer and a double convolution module; each double convolution module consists of a convolution layer, layer normalization, and GeLU layer; in the convolution module, the convolution layer increases the number of channels of the original image while keeping the image length and width unchanged, then uses layer normalization to stabilize the network, and uses the GeLU layer for activation; after the multi-atlas input, it first passes through a convolution module to increase the number of channels to 16 while keeping the length and width of the multi-atlas unchanged, and then after 4 downsamplings, the number of channels is gradually increased to 256 and the length and width of the multi-atlas are reduced to 1 / 16 of the original; the features extracted by each downsampling module will be retained to help restore the predicted low-resolution segmentation results to the original size.

6. The spine image segmentation method based on multi-atlas and SAM model according to claim 5, characterized in that: In the maximum pooling layer, the convolution kernel size is 2 and the stride is 2, which reduces the length and width of the image by half while keeping the original number of channels unchanged.

7. The spine image segmentation method based on multi-atlas and SAM model according to claim 5, characterized in that: The convolution kernel size of the convolution layer is 3, the stride is 1, and the padding is 1.

8. The spine image segmentation method based on multi-atlas and SAM model according to claim 1, characterized in that: The mask decoder of step (4) includes a decoding module, four upsampling modules, a cross attention layer, a multi-layer MLP and a dot product operation; the decoding module fuses the features from the image encoder and the features from the prompt encoder to generate updated image features and prompt features; in the upsampling module, after the deconvolution layer, layer normalization and GeLU layer, the updated image features are concatenated with the corresponding down-sampled features retained by the prompt encoder, and then pass through two layers of convolution modules to obtain higher resolution results; in the token-to-image cross attention layer, the prompt features generated by the decoding module are updated again, and then pass through the multi-layer MLP to generate a feature vector; finally, the feature vector generated by the MLP layer is dot-multiplied with the high-resolution image features generated by the upsampling module to obtain the final segmentation result.

9. The method for segmenting spinal images based on multi-atlas and SAM model according to claim 8, characterized in that: The decoding module uses the decoding module of the SAM decoder, which has two layers, and each layer has the structure of: self-attention layer, token-to-image cross-attention layer, MLP layer, and image-to-token cross-attention layer.

10. The spine image segmentation method based on multi-atlas and SAM model according to claim 8, characterized in that: The convolution kernel size of the deconvolution layer is 2, and the step size is 2. After deconvolution, the number of channels remains unchanged, and the length and width of the image become twice as large as before.

Images