A multi-modal medical image segmentation method and system based on Transformer and UNet

Abstract

The application discloses a kind of multi-modal medical image segmentation method and system based on Transformer and UNet.To the framework of the encoder-decoder structure UNet, two Transformer modules are designed to replace the skip connection in UNet.The two modules are respectively a channel-level cross-fusion Transformer module for multi-scale feature information fusion and a channel-level cross-attention mechanism module for encoder-decoder information fusion, and the proposed method is called CTransUNet.For each segmentation task, first, train the network with a large amount of external data;then, the image to be segmented is preprocessed and sent to the trained network for prediction;finally, the prediction results are post-processed to eliminate too few false categories.The application uses an end-to-end overall framework to segment medical images, and the input image can directly obtain the segmentation result.

Description

TECHNICAL FIELD

[0001] The application belongs to the field of computer vision, deep learning technology and medical image segmentation, and particularly relates to a multi-modal medical image segmentation method and system based on a Transformer and UNet. BACKGROUND

[0002] Medical image segmentation is a popular research direction in computer vision in recent years and has important applications in clinical diagnosis, treatment plan formulation and medical research. With the rapid development of artificial intelligence technology, deep learning technology has gradually become the mainstream method for image segmentation research. Because of the imaging principle of medical images such as MRI, CT and ultrasound images, there are natural noise, artifacts and other interference, as well as fuzzy background, unclear boundary and small target segmentation area, which makes it difficult to accurately segment. The UNet with an encoder-decoder structure and a skip connection is a relatively effective solution for many medical image segmentation tasks, but the simple nature of the skip connection still cannot learn the multi-scale features of the encoder and the information fusion between the two branches of the encoder-decoder, which limits the learning ability of the network and makes the network unable to obtain more accurate segmentation results. SUMMARY

[0003] The application aims to provide a multi-modal medical image segmentation method and system based on a Transformer and UNet to solve the technical problem that the skip connection module in the existing UNet network structure cannot learn the multi-scale features of the encoder and the information fusion between the two branches of the encoder-decoder, resulting in a significant semantic gap.

[0004] To solve the above technical problems, the specific technical solutions of the application are as follows:

[0005] A multi-modal medical image segmentation method based on a Transformer and UNet, comprising the following steps:

[0006] Step 1, data preprocessing step: each specific medical image segmentation task has corresponding training data set and segmentation image data; for each image of the training data set or the segmentation image, use preprocessing methods such as empty background cropping, intensity value normalization and random cropping to obtain input data structure suitable for the proposed network;

[0007] Step 2, build the network structure of CTrans UNet and train the network, including the following sub-steps:

[0008] Step 2.1, design the network structure of CTrans UNet;

[0009] Step 2.2, select and construct training samples from the processed data set;

[0010] Step 2.3, design a good training and prediction strategy, set network hyperparameters and loss functions, use the training samples selected and constructed in step 2.2 to effectively train the CTrans UNet network, optimize the parameters of the CTrans UNet network until the training converges;

[0011] Step 3, for the image to be segmented, use sliding prediction to input the trained CTrans UNet network for segmentation calculation to obtain the final segmentation result.

[0012] Further, the data preprocessing step in step 1 is as follows:

[0013] For the training set, first perform empty background cropping on each image to remove the background area to reduce invalid computation; then perform intensity value normalization to reduce the variance of intensity values between samples; finally, perform random cropping, which needs to meet the network input size, and because of randomness, it has the effect of data enhancement;

[0014] For the image to be segmented, first perform empty background cropping on each image to remove the background area to reduce invalid computation, and record the cropping position; then perform intensity value normalization to reduce the variance of intensity values between samples.

[0015] Further, the network structure of CTrans UNet in step 2.1 is designed, which includes the following steps:

[0016] Step 2.1.1, design a channel-level cross-fusion Transformer module, referred to as a CCFT module, which uses the universal advantage of Transformer structure in long-range dependence to make up for the defects of local learning and fuse multi-scale encoder feature information;

[0017] Step 2.1.2, design a channel-level cross-attention mechanism module, referred to as a CCA module, which effectively fuses the encoder features of the current layer and the decoder features of the upper layer, and makes up for the semantic gap between the two branches;

[0018] Step 2.1.3, based on the UNet network structure, replace the skip connection with the channel-level cross-fusion Transformer module and the channel-level cross-attention mechanism module for better feature fusion, and form a hybrid network structure that combines the advantages of Transformer and UNet, called CTrans UNet.

[0019] Further, the training sample is selected and constructed from the processed data set in step 2.2, specifically including the following steps:

[0020] According to the conditions of the hardware environment and the size of the single image, the same number of images are randomly selected from the data set for network training in each round of training, and the number needs to be a multiple of 2.

[0021] Further, in step 2.3, the training and prediction strategy is designed, the network hyperparameters and loss function are set, the training sample selected and constructed in step 2.2 is used to effectively train CTrans UNet, and the parameters of CTrans UNet network are optimized until the training converges, specifically including the following steps:

[0022] Step 2.3.1, using 5-fold cross-validation training method, 5 trained CTrans UNet networks are obtained, and the image to be segmented is respectively passed through the 5 networks to obtain 5 segmentation results, and finally the average of the 5 results is taken as the final result;

[0023] Step 2.3.2, the necessary settings of network hyperparameters are as follows: set the initial learning rate as lr, the optimizer as Adam, the learning round as epoch, and the learning rate adjustment strategy as ReduceLRonPlateau. If the loss does not decrease within patience rounds, the learning rate is reduced to the original factor times;

[0024] Step 2.3.3, using generalized Dice Loss function as the loss function of most medical image segmentation methods, different weights are given to different categories in the loss function, and the formula is:

[0025]

[0026] Where GDL represents generalized Dice Loss, K is the number of categories, r kn is the true label of pixel n in the kth category, p kn is the prediction result of pixel n in the kth category, w k is the weight of each category when calculating the Dice Loss function, based on the principle that the weight of the category with more area should be smaller, and the calculation method of w k is:

[0027]

[0028] The reciprocal of the square of the sum of the voxel values is used to calculate the weight, and the weight of the category with more voxels is smaller, and the weight of the category with less voxels is larger, which can effectively solve the problem of class imbalance.

[0029] Further, step 3 adopts a sliding prediction mode to send the to-be-segmented image into the trained CTrans UNet network for segmentation calculation to obtain a final segmentation result, and specifically includes the following steps:

[0030] Step 301, slidingly cutting the preprocessed to-be-segmented image, and cutting out a plurality of image blocks along each axis of the to-be-segmented image according to a fixed step length step, wherein the size of the image block needs to meet the network input size requirement;

[0031] Step 302, sending the plurality of image blocks obtained by sliding cutting into the trained CTrans UNet network for prediction to obtain the segmentation result of each block;

[0032] Step 303, for the segmentation result output of each block, an argmax function needs to be used to classify each pixel point of the segmentation result output to obtain the classified segmentation result, and the argmax function is as follows:

[0033]

[0034] For a function f(x), when x takes a value range S, the result of argmax is the x point set that makes f(x) take the maximum value; wherein x and y represent variables, S represents the value range of the variable, and f(x) and f(y) represent the function values corresponding to x and y;

[0035] Step 304, placing the segmentation result of each block in the corresponding position of the to-be-segmented image according to the position when cutting, to combine into a complete segmentation result.

[0036] The application further discloses a multi-modal medical image segmentation system based on a Transformer and a UNet, which is characterized by comprising a data preprocessing unit, a CTrans UNet network training unit and a to-be-segmented image prediction unit.

[0037] The data preprocessing unit performs data preprocessing steps such as air background removal, intensity value normalization and random cropping on the training set and the to-be-segmented image.

[0038] The CTrans UNet network training unit is configured to perform the following steps:

[0039] Step A, designing the network structure of the CTrans UNet network;

[0040] Step B, selecting and constructing training samples as the input of each round of network training;

[0041] Step C, design training and prediction strategy, set network hyperparameters, loss function, use the training sample selected and constructed in step B to effectively train the CTrans UNet network, optimize the parameters of the CTrans UNet network, until the training reaches convergence;

[0042] The prediction unit of the image to be segmented is configured to perform the following actions:

[0043] Step A, slidingly cut the preprocessed image to be segmented along each axis according to a fixed step length step to obtain a plurality of image blocks;

[0044] Step B, the image blocks are respectively sent into the trained CTrans UNet network for prediction to obtain the segmentation result of each block;

[0045] Step C, use the argmax function to classify each pixel point in the segmentation result of each image block, and the result after classification is the final segmentation result of the image block;

[0046] Step D, place the final segmentation result of the image block to the corresponding position of the image to be segmented according to the position when cutting, and combine to form the complete segmentation result of the image to be segmented.

[0047] The multi-modal medical image segmentation method and system based on the Transformer and UNet have the following advantages:

[0048] 1. The multi-modal medical image segmentation method based on the Transformer and UNet has universality in medical image segmentation, and can be applied in various medical images and various segmentation scenarios;

[0049] 2. The application designs a channel-level cross-fusion Transformer module, so that the network can fuse multi-scale feature information on the encoder, and because of the long-range dependence characteristics of the Transformer, global information can be learned;

[0050] 3. The application designs a channel-level cross-attention mechanism module to fuse the features of the encoder and the decoder, so that the semantic gap between the two branches of the network is further bridged;

[0051] 4. The application adopts a good preprocessing step, training and prediction strategy, sets network hyperparameters, and defines a general loss function, which is very effective for the effective training of the network, improves the accuracy of the network, and reduces the training time as much as possible. BRIEF DESCRIPTION OF DRAWINGS

[0052] Figure 1A schematic diagram of a multi-modal medical image segmentation method based on a Transformer and a UNet according to the present application.

[0053] Figure 2 A schematic diagram of a CTrans UNet network overall framework according to the present application. DETAILED DESCRIPTION

[0054] In order to better understand the purpose, structure and function of the present application, a multi-modal medical image segmentation method and system based on a Transformer and a UNet according to the present application will be described in further detail below with reference to the accompanying drawings.

[0055] The present application first proposes a multi-modal medical image segmentation method based on a Transformer and a UNet, as shown in the accompanying drawings, which comprises the following steps: Figure 1

[0056] Step 1, data preprocessing step: each specific medical image segmentation task has two parts of data, a corresponding training data set and an image to be segmented. For each image of the training data set or the image to be segmented, the preprocessing method of empty background cropping, intensity value normalization and random cropping is used to obtain the input data structure suitable for the proposed network;

[0057] For the training set, first, empty background cropping is performed on each image to remove the background area to reduce the invalid calculation amount; then, intensity value normalization is performed to reduce the variance of intensity values between samples; and finally, random cropping is performed, and the cropping size needs to meet the network input size, and because of the randomness of the cropping position selection, the effect of data enhancement can be achieved;

[0058] For the image to be segmented, first, empty background cropping is performed on each image to remove the background area to reduce the invalid calculation amount, and the cropping position is recorded, which is needed to place the segmentation result in the same position after the subsequent prediction is completed; then, intensity value normalization is performed to reduce the variance of intensity values between samples.

[0059] Step 2, build the network structure of CTrans UNet and perform network training, including the following sub-steps:

[0060] Step 2.1, design the network structure of CTrans UNet, specifically including the following steps:

[0061] Step 201, design a channel-level cross-fusion Transformer module, referred to as a CCFT module, which utilizes the universal advantage of the Transformer structure in long-range dependence to make up for the defects of local learning, and can fuse multi-scale encoder feature information;

[0062] ​Step 202, design a channel-level cross-attention mechanism, referred to as CCA module, to effectively fuse the encoder features of the current layer and the decoder features of the upper layer, and to bridge the semantic gap between the two branches;

[0063] Step 203, based on the UNet structure, replace the skip connection with a channel-level cross-fusion Transformer module and a channel-level cross-attention mechanism module for better feature fusion, forming a hybrid network structure that combines the advantages of Transformer and UNet, namely CTrans UNet, as shown in Figure 2 .

[0064] Step 2.2, select and construct training samples from the processed data set: according to the conditions of the hardware environment and the size of individual images, randomly select the same number of images from the data set for each round of training, and the number must be a multiple of 2.

[0065] Step 2.3, design a good training and prediction strategy, set network hyperparameters and loss functions, use the training samples selected and constructed in step 2.2 to effectively train CTrans UNet, optimize the parameters of CTrans UNet network until the training converges, and the specific steps are as follows:

[0066] Step 204, use 5-fold cross-validation training method to obtain 5 trained CTrans UNet networks, and use the 5 networks to obtain 5 segmentation results for the image to be segmented, and finally take the average as the final result;

[0067] Step 205, the necessary settings of network hyperparameters are as follows: set the initial learning rate to 1×10 -3 , the optimizer is Adam, the learning round is 300, and the learning rate adjustment strategy is ReduceLRonPlateau. If the loss does not decrease within 30 rounds, the learning rate is reduced to 0.1 times the original;

[0068] Step 206, use the generalized Dice Loss loss function as the loss function for medical image segmentation tasks, and assign different weights to different classes in the loss function. The general formula is:

[0069]

[0070] where GDL represents the generalized Dice Loss, K is the number of classes, r kn is the true label of pixel n in the kth class, p kn is the predicted result of pixel n in the kth class, and w kThat is, the weight of each category when calculating the Dice Loss, based on the principle that the larger the category weight should be smaller, w k The calculation method is as follows:

[0071]

[0072] The reciprocal of the square of the sum of the voxel values is used to calculate the weight, the weight of the class with more voxels is smaller, and the weight of the class with less voxels is larger, which can effectively solve the problem of class imbalance.

[0073] Step 3, for the image to be segmented, a sliding prediction method is used to input the trained CTrans UNet network for segmentation calculation to obtain the final segmentation result, which specifically includes the following steps:

[0074] Step 301, sliding cutting of the preprocessed image to be segmented is performed, and a plurality of image blocks are cut out along the axes of the image to be segmented according to a fixed step length step, and the size of the image block needs to meet the network input size requirement;

[0075] Step 302, a plurality of image blocks obtained by sliding cutting are respectively input into the trained CTrans UNet network for prediction to obtain the segmentation result of each block;

[0076] Step 303, for the segmentation result output of each block, the argmax function needs to be used to classify each pixel point of the segmentation result output, that is, the class with the maximum probability is selected to obtain the classification result after classification, and the argmax function is as follows:

[0077]

[0078] For a function f(x), when x takes a value range S, the result of argmax is the set of x points that make f(x) take the maximum value. Where x, y represent variables, S represents the value range of the variable, f(x) and f(y) represent the function values corresponding to x and y.

[0079] Step 304, the segmentation result of each block is placed in the corresponding position of the image to be segmented according to the position when cutting, and combined into a complete segmentation result.

[0080] The application also provides a multi-modal medical image segmentation system based on Transformer and UNet, which comprises a data preprocessing unit, a CTrans UNet network training unit and a prediction unit of an image to be segmented.

[0081] The data preprocessing unit performs the following data preprocessing steps: background removal, intensity value normalization and random cropping, for the training set and the image to be segmented.

[0082] The CTrans UNet network training unit is configured to train the designed CTrans UNet network and perform the following steps:

[0083] Step A, designing the network structure of the CTrans UNet network;

[0084] Step B, selecting and constructing training samples as the input of each round of network training;

[0085] Step C, designing a good training and prediction strategy, setting network hyperparameters and loss functions, using the training samples selected and constructed in step B to effectively train the CTrans UNet network, optimizing the parameters of the CTrans UNet network until the training converges;

[0086] The prediction unit of the image to be segmented is configured to perform the following actions:

[0087] Step A, slidingly cutting the preprocessed image to be segmented along each axis according to a fixed step size step to obtain a plurality of image patches;

[0088] Step B, sending the image patches obtained by slidingly cutting in step A into the trained CTrans UNet network for prediction to obtain the segmentation result of each patch;

[0089] Step C, using the argmax function to classify each pixel point in the segmentation result of each image patch to obtain the final segmentation result after classification as the final segmentation result of the image patch;

[0090] Step D, placing the final segmentation result of the image patch into the corresponding position of the image to be segmented according to the position when cutting, and combining to form the complete segmentation result of the image to be segmented.

[0091] Those skilled in the art can understand that each block in these structural diagrams and / or block diagrams and / or flowcharts and combinations of blocks in these structural diagrams and / or block diagrams and / or flowcharts can be implemented by computer program instructions. These computer program instructions can be provided to a general-purpose computer, a professional computer, or a processor of other programmable data processing method to generate a machine, so that the instructions executed by the computer or the processor of other programmable data processing method create a method specified in the block or blocks of the structural diagram and / or block diagram and / or flowchart.

[0092] It is to be understood that the present application is described by way of example only, and that modifications or alterations can be made to the features and embodiments described without departing from the spirit and scope of the application. In addition, modifications can be made to the features and embodiments described to accommodate specific situations and materials without departing from the spirit and scope of the application. Accordingly, the application is not limited to the specific embodiments disclosed herein, but rather, the scope of the application includes all embodiments falling within the scope of the claims.

Claims

1. A multimodal medical image segmentation method based on Transformer and UNet, characterized in that, Includes the following steps: Step 1: Data Preprocessing Steps: Each specific medical image segmentation task has two parts of data: a corresponding training dataset and images to be segmented. For each image in the training dataset or the images to be segmented, preprocessing methods such as empty background cropping, intensity value normalization, and random cropping are used to obtain an input data structure suitable for the proposed network. Step 2: Build the CTrans UNet network structure and train the network, including the following sub-steps: Step 2.1: Design the network structure of CTrans UNet; Step 2.2: Select and construct training samples from the processed dataset; Step 2.3: Design training and prediction strategies, set network hyperparameters and loss functions, and use the training samples selected and constructed in Step 2.2 to effectively train the CTrans UNet network and optimize the parameters of the CTrans UNet network until the training converges. Step 3: For the image to be segmented, the sliding prediction method is used to feed it into the trained CTrans UNet network for segmentation calculation to obtain the final segmentation result; Step 2.1, which designs the CTrans UNet network structure, specifically includes the following steps: Step 2.1.1: Design a channel-level cross-fusion Transformer module, referred to as the CCFT module, to fuse multi-scale encoder feature information; Step 2.1.2: Design a channel-level cross-attention mechanism module, referred to as the CCA module, to effectively fuse the features of the encoder at this layer and the features of the decoder upsampled at the upper layer, thereby bridging the semantic gap between the two branches. Step 2.1.3: Based on the UNet network structure, replace skip connections with channel-level cross-fusion Transformer modules and channel-level cross-attention mechanism modules to achieve better feature fusion, forming a hybrid network structure that combines the advantages of both Transformer and UNet, called CTrans UNet; Step 2.3 involves designing training and prediction strategies, setting network hyperparameters and loss functions, and effectively training the CTrans UNet using the training samples selected in step 2.2 to optimize the parameters of the CTrans UNet network until training converges. Specifically, this includes the following steps: Step 2.3.1: Use the 5-fold cross-validation training method to obtain 5 trained CTrans UNet networks, and pass the image to be segmented through these 5 networks to obtain 5 segmentation results. Finally, take the average of them as the final result. Step 2.3.2: The necessary settings for network hyperparameters are as follows: set the initial learning rate to lr, select Adam as the optimizer, set the number of learning epochs to epochs, and set the learning rate adjustment strategy to ReduceLRonPlateau. If the loss does not decrease within the patience epoch, reduce the learning rate to factor times the original value. Step 2.3.3: Use the generalized Dice Loss function as the loss function for the medical image segmentation method, assigning different weights to different categories on the loss function. The formula is as follows: Where GDL stands for Generalized Dice Loss, K is the number of categories, and r kn Let p be the true label of pixel n in the k-th class. kn For the prediction result of pixel n in the kth class, w k That is, the weight of each category when calculating the Dice Loss function, w k The calculation method is as follows:

2. The multimodal medical image segmentation method based on Transformer and UNet according to claim 1, characterized in that, The data preprocessing steps in step 1 are as follows: For the training set, each image is first cropped with an empty background to remove background areas and reduce unnecessary computation; then the intensity values ​​are normalized to reduce the variance of intensity values ​​between samples; finally, random cropping is performed, and the cropping size needs to meet the network input size. For each image to be segmented, the background is first cropped to remove the background area to reduce unnecessary computation, and the cropping position is recorded. Then, the intensity values ​​are normalized to reduce the variance of intensity values ​​between samples.

3. The multimodal medical image segmentation method based on Transformer and UNet according to claim 1, characterized in that, Step 2.2, which involves selecting and constructing training samples from the processed dataset, specifically includes the following steps: Depending on the hardware environment and the size of a single image, each training round randomly selects the same number of images from the dataset and feeds them into the network for training, and this number must be a multiple of 2.

4. The multimodal medical image segmentation method based on Transformer and UNet according to claim 1, characterized in that, Step 3 uses a sliding prediction method to feed the image to be segmented into the trained CTrans UNet network for segmentation calculation to obtain the final segmentation result. Specifically, it includes the following steps: Step 301: Perform sliding slicing on the preprocessed image to be segmented, cutting out a number of small image blocks along each axis of the image to be segmented at a fixed step size. The size of the small image blocks must meet the network input size requirements. Step 302: Feed the image blocks obtained by sliding slicing into the trained CTrans UNet network for prediction to obtain the segmentation results of each block. Step 303: For the segmentation results of each small block, the argmax function needs to be used to classify each pixel in the segmentation results to obtain the classified segmentation results. The argmax function is shown below: For a function f(x), when the range of x is S, the result of argmax is the set of x points that make f(x) reach its maximum value; where x and y are variables, S represents the range of values ​​of the variables, and f(x) and f(y) represent the function values ​​corresponding to x and y. Step 304: Place the segmentation results of each small block into the corresponding positions in the image to be segmented according to the positions when the blocks were cut, and combine them into a complete segmentation result.

5. A multimodal medical image segmentation system based on Transformer and UNet according to any one of claims 1-4, characterized in that, It includes a data preprocessing unit, a CTrans UNet network training unit, and a prediction unit for the image to be segmented; The data preprocessing unit performs the following preprocessing steps for the training set and the image to be segmented: background removal, intensity value normalization, and random cropping. The CTrans UNet network training unit, used to train the designed CTrans UNet network, is configured to perform the following steps: Step A: Design the network structure of the CTrans UNet network; Step B: Select and construct training samples as input for each round of network training; Step C: Design training and prediction strategies, set network hyperparameters and loss functions, and use the training samples selected and constructed in Step B to effectively train the CTrans UNet network, optimize the parameters of the CTrans UNet network, until the training converges. The prediction unit of the image to be segmented is configured to perform the following actions: Step A: The preprocessed image to be segmented is slid along each axis at a fixed step size to obtain several small image blocks; Step B: Feed each small image patch into the trained CTrans UNet network for prediction to obtain the segmentation result of each patch; Step C: Use the argmax function to classify each pixel in the segmentation result of each image patch, and use the classification result as the final segmentation result of the image patch. Step D: Place the final segmentation results of the image patches into the corresponding positions of the image to be segmented according to their original positions when they were cut, and combine them to form the complete segmentation result of the image to be segmented.

Images