A spine image segmentation method based on multi-atlas and SAM model

Abstract

The application discloses a kind of spinal column image segmentation methods based on multi-atlas and SAM model, using simple network to carry out preliminary segmentation to the image to be segmented, then prompt image is registered with rough segmentation result, and it is used as multi-atlas;Using the ViT model of SAM pre-training weight as image encoder extracts input image feature, and adapter is inserted therein to fine-tune;Using CNN as prompt encoder extracts multi-atlas feature, and retains the feature map of each down-sampling;Using the decoding module of mask decoder fuses the features extracted by two encoders, and obtains feature vector through multi-layer MLP, finally, the prediction result is obtained by dot product.The application automatically generates segmentation result without manual intervention through a multi-atlas generation subnetwork and a segmentation subnetwork, improves the inefficient operation of SAM model in the field of medical segmentation, which needs manual labeling points and frames, and combines the prior knowledge of multi-atlas to make the segmentation result more accurate.

Description

TECHNICAL FIELD

[0001] The application belongs to the technical field of medical images, and particularly relates to a spine image segmentation method based on a multi-atlas and SAM model. BACKGROUND

[0002] The spine is an indispensable important result in the human body, which provides important support to the human body. Accurate spine segmentation results are also important for the treatment of spinal diseases, which can help doctors better understand the spinal anatomic morphology of patients so as to better diagnose and treat the patients. However, accurate manual annotation of the spine is a very time-consuming and inefficient work. And due to the different personal experiences of the annotators, manual annotation is also prone to errors. For automatic annotation methods based on deep learning, it is also difficult. Due to the complex anatomical structure of the spine region, noise of medical images, low contrast of neural regions in CT and other reasons, the segmentation result of the spine region cannot be very accurate.

[0003] Multi-atlas segmentation has the characteristics of stability and accuracy. Traditional multi-atlas segmentation uses multiple images and their labels as templates, and then performs registration on the to-be-segmented image and the labeled template image through SyN or Demons algorithm to obtain multiple registered labeled images. Then the multiple labeled results are fused through weighted voting or other methods to generate an image, and then some post-processing is performed to obtain the final segmentation result of the to-be-segmented image. With the development of deep learning, multi-atlas segmentation has also been combined with deep learning to achieve the combination of advantages.

[0004] The SAM model provides a basic model in the field of image segmentation, and its high accuracy and convenience provide a new idea for segmentation. It uses the MAE pre-trained ViT model as an image feature extractor. Because the ViT model has a size, the SAM also provides three models with different parameter quantities. The larger the parameter quantity, the larger the model volume, and the more accurate the segmentation result. In addition, it provides three ways of points, boxes and texts as prompts for the to-be-segmented region of the image, and the operator only needs to click or draw a box or input text to accurately complete the segmentation of the two-dimensional image. However, the training data of the SAM model is natural images, while medical images are obtained by X-ray, nuclear magnetic resonance and other methods, so there are great differences in pixel intensity, texture and other properties between medical images and natural images. Therefore, the SAM model does not perform well in medical image segmentation. In addition, for CT, MRI and other three-dimensional images, which are stacked by hundreds of two-dimensional images, the segmentation of the entire three-dimensional image needs to be performed on each two-dimensional image. If the operator manually clicks, draws a box and other operations on each two-dimensional image in the three-dimensional image, it is still a complex and energy-consuming task.

[0005] Therefore, how to combine multi-atlas segmentation and SAM model, improve the shortcomings of SAM, and use the high accuracy of SAM and multi-atlas to complete the accurate segmentation of the spine region is the key task of the present application. SUMMARY

[0006] The application proposes a spine image segmentation method based on multi-atlas and SAM model, which improves the low efficiency of manual labeling points and frames in the medical segmentation field of SAM model, and combines the prior knowledge generated by multi-atlas to make the segmentation result more accurate.

[0007] SUMMARY: The spine image segmentation method based on multi-atlas and SAM model comprises the following steps:

[0008] (1) using a simple network to preliminarily segment the image to be segmented to obtain a coarse segmentation result, then registering the prompt image with the coarse segmentation result, and taking it as a multi-atlas;

[0009] (2) using a ViT model with SAM pre-training weight as an image encoder to extract input image features, and inserting an adapter for fine-tuning;

[0010] (3) using CNN as a prompt encoder to extract multi-atlas features, and retaining each down-sampled feature map to help subsequent image up-sampling to a higher resolution;

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

[0012] Further, the step (1) is implemented as follows:

[0013] Use the pictures randomly reserved in the data set as prompt images, and input the image into the UNet segmentation network to generate a coarse segmentation result; then use the SyN algorithm to register the label of the prompt image to the coarse segmentation result, and take the corresponding registered result as a multi-atlas.

[0014] Further, the step (2) is implemented as follows:

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

[0016] Further, the step (2) inserts the adapter to fine-tune the implementation process as follows:

[0017] After global pooling, the input feature map is compressed by channel, and then two fully connected layers are used to compress and recover the feature map, and then the sigmoid layer is used to obtain the weight of the feature map by channel weighting. After obtaining the weighted feature map, compression and recovery are performed through convolution and deconvolution, and the final output feature map is obtained through layer normalization after residual connection.

[0018] Further, the step (3) is implemented as follows:

[0019] Use multiple atlases as input and use four down-sampling modules for feature extraction; each down-sampling module contains a max-pooling layer and a double convolution module; each double convolution module is composed of a convolution layer, a layer normalization, and a GeLU layer; in the convolution module, the convolution layer increases the channel number of the original image while keeping the image length and width unchanged, then the layer normalization is used to stabilize the network, and the GeLU layer is used for activation; after the multi-atlas input, first pass through a convolution module to change the channel number to 16 and keep the multi-atlas length and width unchanged, then pass through 4 times down-sampling to gradually increase the channel number to 256 and reduce the multi-atlas length and width to 1 / 16 of the original size; the features extracted by each down-sampling module will be retained to help the low-resolution segmentation result recovered to the original size.

[0020] Further, in the max-pooling layer, the convolution kernel size is 2, the step is 2, and the image length and width are reduced by half while keeping the original channel number unchanged.

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

[0022] Further, step (4) the mask decoder includes a decoding module, four up-sampling modules, a cross-attention layer, a multi-layer MLP, and a point multiplication operation; the decoding module fuses features from the image encoder and features from the prompt encoder to produce updated image features and prompt features; in the up-sampling module, after the deconvolution layer, layer normalization, and GeLU layer, the updated image features are spliced with the corresponding down-sampling features retained by the prompt encoder, and then two convolution modules are used 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 a feature vector is generated after passing through the multi-layer MLP; finally, the feature vector generated by the MLP layer is multiplied with the high-resolution image features generated by the up-sampling module to obtain the final segmentation result.

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

[0024] Further, the convolution kernel size of the deconvolution layer is 2, the step is 2, the channel number remains unchanged after deconvolution, and the length and width of the image become twice the original.

[0025] Advantages: Compared with the prior art, the advantages of the present application: the present application uses anatomical prior knowledge contained in multiple atlases to make the segmentation result more accurate; the present application improves the low-efficiency operation of manually labeling points and frames in the SAM model in the medical segmentation field, adopts a fine-tuned original SAM model image encoder and automatically generates a segmentation result through a prompt encoder and a mask decoder; the present application achieves high precision of SAM, and is fully automatic and does not require manual labeling. BRIEF DESCRIPTION OF DRAWINGS

[0026] Figure 1 is a schematic diagram of a spine image segmentation network structure proposed by the present application; DETAILED DESCRIPTION

[0027] The present application will be further described below in conjunction with the drawings.

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

[0029] Step 1: using a simple network to preliminarily segment the to-be-segmented image to obtain a coarse segmentation result, then registering the prompt image with the coarse segmentation result, and taking it as a multiple atlas.

[0030] In the multiple atlas generation subnetwork, 20 pictures reserved in the data set are used as prompt pictures. The input image generates a coarse segmentation result after passing through a simple segmentation network, for example, UNet. Then the original prompt image is registered to the coarse segmentation result using the SyN algorithm, and the corresponding registered result is taken as a multiple atlas.

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

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

[0033] The picture is first divided into multiple picture blocks by the picture block operation, and the position encoding is added. After passing through multiple Transformer modules, the image features are extracted. In the image encoder, 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 to fine-tune. Each layer normalization uses a residual connection to improve accuracy. When using the SAM model with the smallest parameter amount, the level of the Transformer is 12. During training, the Transformer weights of the original SAM 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 the channel-weighted feature map is obtained by element-wise multiplication. The implementation is as follows: the input feature map is compressed by channel through global pooling, then two fully connected layers are used to compress and restore the feature map, and then a sigmoid layer is used to obtain the weight of the feature map by channel. After obtaining the weighted feature map, compression and restoration are performed through convolution and deconvolution, and the final output feature map is obtained through layer normalization after residual connection.

[0035] Step 3: Use CNN as the prompt encoder to extract multi-atlas features and keep the feature map of each down-sampling to help subsequent image up-sampling to higher resolution.

[0036] The multi-atlas generated in step 1 is used as input, and feature extraction is performed using 4 down-sampling modules. Each down-sampling module contains a max-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 is composed of a convolution layer, a layer normalization, and a GeLU layer. In the max-pooling layer, the convolution kernel size is 2, and the step is 2, which reduces the image length and width by half while keeping the original number of channels unchanged. In the convolution module, the convolution kernel size of the convolution layer is 3, the step is 1, and the padding is 1, which 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 change the number of channels to 16 and keep the multi-atlas length and width unchanged, then passes through 4 down-sampling modules to gradually increase the number of channels to 256 and reduce the multi-atlas length and width to 1 / 16 of the original size. The features extracted by each down-sampling module will be retained to help restore the low-resolution segmentation result to the original size.

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

[0038] The architecture of the mask decoder is a decoding module, four up-sampling modules, a cross-attention layer, a three-layer MLP, and a point multiplication operation. In the mask decoder, the decoding module uses the decoding module of the SAM decoder, which has two layers, and each layer has the following structure: self-attention layer, token-to-image cross-attention layer, MLP layer, and image-to-token cross-attention layer. 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 up-sampling module, after the deconvolution layer, layer normalization, and GeLU layer, the updated image features are spliced with the corresponding down-sampling features retained in the prompt encoder, and then two convolution modules are used to obtain higher resolution results. The convolution kernel size of the deconvolution layer is 2, and the step is 2, after deconvolution, the number of channels remains unchanged, and the length and width of the image become twice the original. In the token-to-image cross-attention layer, the prompt features generated by the decoding module are updated again, then a multi-layer MLP is used to generate a feature vector, and finally the feature vector generated by the MLP layer is multiplied with the high-resolution image features generated by the up-sampling module to obtain the final segmentation result.

[0039] Validation is performed on the spine dataset and the nerve dataset, wherein the spine dataset is four categories: background, nerve, intervertebral disc, cone, and each 2D picture size is 512*512; the nerve dataset is two categories: background, nerve, and each 2D picture size is 160*320. The present application is compared with the point and box of the SAM model; wherein the generation of the point and the box is generated in advance based on the real label to simulate the manual selection of the point and the box: the real label is divided into a foreground point set and a background point set, and some points are randomly selected in each category of a picture as prompt points; in the foreground point set, a box that can just frame all the foreground points is selected as a narrow box, and a box that is randomly expanded in terms of width and height by a certain voxel is selected as a prompt box. As shown in Table 1, the segmentation accuracy of the present application is better than the accuracy of SAM using 1, 3 and 5 points, and is also better than the accuracy of using a tight box, a 10-voxel expansion and a 20-voxel expansion box (i.e. DICE = 0.9165, MIOU = 0.8491, 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 application is better than the accuracy of SAM using 1, 3 and 5 points, and is also better than the accuracy of using a tight box, a 10-voxel expansion and a 20-voxel expansion box (i.e. DICE = 0.8525, MIOU = 0.7484, ASD = 0.7596). From the above results, it can be seen that the present application fuses prior knowledge from multiple atlases and thus achieves high accuracy exceeding the best performance of SAM. At the same time, the present method does not need any manually labeled points or boxes as prompts, so that the present application is fully automated and very efficient.

[0043] Table 2 Segmentation results on the nerve dataset

[0044]

[0045] The present application realizes accurate segmentation of the spine based on multiple atlases and a SAM model. The present application uses two encoders to extract features of a to-be-segmented image and multiple atlases, and fuses the features through a mask decoder to finally generate a high-resolution segmentation result. Due to the use of multiple atlas prior knowledge and the SAM framework, the present method has high accuracy. At the same time, the present method also does not need any manual labeling and is very efficient.

[0046] The above only describes the preferred embodiments of the present application, and it should be noted that for ordinary skilled persons in the art, some improvements and refinements can be made without departing from the principles of the present application, and these improvements and refinements should also be regarded as the protection scope of the present application.

Claims

1. A method for spinal image segmentation based on multi-atlas and SAM model, characterized in that, The method comprises the following steps: (1) using a simple network to preliminarily segment the image to be segmented to obtain a coarse segmentation result, then registering the prompt image with the coarse segmentation result, and taking the registration result as a multi-atlas; (2) using a ViT model with SAM pre-training weights as an image encoder to extract input image features, and inserting an adapter module in the ViT model to fine-tune the ViT model; (3) using a CNN as a prompt encoder to extract multi-atlas features, and retaining each down-sampled feature map to help subsequent image up-sampling to a higher resolution; (4) using a decoding module of a mask decoder to fuse the features extracted by the two encoders, and obtaining a feature vector through a multi-layer MLP, and finally obtaining a prediction result through point multiplication; The mask decoder of step (4) comprises a decoding module, four up-sampling modules, a cross-attention layer, a multi-layer MLP and a point multiplication 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 up-sampling module, after the inverse convolution layer, the layer normalization and the GeLU layer, the updated image features are spliced with the corresponding down-sampled features retained by the prompt encoder, and then two convolution modules are used to obtain a higher resolution result; in the token-to-image cross-attention layer, the prompt features generated by the decoding module are updated again, and then a feature vector is generated after the multi-layer MLP; finally, the feature vector generated by the MLP layer is multiplied with the high-resolution image features generated by the up-sampling module to obtain the final segmentation result.

2. The spinal image segmentation method based on multi-atlas and SAM model according to claim 1, characterized in that, The step (1) is implemented as follows: The pictures randomly reserved in the data set are used as prompt images, and the input images are segmented by a UNet segmentation network to generate coarse segmentation results; then, the label of the prompt image is registered to the coarse segmentation result by using a SyN algorithm, and the registration result is taken as a multi-atlas.

3. The spinal image segmentation method based on multi-atlas and SAM model according to claim 1, characterized in that, The step (2) is implemented as follows: The two-dimensional picture is divided into multiple picture blocks and position encoding is added, and the image features are extracted after passing through multiple 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; an adapter module is arranged in parallel with the MLP after the second layer normalization to fine-tune; each layer normalization uses a residual connection to improve the accuracy.

4. The spinal image segmentation method based on multi-atlas and SAM model according to claim 1, characterized in that, The step (2) is implemented as follows: After global pooling, the input feature map is compressed by channel, then two fully connected layers are used to compress and restore the feature map, and a sigmoid layer is used to obtain the weight of the feature map by channel weighting; after obtaining the weighted feature map, convolution and deconvolution are used for compression and restoration, and a residual connection is used to obtain the final output feature map through layer normalization.

5. The spinal image segmentation method based on multi-atlas and SAM model according to claim 1, characterized in that, The step (3) is implemented as follows: Use multiple atlas as input, and use four down-sampling modules for feature extraction; each down-sampling module contains a max-pooling layer and a double convolution module; each double convolution module is composed of a convolution layer, a layer normalization, and a GeLU layer; in the convolution module, the convolution layer increases the channel number of the original image while keeping the image length and width unchanged, then uses the layer normalization to stabilize the network, and uses the GeLU layer for activation; after the multiple atlas input, first pass through a convolution module to change the channel number to 16 and keep the multiple atlas length and width unchanged, then pass through 4 times down-sampling to gradually increase the channel number to 256 and reduce the multiple atlas length and width to 1 / 16 of the original; the features extracted by each down-sampling module will be retained to help the low-resolution segmentation result recovered to the original size.

6. The spinal image segmentation method based on multi-atlas and SAM model according to claim 5, characterized in that, In the max-pooling layer, the convolution kernel size is 2, the step is 2, the image length and width are reduced by half while keeping the original channel number unchanged.

7. The spinal 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 step is 1, and the padding is 1.

8. The spinal image segmentation method based on multi-atlas and SAM model according to claim 1, characterized in that, The decoding module uses the decoding module of the SAM decoder, a total of two layers, and each layer structure is: a self-attention layer, a token-to-image cross-attention layer, an MLP layer, and an image-to-token cross-attention layer.

9. The spinal image segmentation method based on multi-atlas and SAM model according to claim 1, characterized in that, The convolution kernel size of the deconvolution layer is 2, the step is 2, and after deconvolution, the channel number remains unchanged, and the image length and width become twice the original.

Images