Medical image segmentation method based on KAN network

Abstract

The invention relates to the technical field of image segmentation, and discloses a medical image segmentation method based on a KAN network, and the method comprises the following steps: 1, carrying out the data enhancement and preprocessing of a medical image data set, and constructing a training data set; 2, building an initial network model; 3, inputting the training data set into the constructed initial medical image segmentation model for training until the loss function converges, and obtaining a trained medical image segmentation model; according to the KAN channel attention module and the KAN space attention module, the KAN network replaces a traditional MLP structure, input features directly pass through a nonlinear activation function, compared with the MLP and the KAN network, the number of nodes can be greatly reduced, the model complexity is reduced, the computing power cost is reduced, and the operation efficiency is improved. The KAN channel attention module provided by the invention can solve the problem of excessive segmentation of an image in a decoder, and deletes invalid channel information. The KAN provided by the invention can effectively improve the medical image segmentation precision.

Description

Technical Field

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

[0002] In recent years, with the continuous development of deep learning, computer-aided diagnosis technology based on artificial intelligence has also made many significant progress in the medical field, buying time for the diagnosis and treatment of diseases. Among them, many models for medical image segmentation based on deep learning have also been proposed. In 2015, Long et al. proposed a fully convolutional neural network (FCN), which replaced the fully connected layer in the previous network with deconvolution, and upsampled the last layer of feature maps to restore it to its original size, ultimately achieving the purpose of accurate pixel-by-pixel classification and prediction. In the same year, Ronneberger et al. proposed the U-Net network. The U-Net network consists of a unique U-shaped symmetrical structure and jump connection parts, which can accurately capture the feature information in the available images.

[0003] However, the U-Net network is a pure convolutional network, which focuses on the features of local details, pays insufficient attention to the global whole, and has insufficient global feature modeling capabilities. Therefore, the attention mechanism was introduced into the medical image segmentation network. In 2021, Chen et al. proposed the TransUNet network. However, due to the addition of the attention mechanism, the number of parameters and the amount of calculation in the segmentation network increased greatly, and the requirements for hardware resources such as computing power and video memory also increased sharply. Therefore, it is urgent to provide a medical image segmentation method based on the KAN network. Summary of the invention

[0004] 1. Technical issues to be resolved

[0005] In view of the shortcomings of the prior art, the present invention provides a medical image segmentation method based on a KAN network, which has the advantages of using the KAN network technology to improve the decoder part and the jump connection part in the codec, and can improve the segmentation accuracy of the target image while reducing the calculation complexity of the model. It solves the problem that due to the addition of the attention mechanism, the number of parameters and the amount of calculation in the segmentation network are greatly increased, and the requirements for hardware resources such as computing power and video memory are also sharply increased.

[0006] (II) Technical solution

[0007] To achieve the above object, the present invention provides the following technical solution: a medical image segmentation method based on a KAN network, comprising the following steps:

[0008] Step 1: Perform data enhancement and preprocessing on the medical image dataset to construct a training dataset;

[0009] Step 2: Build the initial network model;

[0010] Step 3, inputting the training data set into the constructed initial medical image segmentation model for training until the loss function converges, thereby obtaining a trained medical image segmentation model;

[0011] Step 4: segment the image, obtain the medical image to be segmented and input it into the trained medical image segmentation model, and output the segmentation result.

[0012] Preferably, in step 1, the data enhancement method includes random flipping, scaling and cropping.

[0013] Preferably, in step 1, in order to enable the network to capture more features, the image data is preprocessed, the size of the original image is adjusted to 512x512, and the real mask image is interpolated and assigned to adjust its size to 512×512.

[0014] Preferably, in step 2, the initial network model includes: an encoder, a KAN channel attention module, a decoder, a KAN spatial attention module and a skip connection module.

[0015] Preferably, in step 3, according to the binary classification characteristics of each pixel in medical image segmentation, a binary cross entropy loss function is selected, and its specific formula is as follows:

[0016]

[0017] Cross entropy can measure the information difference between the true value and the predicted value distribution. In the formula, N is the total number of samples, that is, the number of pixels, and y is i is the category to which the i-th sample belongs. Medical image segmentation is binary classification. i is the predicted value of the i-th sample. Generally speaking, it is a probability value. This formula directly measures the similarity between two samples.

[0018] (III) Beneficial effects

[0019] Compared with the prior art, the present invention provides a medical image segmentation method based on KAN network, which has the following beneficial effects:

[0020] The KAN channel attention module and KAN spatial attention module proposed in the present invention enable the KAN network to replace the traditional MLP structure, so that the input features directly pass through the nonlinear activation function. Compared with MLP, the number of KAN network nodes can be greatly reduced, the model complexity is reduced, and the computing cost is reduced. The KAN channel attention module proposed in the present invention can solve the problem of over-segmentation of images in the decoder and delete invalid channel information. The KAN network proposed in the present invention can effectively improve the accuracy of medical image segmentation. BRIEF DESCRIPTION OF THE DRAWINGS

[0021] Figure 1 A schematic diagram of the structure of the network provided by the present invention;

[0022] Figure 2 It is a schematic diagram of the ConvNeXt module structure;

[0023] Figure 3 Schematic diagram of the KAN channel attention module provided by the present invention;

[0024] Figure 4 Schematic diagram of the KAN spatial attention module provided by the present invention;

[0025] Figure 5 This is a GLAS comparison test diagram provided by the present invention. DETAILED DESCRIPTION

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

[0027] It should be noted that all directional indications in the embodiments of the present application are only used to explain the relative position relationship, movement status, etc. between the components in a certain specific posture. If the specific posture changes, the directional indication will also change accordingly.

[0028] In this application, unless otherwise clearly specified and limited, the terms "connection", "fixation", etc. should be understood in a broad sense. For example, "fixation" can be a fixed connection, a detachable connection, or an integral connection; it can be a mechanical connection or an electrical connection; it can be a direct connection or an indirect connection through an intermediate medium, it can be the internal connection of two elements or the interaction relationship between two elements, unless otherwise clearly defined. For ordinary technicians in this field, the specific meanings of the above terms in this application can be understood according to specific circumstances.

[0029] In addition, in this application, descriptions such as "first", "second", etc. are only used for descriptive purposes and cannot be understood as indicating or implying their relative importance or implicitly indicating the number of technical features indicated. Therefore, the features defined as "first" and "second" may explicitly or implicitly include at least one of the features. In addition, the technical solutions between the various embodiments can be combined with each other, but they must be based on the ability of ordinary technicians in this field to implement them. When the combination of technical solutions is contradictory or cannot be implemented, it should be deemed that such combination of technical solutions does not exist and is not within the scope of protection required by this application.

[0030] In the related technology, the U-Net network is a pure convolutional network, which focuses on the features of local details, pays insufficient attention to the global whole, and has insufficient global feature modeling capabilities. Therefore, the attention mechanism is introduced into the medical image segmentation network. In 2021, Chen et al. proposed the TransUNet network. However, due to the addition of the attention mechanism, the number of parameters and the amount of calculation in the segmentation network are greatly increased, and the requirements for hardware resources such as computing power and video memory are also sharply increased. Therefore, it is urgent to provide a medical image segmentation method based on the KAN network.

[0031] In order to solve the problems in the related art to a certain extent, the embodiment of the present application provides a medical image segmentation method based on a KAN network. The encoder part uses a ConvNeXt network to extract image features, and uses a ConvNeXtBlock module to replace the 3×3 convolution kernel used to extract image features in the original U-Net network encoder, and selects a structure in which the number of modules per layer is 3, 3, 9, and 3 respectively; in the decoder part, a depth-separable convolution module is used to replace the 3×3 convolution kernel in the original U-Net network decoder, and a KAN channel attention module is added after each layer of the decoder; in the jump connection part, the features extracted by each layer of the encoder are merged with the decoder part features of the corresponding layer after passing through the KAN spatial attention module;

[0032] The KAN spatial attention module is similar to ViT. It copies the input features into three copies X, Q, and V. The size of the three input features is C×H×W. After global pooling of the input feature Q, it is compressed to C×1×1 in the spatial range to meet the KAN network input. Then, the spatial weights are learned through two layers of KAN networks, so that the number of feature channels becomes And through the activation function Softmax, the model calculation amount is reduced. The input feature V becomes Then the feature dimension changes to The processed feature matrix Q and V are multiplied by matrices to obtain a feature matrix of size 1×HW. After the matrix is ​​dimensionally changed and processed with the Sigmoid function, a feature matrix of size 1×H×W is obtained. Then, the feature matrix is ​​Hadamard-producted with the input feature X to obtain the final output feature C×H×W. The KAN spatial attention module reduces the computational complexity of the MLP used in the general Transformer and uses mathematical ideas to reduce computing power costs.

[0033] The KAN channel attention module is based on the channel attention module. After the input features are subjected to the maximum pooling operation and the average pooling operation, the global statistical information of each channel is obtained respectively. The difference is that the channel weights are learned through two layers of KAN networks. The KAN network learns the channel weights by nesting nonlinear functions, and the two processed results are added by matrix element addition, and the weights are normalized to between 0 and 1 using the Sigmoid function. Each channel is scaled, and finally, the scaled channel features are multiplied with the input features. The KAN channel attention module at the end of each layer of the decoder can be used to solve the problem of over-segmentation in the decoder.

[0034] The present application is described below with reference to the accompanying drawings and specific embodiments:

[0035] Combination Figure 1-Figure 5 The embodiment of the present application provides a medical image segmentation method based on a KAN network. Step one is to perform data enhancement and preprocessing on a medical image dataset to construct a training dataset.

[0036] The data enhancement method adopted in this embodiment includes random flipping, scaling and cropping. At the same time, in order to enable the network to capture more features, the image data is preprocessed, the size of the original image is adjusted to 512x512, and the real mask image is interpolated and assigned to adjust its size to 512×512, which is convenient for training the network model.

[0037] Step 2: Build the initial network model.

[0038] Figure 1 It is a schematic diagram of the network structure provided by the present invention. Among them, the encoder part uses the ConvNeXt network to extract image features, and the ConvNeXt Block module is used to replace the 3×3 convolution kernel used to extract image features in the original U-Net network encoder, and the structure with 3, 3, 9, and 3 modules per layer is selected; the decoder part uses a depth-separable convolution module to replace the 3×3 convolution kernel in the original U-Net network decoder, and a KAN channel attention module is added after each layer of the decoder; in the jump connection part, the features extracted by each layer of the encoder are merged with the decoder part features of the corresponding layer after passing through the KAN spatial attention module.

[0039] The ConvNeXt network used in the encoder makes the image pass through a 4×4 convolution kernel with a stride of 4, and then uses the ConvNeXt module to continue extracting features. Figure 2This is a schematic diagram of the structure of the ConvNeXt module, where the features first pass through a depth-separable convolution of size 7×7, then after channel normalization, pass through a 1×1 convolution kernel and a GELU activation function, and finally pass through a 1×1 convolution kernel and a scaling layer. The output features are added to the input features to finally obtain a feature extracted by a ConvNeXt module. The specific structure of the network convolution kernel is shown in Table 1 below:

[0040]

[0041]

[0042] The decoder uses a depth-separable convolution module to replace the 3×3 convolution kernel in the original U-Net network decoder, and adds a KAN channel attention module after each layer of the decoder. The last layer of the decoder is different from the previous layers. The feature information received from the previous layer of the decoder is combined with the feature information output from the first layer of the encoder and then output by the KAN spatial attention module. The feature information is directly input into the KAN channel attention module, and then goes through two layers of transposed convolution and conventional convolution to obtain the segmentation result.

[0043] The specific structure of the KAN channel attention module is as follows Figure 3 As shown in the figure. The KAN channel attention module is based on the channel attention module. After the input features are subjected to the maximum pooling operation and the average pooling operation, the global statistical information of each channel is obtained respectively. The difference is that the channel weights are learned through two layers of KAN networks. The KAN network learns the channel weights by nesting nonlinear functions. The two processed results are added by matrix element addition, and the weights are normalized to between 0 and 1 using the Sigmoid function. Each channel is scaled, and finally, the scaled channel features are multiplied with the input features. The KAN channel attention module at the end of each layer of the decoder can be used to solve the problem of over-segmentation in the decoder, and the KAN network is more mathematically interpretable than the traditional MLP.

[0044] In the skip connection part, the features extracted from each encoder layer are merged with the decoder features of the corresponding layer after passing through the KAN spatial attention module. Figure 4 The specific structural diagram of the KAN spatial attention module. The KAN spatial attention module is similar to ViT. The input feature is copied into three copies X, Q, and V. The size of the three input features is C×H×W. After the input feature Q is globally pooled, it is compressed to C×1×1 in the spatial range to meet the KAN network input. Then, the spatial weights are learned through a two-layer KAN network to reduce the model calculation amount and make the number of output feature channels become Then the activation function Softmax is used. The input feature V becomes Then the feature dimension changes to The processed feature matrix Q and V are multiplied by matrices to obtain a feature matrix of size 1×HW. The matrix is ​​transformed in dimension and processed by Sigmoid function to obtain a feature matrix of size 1×H×W. The feature matrix is ​​then Hadamard-producted with the input feature X to obtain the final output feature C×H×W. The KAN spatial attention module reduces the computational complexity of MLP in general Transformer and uses mathematical ideas to reduce computing power costs.

[0045] Step three, input the training data set into the constructed initial medical image segmentation model for training until the loss function converges to obtain a trained medical image segmentation model.

[0046] According to the binary classification characteristics of each pixel in medical image segmentation, a binary cross entropy loss function is selected, and its specific formula is as follows:

[0047]

[0048] Cross entropy can measure the information difference between the true value and the predicted value distribution. In the formula, N is the total number of samples, that is, the number of pixels, and y i is the category to which the i-th sample belongs. Medical image segmentation is binary classification. i is the predicted value of the ith sample. Generally speaking, it is a probability value. This formula directly measures the similarity between two samples. Using the cross entropy loss function as the loss function for network training can effectively avoid the problem that the network cannot continue training when the gradient is small, and make the weight of the neural network effectively converge to the true value during the training process.

[0049] Step 4: Segment the image

[0050] Obtain the medical image to be segmented and input it into the trained medical image segmentation model, and output the segmentation result.

[0051] In order to evaluate the medical image segmentation network proposed in this embodiment, the GLAS dataset is selected to show its experimental results.

[0052] The GLAS dataset is a set of datasets for colorectal glandular tissue segmentation tasks, originally released by the 2015 MICCAI challenge. The dataset contains 165 images with four different resolutions, each with its own label. These images are from 16 H&E-stained histological sections, which represent colorectal adenocarcinomas at the T3 or T4 stage.

[0053] like Figure 5, is the experimental segmentation results comparison between the proposed network and the other six networks. From the segmentation results, it can be seen intuitively that the proposed network not only pays attention to the overall situation, but also reduces the situation of not segmenting the target, and significantly reduces the situation of over-segmentation in local details, and the segmentation effect is better. The following table shows the specific measurement indicators of the proposed network compared with other networks. From the specific indicators, it can be seen that the proposed network performs better than other networks in the three evaluation indicators of mean intersection over union (mIoU), Dice similarity coefficient (DSC) and recall rate (Recall). It is lower than UNet++ in precision, but it is improved compared with the remaining five networks. In terms of model parameter quantity (Params), it is higher than U-Net network, U2-Net network and UNet++ network, and the model parameter quantity has increased, but in terms of training calculation quantity (FLOPs), it is much lower than other networks due to the use of KAN network instead of MLP structure, and the segmentation performance is also significantly improved. Therefore, the proposed network overall reduces the model complexity and obtains higher segmentation performance.

[0054]

[0055]

[0056] In the description of this specification, the description with reference to the terms "one embodiment", "some embodiments", "example", "specific example", or "some examples" etc. means that the specific features, structures, materials or characteristics described in conjunction with the embodiment or example are included in at least one embodiment or example of the present application. In this specification, the schematic representations of the above terms do not necessarily refer to the same embodiment or example. Moreover, the specific features, structures, materials or characteristics described may be combined in any one or more embodiments or examples in a suitable manner. In addition, those skilled in the art may combine and combine different embodiments or examples described in this specification.

[0057] In addition, the technical solutions between the various embodiments can be combined with each other, but they must be based on the fact that they can be implemented by ordinary technicians in this field. When the combination of technical solutions is contradictory or cannot be implemented, it should be deemed that such combination of technical solutions does not exist and is not within the scope of protection required by this application.

[0058] Although embodiments of the present invention have been shown and described, it will be appreciated by those skilled in the art that various changes, modifications, substitutions and variations may be made to the embodiments without departing from the principles and spirit of the present invention, and that the scope of the present invention is defined by the appended claims and their equivalents.

Claims

1. A medical image segmentation method based on KAN network, characterized in that: The steps include: Step 1: Perform data enhancement and preprocessing on the medical image dataset to construct a training dataset; Step 2: Build the initial network model; Step 3, inputting the training data set into the constructed initial medical image segmentation model for training until the loss function converges, thereby obtaining a trained medical image segmentation model; Step 4: segment the image, obtain the medical image to be segmented and input it into the trained medical image segmentation model, and output the segmentation result.

2. The medical image segmentation method based on KAN network according to claim 1, characterized in that: In the step 1, the data enhancement method includes random flipping, scaling and cropping.

3. The medical image segmentation method based on KAN network according to claim 1, characterized in that: In the step 1, in order to enable the network to capture more features, the image data is preprocessed, the size of the original image is adjusted to 512x512, and the real mask image is interpolated and assigned to adjust its size to 512×512.

4. The medical image segmentation method based on KAN network according to claim 1, characterized in that: In the step 2, the initial network model includes: an encoder, a KAN channel attention module, a decoder, a KAN spatial attention module and a skip connection module.

5. The medical image segmentation method based on KAN network according to claim 1, characterized in that: In step 3, according to the characteristics of the binary classification of each pixel in medical image segmentation, a binary cross entropy loss function is selected, and its specific formula is as follows: Cross entropy can measure the information difference between the true value and the predicted value distribution. In the formula, N is the total number of samples, that is, the number of pixels, and y is i is the category to which the i-th sample belongs. Medical image segmentation is binary classification. i is the predicted value of the i-th sample. Generally speaking, it is a probability value. This formula directly measures the similarity between two samples.

Images