A neural network-based image segmentation method

Abstract

The application provides a kind of neural network-based image segmentation method, it is related to medical image processing technical field, the method comprises: constructing medical image segmentation model;Medical image segmentation model includes: encoder and decoder;Decoder includes: first mask decoder and second mask decoder;Through loss function and self-attention mechanism, encoder is trained, and initial encoder is obtained;Medical space transpose adapter is embedded in initial encoder, and final encoder is obtained, and medical image is encoded by final encoder, and encoding result is obtained;Medical space transpose adapter is embedded in second mask decoder, and final decoder is obtained;Through final decoder, the decoding processing of encoding result is obtained, and image segmentation result is obtained.The method cascades connection between medical space transpose adapter and SAM decoder, and combines loss calculation, improves the segmentation accuracy of model, so that the segmentation result is more accurate and reliable.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of medical image processing, in particular to an image segmentation method based on a neural network. BACKGROUND

[0002] Segment Anything Model (SAM) is a large image segmentation model launched by Meta AI, which is trained based on 11 million images and more than 1 billion masks, and is currently the largest image segmentation model based on the Transformer architecture. SAM faces significant challenges in the field of medical image segmentation. The reason is that medical images are usually accompanied by obvious noise, artifacts and low-contrast regions. These characteristics greatly increase the difficulty of automatically identifying subtle lesions such as tumors, blood vessels or organ boundaries. In addition, the structures and lesions in medical images exhibit great heterogeneity, from small lesions of a few millimeters to large organs of several centimeters, and the changes in shape and size require the segmentation model to have high adaptability and diverse processing capabilities.

[0003] However, as a model designed for general image segmentation, SAM is not optimized for the unique difficulties of medical images during design and training, which leads to the problem that it may not be accurate enough to capture key details when dealing with the complexity of medical images. Further, medical image datasets often have the problem of class imbalance, especially for rare lesions or rare diseases, where labeled data is extremely scarce. This lack of data limits the ability of SAM to learn all potential lesion features, exacerbating its performance bottleneck in medical segmentation tasks. Therefore, it is necessary to design an image segmentation method based on a neural network. SUMMARY

[0004] The purpose of the present application is to provide an image segmentation method based on a neural network, which is connected in cascade with a medical space transpose adapter and a SAM decoder, and combines a loss calculation method matched with the medical space transpose adapter, to improve the segmentation accuracy of the model and simplify the decoder parameters, so that the segmentation result is more accurate and reliable.

[0005] To achieve the above purpose, the present application provides the following scheme:

[0006] An image segmentation method based on a neural network, comprising the following steps:

[0007] constructing a medical image segmentation model; the medical image segmentation model comprises an encoder and a decoder; the encoder comprises a residual connection part, a multi-head attention mechanism part and a linear normalization layer part; the decoder comprises a first mask decoder and a second mask decoder;

[0008] The initial encoder is obtained by training the encoder through a loss function and a self-attention mechanism;

[0009] The medical space transpose adapter is embedded in the initial encoder to obtain a final encoder, and the medical image is encoded through the final encoder to obtain an encoding result;

[0010] The medical space transpose adapter is embedded in the second mask decoder to obtain a final decoder; the final decoder comprises a first mask decoder and an optimized mask decoder.

[0011] The encoding result is decoded through the final decoder to obtain an image segmentation result.

[0012] Optionally, the medical image segmentation model is a segmentation-all model.

[0013] Optionally, the expression of the loss function is: wherein, L total is the loss function, is the loss of the i-th transformer layer, is the loss of the first mask decoder, is the loss of the second mask decoder, and M is the number of transformer layers.

[0014] The loss of the i-th transformer layer is obtained by a Dice loss function; the loss of the first mask decoder is obtained by a BCE loss function, and the loss of the second mask decoder is obtained by an iou loss function.

[0015] Optionally, the expression of the self-attention mechanism is: wherein, Attention is self-attention, Q is query, K is key, V is value, d k is the feature dimension, SOFTMAX is the SOFTMAX function, K T is the transpose of the key.

[0016] Optionally, the medical space transpose adapter comprises a residual adapter, a dimension fusion adapter and a normalization adapter; the residual adapter, the dimension fusion adapter and the normalization adapter are all composed of a LoRA module; the expression of the forward propagation process built-in the LoRA module is: W0+ΔW=W0+BA; wherein, W0 is a pre-training weight matrix, ΔW is a low-rank decomposition, B is an upper projection parameter matrix, and A is a lower projection parameter matrix.

[0017] Optionally, a medical spatial transpose adapter is embedded in the initial encoder to obtain the final encoder, and the medical image is encoded using the final encoder to obtain the encoding result. Specifically, after embedding the dimension fusion adapter into the multi-head attention mechanism, after embedding the normalization adapter into the linear normalization layer, the residual adapter is used as the residual connection part to obtain the final encoder, and the medical image is encoded using the final encoder to obtain the encoding result.

[0018] Optionally, both the first mask decoder and the second mask decoder include: a Transformer decoder, a pixel decoder, and a convolutional layer; the pixel decoders are connected by skip connections.

[0019] Optionally, the encoded result is decoded by the final decoder to obtain the image segmentation result, including:

[0020] A priori probability mask is generated using the first mask decoder;

[0021] The prior probability mask is decoded using a Transformer decoder and a pixel decoder respectively to obtain the first result and the second result.

[0022] The image segmentation result is obtained by merging the first and second results through an optimized mask decoder.

[0023] This invention discloses the following technical effects: The image segmentation method based on neural networks provided by this invention includes: constructing a medical image segmentation model; the medical image segmentation model includes an encoder and a decoder; the encoder includes a residual connection part, a multi-head attention mechanism part, and a linear normalization layer part; the decoder includes a first mask decoder and a second mask decoder; the encoder is trained using a loss function and a self-attention mechanism to obtain an initial encoder; a medical spatial transpose adapter is embedded in the initial encoder to obtain a final encoder, and the medical image is encoded using the final encoder to obtain an encoding result; the medical spatial transpose adapter is embedded in the second mask decoder to obtain a final decoder; the encoding result is decoded using the final decoder to obtain an image segmentation result. This method cascades the medical spatial transpose adapter and the SAM decoder, and combines loss calculation, improving the segmentation accuracy of the model and making the segmentation results more accurate and reliable. Attached Figure Description

[0024] To more clearly illustrate the technical solutions in the embodiments of the present invention or the prior art, the drawings used in the embodiments will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0025] Figure 1 This is a flowchart of the image segmentation method of the present invention;

[0026] Figure 2 This is a structural diagram of the medical space transposition adapter of the present invention;

[0027] Figure 3 This is a diagram of the final decoder structure of the present invention. Detailed Implementation

[0028] The technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.

[0029] To make the above-mentioned objects, features and advantages of the present invention more apparent and understandable, the present invention will be further described in detail below with reference to the accompanying drawings and specific embodiments.

[0030] like Figure 1 As shown, this invention provides an image segmentation method based on a neural network, comprising the following steps:

[0031] Step 100: Construct a medical image segmentation model;

[0032] Specifically, the medical image segmentation model in this embodiment adopts the Segment All Model (SAM), which includes an encoder and a decoder. The encoder includes a residual connection part, a multi-head attention mechanism part, and a linear normalization layer part. The decoder includes a first mask decoder and a second mask decoder.

[0033] Step 200: Train the encoder using a loss function and a self-attention mechanism to obtain the initial encoder;

[0034] Specifically, the expression for the loss function is as follows:

[0035]

[0036] Among them, L total For loss function, The loss of the i-th transformer layer is... For the loss of the first mask decoder, Let M be the loss of the second mask decoder, and M be the number of transformer layers. The loss of the i-th transformer layer is calculated using the Dice loss function; the loss of the first mask decoder is calculated using the BCE loss function; and the loss of the second mask decoder is calculated using the Iou loss function.

[0037] Specifically, during encoder training, the input to the encoder is the original medical image, such as a CT scan image. First, the image X is expanded into multiple segment-length sequences, which are then fed in batches into the Transformer module of the SAM encoder. Through the self-attention mechanism within the Transformer module, the contextual relationships between pixels in the image are captured, and the information extraction method is remembered by the feedforward neural network within the Transformer module. Finally, after passing through multiple Transformer modules, fine-grained medical image features are obtained. The expression for the self-attention mechanism is as follows:

[0038]

[0039] Q = W Q X, K = W K X, V = W V X;

[0040] Where Attention is self-attention, Q is the query, K is the key, V is the value, and d is the input key. k The feature dimension is determined by the sequence length, SOFTMAX is the SOFTMAX function, and K T W is the transpose of the key. Q W K and W V These are the computation matrices for Q, K, and V, used to calculate the specific values ​​of Q, K, and V. Furthermore, the original parameters of the SAM encoder are frozen during encoder training.

[0041] It should be noted that the inputs for modules at different levels can be obtained through multiple levels of Transformer operations within the SAM encoder. This effectively converts information from the input image at different levels into embedding vectors, which can then be used in subsequent processing steps of the model, such as feature fusion, classification, or other visual tasks.

[0042] Step 300: Embed a medical space transpose adapter in the initial encoder to obtain the final encoder, and encode the medical image through the final encoder to obtain the encoding result;

[0043] like Figure 2As shown, the medical spatial transpose adapter includes a residual adapter, a dimension fusion adapter, and a normalization adapter, which respectively act on the residual connection part, the multi-head attention mechanism part, and the linear normalization layer part in each transformer layer of the SAM encoder. The residual adapter, dimension fusion adapter, and normalization adapter are all designed based on the LoRA module as the basic adapter module. The English explanations in the figure are: NormalizedAdapter, MultilayerPerceptron, LayerNormalization, DimensionFusionAdapter, Feature fusion, Depth feature, Spatial feature, Multi-HeadAttention, Depth Part, Spatial Part, Upperprojection, Lowerprojection, and ResidualAdapter.

[0044] Specifically, the LoRA module consists of a lower projection matrix, a ReLU function, and an upper projection parameter matrix in that order. First, the lower projection parameter matrix compresses the input sequence into a sequence of length r. Then, the sequence is expanded using the ReLU function activation and the upper projection parameter matrix. This process significantly compresses the number of parameters during training, thus enabling efficient tuning of models with large parameters. The expression for the forward propagation process of the LoRA module during training is as follows:

[0045] W0+ΔW=W0+BA;

[0046] Where W0 is the pre-trained weight matrix, ΔW is the low-rank decomposition, B is the upper projection parameter matrix, and A is the lower projection parameter matrix.

[0047] Furthermore, the dimension fusion adapter is embedded after the multi-head attention mechanism and before the linear normalization layer to perform modal fusion between the spatial feature transformer and the intermediate data of the deep feature transformer; the normalization adapter is embedded after the linear normalization layer to perform linear normalization adjustment; the residual adapter is placed in the bypass as the residual connection part to fuse the features before and after embedding the residual adapter to obtain the final encoder, and the medical image is encoded through the final encoder to obtain the encoding result.

[0048] Furthermore, the medical spatial transpose adapter divides the attention mechanism of the SAM encoder into spatial and depth branches. In a 3D sample of depth D, the spatial branch processes an input of size D×N×L, where N is the number of embeddings and L is the embedding length. In this way, the spatial branch can capture spatial correlations. In the depth branch, this embodiment transposes the input matrix to N×D×L and performs the same multi-head attention operation, where interactions occur between D×L values, thereby learning depth correlations. Finally, the result of the depth branch is transposed back to its original shape and added to the result of the spatial branch to fuse depth information.

[0049] Step 400: Embed the medical space transpose adapter into the second mask decoder to obtain the final decoder;

[0050] Specifically, both the first mask decoder and the second mask decoder include: a Transformer decoder, a pixel decoder, and a convolutional layer; the final decoder includes: the first mask decoder and the optimized mask decoder; the optimized mask decoder is obtained by cascading the medical space transpose adapter with the second mask decoder, which enhances the segmentation accuracy while simplifying the decoder parameters.

[0051] like Figure 3 As shown, the medical space transpose adapter is embedded in the two Transformer decoders of the second mask decoder, resulting in a simplified mask decoder and improving its decoding performance. Simultaneously, the pixel decoders of the first and second mask decoders exchange embedded features through skip connections, improving model accuracy and enhancing the ability to capture local details. The English explanations in the figure are: Image Embedding, SAM Mask Decoder, Transformer Decoder, PixelDecoder, Conv Layer, Medical Space TransposeAdapter, and Skipping Connections.

[0052] Step 500: The encoded result is decoded by the final decoder to obtain the image segmentation result. Specifically, a prior probability mask is generated by the first mask decoder to guide the optimized mask decoder in a more complex decoding process. The prior probability mask is then decoded by the Transformer decoder and pixel decoder of the first mask decoder to obtain the first result and the second result. The first result and the second result are then merged by the optimized mask decoder to obtain the image segmentation result.

[0053] Furthermore, the Transformer decoder can be viewed as having a structure essentially the same as the Transformer layer in the encoder. The difference lies in the fact that the calculation of Q, K, and V values ​​in the self-attention layer of the Transformer decoder originates from the combined effect of the cue embedding vector and the image embedding vector. The cue embedding vector is often a location label of a point, used to better guide image segmentation. Specifically, the sequence for calculating Q values ​​comes from the cue embedding vector, while the sequences for calculating K and V values ​​come from the image embedding vector. In contrast, the calculation of Q, K, and V values ​​for self-attention in the Transformer layer uses only the image embedding vector for computation.

[0054] Furthermore, for the pixel decoder, it's simply a cross-attention operation, derived from the combined effect of the cue embedding vector and the image embedding vector. However, the sequence for calculating the Q value comes from the image embedding vector, while the sequences for calculating the K and V values ​​come from the cue embedding vector. After two layers of pixel decoders, the output vector undergoes two more convolutional layer operations to obtain the pixel decoder's output. After two layers of Transformer decoders, the Transformer decoder's output is directly output.

[0055] The beneficial effects of this invention are as follows:

[0056] 1) A cascaded connection was adopted between the SAM decoder and the medical space transpose adapter, which improved the segmentation accuracy while simplifying the decoder parameters;

[0057] 2) In the design of the efficient transpose adapter for medical space, the spatial dimension of natural images and the depth dimension unique to medical images are integrated, enabling the SAM model to capture multi-dimensional information simultaneously. This allows the model to better understand the complexity of medical images and enhances the model's learning ability.

[0058] 3) By combining spatial and depth dimensions, the accuracy and robustness of segmentation are significantly improved;

[0059] 4) First, freeze the parameters of the SAM model after training, and then only modify and train the parameters in the SAM architecture, which saves a lot of computing resources.

[0060] 5) By using a parallel transformer training strategy in conjunction with spatial transpose technology, the segmentation accuracy of complex medical images has been improved, and the problem of imbalance in medical image data categories has been solved.

[0061] The various embodiments in this specification are described in a progressive manner, with each embodiment focusing on the differences from other embodiments. The same or similar parts between the various embodiments can be referred to each other.

[0062] Specific examples have been used to illustrate the principles and implementation methods of this invention. The descriptions of the above embodiments are only for the purpose of helping to understand the method and core ideas of this invention. Furthermore, those skilled in the art will recognize that, based on the ideas of this invention, there will be changes in the specific implementation methods and application scope. Therefore, the content of this specification should not be construed as a limitation of this invention.

Claims

1. An image segmentation method based on neural networks, characterized in that, Includes the following steps: Constructing a medical image segmentation model; The medical image segmentation model includes an encoder and a decoder; the encoder includes a residual connection part, a multi-head attention mechanism part, and a linear normalization layer part; the decoder includes a first mask decoder and a second mask decoder; the medical image segmentation model is a segmentation model for all images. The encoder is trained using a loss function and a self-attention mechanism to obtain an initial encoder; the initial encoder is a SAM encoder. A medical spatial transpose adapter is embedded in the initial encoder to obtain the final encoder, and the medical image is encoded through the final encoder to obtain the encoding result. The medical spatial transpose adapter includes a residual adapter, a dimension fusion adapter, and a normalization adapter, which respectively act on the residual connection part, the multi-head attention mechanism part, and the linear normalization layer part in each transformer layer of the SAM encoder. The medical space transpose adapter is embedded in the second mask decoder to obtain an optimized mask decoder, thereby obtaining the final decoder; the final decoder includes: the first mask decoder and the optimized mask decoder; The encoding result is decoded by the final decoder to obtain the image segmentation result; The expression for the loss function is: ;in, For loss function, For the loss of the i-th transformer layer, The loss of the first mask decoder, denoted as the loss of the second mask decoder, and M as the number of transformer layers.

2. The image segmentation method based on neural networks according to claim 1, characterized in that, The loss of the i-th transformer layer is calculated using the Dice loss function; the loss of the first mask decoder is calculated using the BCE loss function; and the loss of the second mask decoder is calculated using the IOU loss function.

3. The image segmentation method based on neural networks according to claim 1, characterized in that, The expression for the self-attention mechanism is: ;in, For self-attention, Q is the query, K is the key, and V is the value. For feature dimension, for function, This is the transpose of the key.

4. The image segmentation method based on neural networks according to claim 1, characterized in that, The residual adapter, the dimension fusion adapter, and the normalization adapter are all composed of LoRA modules; the expression for the forward propagation process built into the LoRA module is: ;in, For the pre-trained weight matrix, It is a low-rank decomposition. B Let be the upward projection parameter matrix. A Let be the downward projection parameter matrix.

5. The image segmentation method based on neural networks according to claim 4, characterized in that, A medical spatial transpose adapter is embedded in the initial encoder to obtain a final encoder. The final encoder is then used to encode the medical image to obtain the encoding result. Specifically, the dimension fusion adapter is embedded in the multi-head attention mechanism, the normalization adapter is embedded in the linear normalization layer, and the residual adapter is used as the residual connection part to obtain the final encoder. The final encoder is then used to encode the medical image to obtain the encoding result.

6. The image segmentation method based on neural networks according to claim 1, characterized in that, Both the first mask decoder and the second mask decoder include: a Transformer decoder, a pixel decoder, and a convolutional layer; the pixel decoders are connected by skip connections.

7. The image segmentation method based on neural networks according to claim 6, characterized in that, The encoding result is decoded by the final decoder to obtain the image segmentation result, including: A priori probability mask is generated using the first mask decoder; The prior probability mask is decoded by the Transformer decoder and the pixel decoder of the first mask decoder to obtain a first result and a second result. The optimized mask decoder performs a merging operation on the first result and the second result to obtain the image segmentation result.

Images