Image segmentation method based on neural network

Abstract

The invention provides an image segmentation method based on a neural network, and relates to the technical field of medical image processing, and the method comprises the steps: constructing a medical image segmentation model; the medical image segmentation model comprises an encoder and a decoder; the decoder comprises a first mask decoder and a second mask decoder; training an encoder through a loss function and a self-attention mechanism to obtain an initial encoder; embedding a medical space transpose adapter into the initial encoder to obtain a final encoder, and performing encoding operation on the medical image through the final encoder to obtain an encoding result; embedding the medical space transposition adapter into the second mask decoder to obtain a final decoder; and decoding the coding result through a final decoder to obtain an image segmentation result. According to the method, the medical space transposition adapter and the SAM decoder are in cascade connection, loss calculation is combined, the segmentation precision of the model is improved, and the segmentation result is more accurate and reliable.

Description

Technical Field

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

[0002] SegmentAnything Model (SAM) is a large-scale image segmentation model launched by MetaAI. It is trained based on 11 million images and more than 1 billion masks. It 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 areas. 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 show great heterogeneity. From lesions as small as a few millimeters to organs as large as tens of centimeters, the unpredictable changes in morphology and scale require segmentation models to have extremely high adaptability and diverse processing capabilities.

[0003] However, as a model for general image segmentation, SAM is not optimized for the unique challenges of medical images during its design and training, which may cause it to suffer from insufficient accuracy or inability to accurately capture key details in dealing with the complexity of medical images. Furthermore, medical image datasets often have class imbalance problems, especially for a few lesions or rare diseases, where labeled data is extremely scarce. This lack of data limits SAM's ability 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 neural networks. Summary of the invention

[0004] The purpose of the present invention is to provide an image segmentation method based on a neural network, which is cascaded with a SAM decoder through a medical space transposition adapter, and combined with a loss calculation method matching the medical space transposition adapter, so as 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 object, the present invention provides the following solutions:

[0006] A neural network-based image segmentation method comprises the following steps:

[0007] Construct 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;

[0008] The encoder is trained through the loss function and self-attention mechanism to obtain the initial encoder;

[0009] A medical space transposition adapter is embedded in the initial encoder to obtain a final encoder, and the final encoder is used to perform an encoding operation on the medical image to obtain an encoding result;

[0010] Embedding the medical space transposition adapter into the second mask decoder to obtain a final decoder; the final decoder includes: a first mask decoder and an optimized mask decoder;

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

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

[0013] Optionally, the loss function is expressed as: Among them, 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, M is the number of transformer layers;

[0014] 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.

[0015] Optionally, the expression of the self-attention mechanism is: Among them, 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 transposition adapter includes: 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 into the LoRA module is: W 0 +ΔW=W 0 +BA; where W 0 is the pre-trained weight matrix, ΔW is the low-rank decomposition, B is the up-projection parameter matrix, and A is the down-projection parameter matrix.

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

[0018] Optionally, the first mask decoder and the second mask decoder both include: a Transformer decoder, a pixel decoder and a convolutional layer; and a jump connection is used between the pixel decoders.

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

[0020] generating a priori probability mask by a first mask decoder;

[0021] The prior probability mask is decoded by a Transformer decoder and a pixel decoder to obtain a first result and a second result;

[0022] The first result and the second result are combined by optimizing the mask decoder to obtain an image segmentation result.

[0023] The present invention discloses the following technical effects: the neural network-based image segmentation method provided by the present invention comprises: 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; the encoder is trained by a loss function and a self-attention mechanism to obtain an initial encoder; a medical space transposition adapter is embedded in the initial encoder to obtain a final encoder, and the final encoder is used to encode the medical image to obtain an encoding result; the medical space transposition adapter is embedded in the second mask decoder to obtain a final decoder; the encoding result is decoded by the final decoder to obtain an image segmentation result. The method cascades the medical space transposition adapter with the SAM decoder, and combines the loss calculation to improve the segmentation accuracy of the model, so that the segmentation result is more accurate and reliable. BRIEF DESCRIPTION OF THE DRAWINGS

[0024] In order to more clearly illustrate the embodiments of the present invention or the technical solutions in the prior art, the drawings required for use in the embodiments will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying creative labor.

[0025] Figure 1 It is a flow chart of the image segmentation method of the present invention;

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

[0027] Figure 3 This is the final decoder structure diagram of the present invention. DETAILED DESCRIPTION

[0028] The following will be combined with the drawings in the embodiments of the present invention to clearly and completely describe the technical solutions in the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention.

[0029] In order to make the above-mentioned objects, features and advantages of the present invention more obvious and understandable, the present invention is further described in detail below with reference to the accompanying drawings and specific embodiments.

[0030] like Figure 1 As shown, the present 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 of this embodiment adopts a segmentation everything model (SAM), including: 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 through the loss function and self-attention mechanism to obtain an initial encoder;

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

[0035]

[0036] Among them, 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. 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, when training the encoder, the input of the encoder is the original medical image, such as a CT scan image. First, the image X is widened into multiple segment length sequences, and then sent to the Transformer module in the SAM encoder in batches. Through the self-attention mechanism in the Transformer module, the contextual relationship between the pixels of the image is captured and the information extraction method is memorized through the feedforward neural network in the Transformer module. Finally, after passing through multiple layers of Transformer modules, fine-grained medical image features are obtained. The expression of the self-attention mechanism is as follows:

[0038]

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

[0040] Among them, Attention is self-attention, Q is query, K is key, V is value, d k is the feature dimension, which is determined by the sequence length, SOFTMAX is the SOFTMAX function, and K T is the transpose of the key, W Q , W K and W V They are the calculation matrices of Q, K and V, respectively, used to calculate the specific values ​​of Q, K and V. And freeze the original parameters of the SAM encoder during the encoder training process.

[0041] It should be noted that the corresponding inputs of modules at different levels can be obtained through the Transformer operations at multiple levels in the SAM encoder. This can effectively convert the information of the input image at different levels into embedded 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: embedding a medical space transposition adapter in an initial encoder to obtain a final encoder, and performing an encoding operation on the medical image through the final encoder to obtain an encoding result;

[0043] like Figure 2As shown, the medical space transposition adapter includes: residual adapter, dimension fusion adapter and normalized adapter, which act on the residual connection part, multi-head attention mechanism part and linear normalization layer part in each transformer layer of the SAM encoder respectively. The residual adapter, dimension fusion adapter and normalized adapter are all designed with the LoRA module as the basic adapter module. The English explanation in the figure is: NormalizedAdapter (basic adapter), MultilayerPerceptron (multi-layer perceptron), LayerNormalization (linear normalization layer), DimensionFusionAdapter (dimensional fusion adapter), Feature fusion (feature fusion), Depth feature (depth feature), Spatial feature (spatial feature), Multi-HeadAttention (multi-head attention mechanism), Depth Part (depth branch), Spatial Part (spatial branch), Upperprojection (upper projection), Lowerprojection (lower projection), ResidualAdapter (residual adapter).

[0044] Specifically, the LoRA module consists of a down-projection matrix, a ReLU function, and an up-projection parameter matrix. First, the down-projection parameter matrix compresses the input sequence into a sequence of length r, and then expands the sequence through the ReLU function activation and the up-projection parameter matrix. This process can greatly compress the amount of parameters in the training process, thereby achieving efficient adjustment of larger parameter models. The expression of the forward propagation process of the LoRA module during training is:

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

[0046] Among them, W 0 is the pre-trained weight matrix, ΔW is the low-rank decomposition, B is the up-projection parameter matrix, and A is the down-projection parameter matrix.

[0047] Furthermore, the dimension fusion adapter is embedded after the multi-head attention mechanism part and before the linear normalization layer part to perform modal fusion of the intermediate data of the transformer for spatial features and the transformer for processing deep features; the normalization adapter is embedded after the linear normalization layer part 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 the embedded 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 splits the attention mechanism of the SAM encoder into a spatial branch and a deep branch. In a three-dimensional sample with a depth of 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 is able to capture spatial correlations. In the deep branch, this embodiment transposes the input matrix to N×D×L and performs the same multi-head attention operation, where the interaction occurs between D×L, thereby learning deep correlations. Finally, the result of the deep branch is transposed back to the original shape and added to the result of the spatial branch to fuse the depth information.

[0049] Step 400: embedding the medical space transposition adapter into the second mask decoder to obtain a final decoder;

[0050] Specifically, the first mask decoder and the second mask decoder both include: a Transformer decoder, a pixel decoder and a convolutional layer; the final decoder includes: a first mask decoder and an optimized mask decoder; the optimized mask decoder is obtained by cascading a 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 in the figure, the Medical Space Transpose Adapter is embedded in the two layers of Transformer decoders of the second mask decoder to obtain a mask decoder, which improves the decoding efficiency of the second mask decoder. At the same time, the pixel decoders of the first mask decoder and the second mask decoder also exchange embedded features through skip connections, which improves the accuracy of the model and enhances the ability to capture local details. The English explanation in the figure is: Image Embedding, SAM Mask Decoder, Transformer Decoder, PixelDecoder, Conv Layer, Medical Space TransposeAdapter, Skipping Connections.

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

[0053] Furthermore, the Transformer decoder can be considered to have a structure that is basically the same as the transformer layer in the encoder, except that the calculation of the Q, K, and V values ​​in the self-attention layer in the Transformer decoder comes from the combined effect of the prompt embedding vector and the image embedding vector. The prompt embedding vector is often a point position label that is used to better guide the segmentation of the image. Specifically, the sequence for calculating the Q value comes from the prompt embedding vector, and the sequence for calculating the K and V values ​​comes from the image embedding vector. The calculation of the Q, K, and V values ​​of the self-attention in the transformer layer is simply performed using the image embedding vector.

[0054] Furthermore, for the pixel decoder, it is just a simple cross-attention operation, which also comes from the joint action of the prompt embedding vector and the image embedding vector, but the sequence for calculating the Q value comes from the image embedding vector, and the sequence for calculating the K and V values ​​comes from the prompt embedding vector. After passing through the two-layer pixel decoder, the output vector will undergo two more convolutional layer operations to obtain the output result of the pixel decoder. After passing through the two-layer Transformer decoder, the output result of the Transformer decoder is directly output.

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

[0056] 1) The SAM decoder and the medical space transposition adapter are connected in cascade form, which enhances the segmentation accuracy and simplifies the decoder parameters;

[0057] 2) In the design of the efficient transposition adapter for medical space, the spatial dimension of natural images and the depth dimension unique to medical images are integrated, so that the SAM model can capture multi-dimensional information at the same time, allowing the model to more fully understand the complexity of medical images and enhance the model's learning ability;

[0058] 3) By combining the spatial dimension with the depth dimension, 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, saving a lot of computing resources;

[0060] 5) Through the parallel transformer training strategy and the spatial transposition technology, the segmentation accuracy of complex medical images is improved and the problem of imbalance in medical image data categories is solved.

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

[0062] The present invention uses specific examples to illustrate the principles and implementation methods of the present invention. The above examples are only used to help understand the method and core ideas of the present invention. At the same time, for those skilled in the art, according to the ideas of the present invention, there will be changes in the specific implementation methods and application scope. In summary, the content of this specification should not be understood as limiting the present invention.

Claims

1. A neural network-based image segmentation method, characterized in that: The steps include: Build 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 by a loss function and a self-attention mechanism to obtain an initial encoder; Embedding a medical space transposition adapter in the initial encoder to obtain a final encoder, and performing an encoding operation on the medical image by the final encoder to obtain an encoding result; Embedding the medical space transposition adapter into the second mask decoder to obtain a final decoder; the final decoder includes: the first mask decoder and an optimized mask decoder; The encoding result is decoded by the final decoder to obtain an image segmentation result.

2. The neural network-based image segmentation method according to claim 1, characterized in that: The medical image segmentation model is a segmentation-everything model.

3. The neural network-based image segmentation method according to claim 1, characterized in that: The expression of the loss function is: Among them, 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, M is the number of transformer layers; The loss of the i-th transformer layer is calculated by the Dice loss function; the loss of the first mask decoder is calculated by the BCE loss function, and the loss of the second mask decoder is calculated by the iou loss function.

4. The neural network-based image segmentation method according to claim 1, characterized in that: The expression of the self-attention mechanism is: Among them, 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.

5. The neural network-based image segmentation method according to claim 1, characterized in that: The medical space transpose adapter includes: a residual adapter, a dimensionality fusion adapter and a normalization adapter; the residual adapter, the dimensionality fusion adapter and the normalization adapter are all composed of a LoRA module; the expression of the built-in forward propagation process of the LoRA module is: W0+ΔW=W0+BA; wherein W0 is a pre-trained weight matrix, ΔW is a low-rank decomposition, B is an upper projection parameter matrix, and A is a lower projection parameter matrix.

6. The neural network-based image segmentation method according to claim 5, characterized in that: A medical space transposition adapter is embedded in the initial encoder to obtain a final encoder, and the medical image is encoded by the final encoder to obtain an encoding result. Specifically, after the dimension fusion adapter is embedded in the multi-head attention mechanism part, the normalization adapter is embedded in the linear normalization layer part, and the residual adapter is used as the residual connection part to obtain the final encoder, and the medical image is encoded by the final encoder to obtain an encoding result.

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

8. The neural network-based image segmentation method according to claim 7, characterized in that: The encoding result is decoded by the final decoder to obtain an image segmentation result, including: generating a priori probability mask by the first mask decoder; Decoding the prior probability mask by the Transformer decoder and the pixel decoder to obtain a first result and a second result; The first result and the second result are combined by the optimized mask decoder to obtain the image segmentation result.

Images