Three-dimensional medical image segmentation method and device

Abstract

The invention discloses a three-dimensional medical image segmentation method and device, and belongs to the technical field of computer vision and medical image processing, and the method comprises the steps: obtaining a to-be-segmented human spleen CT image, and constructing a data set which comprises a training set, a test set and a verification set; according to task requirements and data set features, corresponding models and parameters thereof are selected for instantiation; performing image preprocessing on the training set and the verification set, and training the selected model; reasoning the test set data by using the trained model, and outputting a prediction result; and performing channel conversion on a prediction result, and performing visualization of three-dimensional medical image segmentation. According to the method, the SwinCAS-UNETR network is provided and evaluated, it is verified that the network has remarkable effectiveness and superiority in MSD human spleen CT image three-dimensional segmentation tasks, medical image information with large foreground size changes can be better dealt with, and segmentation precision and model efficiency are improved.

Description

Technical Field

[0001] The invention relates to a method and a device for three-dimensional medical image segmentation, and belongs to the technical field of computer vision and medical image processing. Background Art

[0002] In the medical field, image segmentation technology plays a vital role in clinical diagnosis and is also the basis for other medical image processing methods. With the continuous development of big data, deep learning algorithms and high-performance computing technology, deep learning has made major breakthroughs in some of the most difficult problems in the field of artificial intelligence and has been widely used in various fields of society. With the improvement of computer hardware performance, image segmentation technology based on deep learning has been widely used in medical image processing. This technology can not only improve the accuracy and efficiency of medical image analysis, but also help doctors better understand the condition and develop more effective treatment plans.

[0003] In biomedical research, high-precision image segmentation can help researchers better understand the biological characteristics of diseases and promote the development of medical frontiers. Therefore, 3D segmentation has numerous applications in medical image segmentation, including tumor identification and organ localization for diagnostic purposes, which is usually solved by using a U-Net similar to the encoder-decoder architecture, where the encoder generates a hierarchical low-dimensional representation of the 3D image and the decoder maps this learned representation to voxel segmentation. Early CNN-based methods used convolution and deconvolution in the encoder and decoder, respectively, but may have difficulty in obtaining accurate segmentation results due to their limited local receptive fields. Swin UNETR designs a hybrid architecture to combine local convolution and global attention. Swin UNETR aims to design hybrid blocks for the encoder and decoder subnetworks. The self-attention mechanism of the SwinTransformer included in the Swin UNETR network allows the model to interact between different positions and capture dependencies, but its own self-attention mechanism plays a role in feature extraction, and no further screening is performed on the extracted features.

[0004] The deep learning framework PyTorch can solve the above problems. However, the deep learning framework PyTorch has the following shortcomings in the task of medical image segmentation: 1) Lack of built-in tools for medical images and insufficient data format support: PyTorch natively only supports basic 2D image formats (such as JPEG, PNG, etc.), while commonly used formats for medical images such as NIfTI, DICOM, etc. require the use of third-party libraries (such as SimpleITK or Nibabel) for processing. For the support of 3D medical images, PyTorch's functions are relatively limited. 2) Insufficient support for 3D convolutional networks: PyTorch supports 3D convolutional networks, but its efficiency is lower than that of optimized libraries dedicated to the field of medical imaging. Medical images usually use 3D data, and the training and inference efficiency of the model is crucial, especially when processing large-scale 3D data. 3) Memory and computing resource requirements: Resource bottlenecks for processing large-size medical images: Medical image data is usually large in size, especially 3D data (such as CT or MRI scans), which requires extremely high memory and computing resources. PyTorch's native memory management and data loading mechanisms may not be optimal when processing such large data sets, and memory overflow or performance bottlenecks may easily occur, requiring users to optimize the data processing process themselves. Summary of the invention

[0005] In order to solve the above problems, the present invention proposes a method and device for three-dimensional medical image segmentation, which can improve the efficiency of model training and simplify the processing flow of medical image data.

[0006] The technical solution adopted by the present invention to solve the technical problem is:

[0007] In a first aspect, an embodiment of the present invention provides a method for segmenting a three-dimensional medical image, comprising the following steps:

[0008] Step S1, obtaining a human spleen CT image to be segmented and constructing a data set, wherein the data set includes a training set, a test set and a validation set;

[0009] Step S2: Select the corresponding model and its parameters for instantiation according to the task requirements and data set characteristics;

[0010] Step S3, performing image preprocessing on the training set and the validation set, and training the selected model to obtain a trained model;

[0011] Step S4, use the trained model to infer the test set data and output the prediction results; if the trained model is a SwinCAS-UNETR network, use the sliding window inference method to perform inference, perform Sigmoid function normalization, and use the threshold method to output the prediction results;

[0012] Step S5, performing channel conversion on the prediction results and performing visualization of the three-dimensional medical image segmentation.

[0013] As a possible implementation manner of this embodiment, the training set includes training image data and training annotation data.

[0014] As a possible implementation of this embodiment, the model includes a V-Net model, a U-Net model, a SwinUNETR model and a SwinCAS-UNETR network model.

[0015] As a possible implementation of this embodiment, the process of establishing the SwinCAS-UNETR network model includes the following steps:

[0016] Add channel attention mechanism to each layer of the decoder part of Swin UNETR network;

[0017] Embed a deep supervision network in the upsampling of the Swin UNETR network;

[0018] Introducing additional loss functions at the intermediate levels of the deep supervised network.

[0019] As a possible implementation manner of this embodiment, the image preprocessing includes loading images and labels, channel priority, intensity scaling, background cropping, direction adjustment, and pixel spacing adjustment.

[0020] As a possible implementation of this embodiment, step S4, using the trained model to infer the test set data and outputting the prediction result, includes the following steps:

[0021] The original 3D data is cropped into sub-images of size (96, 96, 96), and four cropped samples are generated for each data.

[0022] The image to be inferred is cut into this size traversally and inference is performed sequentially to obtain the segmentation result of the entire CT image;

[0023] If the trained model is the SwinCAS-UNETR network, the sliding window inference method is used for inference, and the segmentation results are normalized by the Sigmoid function:

[0024]

[0025] Among them, x is the input value, σ(x) is the normalization result of x;

[0026] Use the threshold method to convert the output of the Sigmoid function into a binary prediction result:

[0027]

[0028] Among them, P(yi=1|xi) is the pixel point The probability of belonging to the target category, τ is the preset threshold.

[0029] As a possible implementation of this embodiment, the method for segmenting a three-dimensional medical image further includes the following steps:

[0030] Step S6, calculate the Dice coefficient and IoU value of the prediction result, and evaluate the prediction result.

[0031] As a possible implementation of this embodiment, step S6, calculating the Dice coefficient and IoU value of the prediction result and evaluating the prediction result, includes the following steps:

[0032] Calculate the Dice coefficient and use it to measure the overlap between the predicted segmentation and the true segmentation;

[0033] Calculate the intersection over union (IoU) and use it to measure the ratio of the intersection area to the union area between the predicted segmentation and the true segmentation;

[0034] The Dice coefficient is calculated as:

[0035]

[0036] The calculation formula of the intersection over union (IoU) is:

[0037]

[0038] Among them, P represents the predicted binary area, and G represents the real binary area.

[0039] In a second aspect, an embodiment of the present invention provides a device for segmenting a three-dimensional medical image, comprising:

[0040] A data set construction module is used to obtain a human spleen CT image to be segmented and construct a data set, wherein the data set includes a training set, a test set and a validation set;

[0041] The model selection module is used to select the corresponding model and its parameters for instantiation according to the task requirements and data set characteristics;

[0042] The model training module is used to perform image preprocessing on the training set and the validation set, and train the selected model to obtain the trained model;

[0043] The data inference module is used to use the trained model to infer the test set data and output the prediction results. If the trained model is the SwinCAS-UNETR network, the sliding window inference method is used for inference, and the Sigmoid function is normalized, and the threshold method is used to output the prediction results.

[0044] The visualization module is used to perform channel conversion on the prediction results and visualize the 3D medical image segmentation.

[0045] As a possible implementation of this embodiment, the three-dimensional medical image segmentation device further includes:

[0046] The result evaluation module is used to calculate the Dice coefficient and IoU value of the prediction results and evaluate the prediction results.

[0047] The beneficial effects of the technical solution of the embodiment of the present invention are as follows:

[0048] A method for three-dimensional medical image segmentation according to the technical solution of an embodiment of the present invention comprises the following steps: step S1, obtaining a human spleen CT image to be segmented and constructing a data set, wherein the data set comprises a training set, a test set and a validation set; step S2, selecting a corresponding model and its parameters for instantiation according to task requirements and data set characteristics; step S3, performing image preprocessing on the training set and the validation set, and training the selected model to obtain a trained model; step S4, using the trained model to infer the test set data and output a prediction result; if the trained model is a SwinCAS-UNETR network, a sliding window inference method is used for inference, and Sigmoid function normalization is performed, and a threshold method is used to output a prediction result; step S5, channel conversion is performed on the prediction result, and visualization of the three-dimensional medical image segmentation is performed. The present invention improves the Swin UNETR network structure, proposes a SwinCAS-UNETR network, and evaluates the SwinCAS-UNETR network, so that it achieves the best effect in the test set, verifies the effectiveness and superiority of the SwinCAS-UNETR network in the MSD human spleen CT image three-dimensional segmentation task, and can better cope with medical image information with large foreground size changes; the present invention develops a complete set of three-dimensional medical image segmentation technology in a unified MONAI framework, which improves segmentation accuracy and model efficiency.

[0049] A three-dimensional medical image segmentation device according to the technical solution of an embodiment of the present invention has the same beneficial effects as a three-dimensional medical image segmentation method according to the technical solution of an embodiment of the present invention. BRIEF DESCRIPTION OF THE DRAWINGS

[0050] Figure 1is a flow chart of a method for three-dimensional medical image segmentation according to an exemplary embodiment;

[0051] Figure 2 is a schematic structural diagram of a device for three-dimensional medical image segmentation according to an exemplary embodiment;

[0052] Figure 3 is a schematic diagram of a SwinCAS-UNETR network structure according to an exemplary embodiment;

[0053] Figure 4 is a schematic diagram of a channel attention module according to an exemplary embodiment;

[0054] Figure 5 is a schematic diagram of a deep supervision network structure according to an exemplary embodiment;

[0055] Figure 6 The figure is a schematic diagram of a reasoning method of a sliding window reasoning method according to an exemplary embodiment. DETAILED DESCRIPTION

[0056] In order to more clearly illustrate the technical features of the solution of the present invention, the present invention is described in detail below through specific implementation methods and in conjunction with the accompanying drawings.

[0057] like Figure 1 As shown, a method for three-dimensional medical image segmentation provided by an embodiment of the present invention includes the following steps:

[0058] Step S1, obtaining a human spleen CT image to be segmented and constructing a data set, wherein the data set includes a training set, a test set and a validation set;

[0059] Step S2: Select the corresponding model and its parameters for instantiation according to the task requirements and data set characteristics;

[0060] Step S3, performing image preprocessing on the training set and the validation set, and training the selected model to obtain a trained model;

[0061] Step S4, use the trained model to infer the test set data and output the prediction results; if the trained model is a SwinCAS-UNETR network, use the sliding window inference method to perform inference, perform Sigmoid function normalization, and use the threshold method to output the prediction results;

[0062] Step S5, performing channel conversion on the prediction results and performing visualization of the three-dimensional medical image segmentation.

[0063] As a possible implementation manner of this embodiment, the training set includes training image data and training annotation data.

[0064] As a possible implementation of this embodiment, the model includes a V-Net model, a U-Net model, a SwinUNETR model and a SwinCAS-UNETR network model.

[0065] As a possible implementation of this embodiment, the process of establishing the SwinCAS-UNETR network model includes the following steps:

[0066] Add channel attention mechanism to each layer of the decoder part of Swin UNETR network;

[0067] Embed a deep supervision network in the upsampling of the Swin UNETR network;

[0068] Introducing additional loss functions at the intermediate levels of the deep supervised network.

[0069] As a possible implementation manner of this embodiment, the image preprocessing includes loading images and labels, channel priority, intensity scaling, background cropping, direction adjustment, and pixel spacing adjustment.

[0070] As a possible implementation of this embodiment, step S4, using the trained model to infer the test set data and outputting the prediction result, includes the following steps:

[0071] The original 3D data is cropped into sub-images of size (96, 96, 96), and four cropped samples are generated for each data.

[0072] The image to be inferred is cut into this size traversally and inference is performed sequentially to obtain the segmentation result of the entire CT image;

[0073] If the trained model is the SwinCAS-UNETR network, the sliding window inference method is used for inference, and the segmentation results are normalized by the Sigmoid function:

[0074]

[0075] Among them, x is the input value, σ(x) is the normalization result of x;

[0076] Use the threshold method to convert the output of the Sigmoid function into a binary prediction result:

[0077]

[0078] Among them, P(yi=1|xi) is the pixel point The probability of belonging to the target category, τ is the preset threshold.

[0079] As a possible implementation of this embodiment, the method for segmenting a three-dimensional medical image further includes the following steps:

[0080] Step S6, calculate the Dice coefficient and IoU value of the prediction result, and evaluate the prediction result.

[0081] As a possible implementation of this embodiment, step S6, calculating the Dice coefficient and IoU value of the prediction result and evaluating the prediction result, includes the following steps:

[0082] Calculate the Dice coefficient and use it to measure the overlap between the predicted segmentation and the true segmentation;

[0083] Calculate the intersection over union (IoU) and use it to measure the ratio of the intersection area to the union area between the predicted segmentation and the true segmentation;

[0084] The Dice coefficient is calculated as:

[0085]

[0086] The calculation formula of the intersection over union (IoU) is:

[0087]

[0088] Among them, P represents the predicted binary area, and G represents the real binary area.

[0089] like Figure 2 As shown, an embodiment of the present invention provides a three-dimensional medical image segmentation device, comprising:

[0090] A data set construction module is used to obtain a human spleen CT image to be segmented and construct a data set, wherein the data set includes a training set, a test set and a validation set;

[0091] The model selection module is used to select the corresponding model and its parameters for instantiation according to the task requirements and data set characteristics;

[0092] The model training module is used to perform image preprocessing on the training set and the validation set, and train the selected model to obtain the trained model;

[0093] The data inference module is used to use the trained model to infer the test set data and output the prediction results. If the trained model is the SwinCAS-UNETR network, the sliding window inference method is used for inference, and the Sigmoid function is normalized, and the threshold method is used to output the prediction results.

[0094] The visualization module is used to perform channel conversion on the prediction results and visualize the 3D medical image segmentation.

[0095] As a possible implementation of this embodiment, the three-dimensional medical image segmentation device further includes:

[0096] The result evaluation module is used to calculate the Dice coefficient and IoU value of the prediction results and evaluate the prediction results.

[0097] The present invention uses the MONAI (Medical Open Network for AI, an open source AI framework for medical research) framework to build a complete set of 3D medical image segmentation system (device), including:

[0098] (1) Dataset construction module.

[0099] The dataset can be a commonly used medical imaging dataset (either NIfTI or DICOM format). The dataset construction module includes a training image storage module (imagesTr) and a training annotation storage module (labelsTr). This facilitates the subsequent replacement and training of different datasets, making the process framework highly versatile and reusable, and can be used for 3D medical imaging training and segmentation of different datasets.

[0100] (2) Model selection module.

[0101] Build a general model selection module, which already includes V-Net, U-Net, Swin UNETR and the proposed SwinCAS-UNETR network models. Select the appropriate model according to the task requirements and dataset characteristics. According to the passed model name, select the corresponding model to instantiate and return it.

[0102] (3) Model training module.

[0103] In this module, the dataset in the dataset construction module is read, and then image preprocessing operations are performed on the training set and validation set, including loading images and labels, channel priority, intensity scaling, background cropping, direction adjustment, and pixel spacing adjustment. Then the loss function, optimizer, and evaluation index set during system development are used by default. Next, the model selected by the user in the detection model selection module is trained for 600 epochs, and then the model inference module is called at regular intervals to evaluate the validation set. If the current indicator is better than the historical best indicator, the current model weight is saved as the best model. After the training is completed, the best weight file (.pth file) obtained by training is automatically saved to the workspace.

[0104] (4) Model reasoning module.

[0105] After the model training is completed, the next step is to perform inference on the test data. Load the best model weights saved during the training process. Use the context manager to turn off gradient calculations to avoid using memory during inference. Then perform inference on the test data, pass the input data into the model, and obtain the prediction results. If the SwinCAS-UNETR network is selected, the sliding window inference method is used for inference, and finally the sigmoid function normalization and threshold method are used to output the results. After the inference is completed, the segmentation results are saved to the workspace.

[0106] (5)Visualization module.

[0107] The segmentation results were converted into channels and then input into the integrated ITK-SNAP software for 3D segmentation visualization.

[0108] (6) Result evaluation module

[0109] Calculate the Dice coefficient and IoU value of the model prediction results, save the results to the workspace and automatically create a CSV table file, and append the average value of each column of the calculation result to the end of the file to facilitate viewing of quantitative indicators.

[0110] In order to prove that the SwinCAS-UNETR network can effectively improve the segmentation accuracy, Task09, the ninth subtask spleen dataset in the MSD (Medical Segmentation Decathlon) dataset, is selected. The goal is to segment the spleen from a three-dimensional CT image, because the size of the foreground of this dataset varies greatly, as shown in Table 1:

[0111] Table 1 Image size statistics

[0112]

[0113]

[0114] Splenic disease is partially manifested as an abnormal increase in spleen volume on 3D CT images, which means that when abnormal spleen enlargement is detected after segmentation, it may indicate disease. Therefore, this dataset includes 3D images of normal-sized spleens and abnormally enlarged diseased pancreases, which can be used to test the segmentation effect of the model and evaluate whether the segmentation accuracy has been improved.

[0115] The Swin UNETR model is modified on a unified framework, and the SwinCAS-UNETR network is proposed. Figure 2As shown in the figure, (1) a channel attention module is added to the network structure of Swin UNETR. The channel attention mechanism targets the channel dimension and calculates the weight matrices of different channels according to the dependencies between channels, focusing the model's attention and training on more important channels, that is, further screening the extracted features to improve the accuracy of the segmentation results. (2) A deep supervision network is introduced into the original Swin UNETR network structure. By introducing additional loss functions in the middle layers of the network, the network is prompted to learn richer and more abstract features. The loss functions of these middle layers are optimized together with the loss function of the final task to enable the network to learn more discriminative feature representations. This effectively alleviates the problems of gradient vanishing and slow convergence during network training, which helps guide the network to more stable learning and better convergence. The loss functions at different positions (or different layers) are multiplied by different coefficients and then summed as the total loss function. The purpose of this is to balance the impact of different layers on model training, so that the loss of each layer can have a certain impact on the training of the model, without the loss of some layers being too large or too small to affect the overall training effect.

[0116] The components of the SwinCAS-UNETR network are as follows:

[0117] In the encoder stage, the input of the model is a 3D image of size H×W×D. First, non-overlapping patches of the input data of size H0×W0×D0 are created, and a patch partitioning layer is used to create windows of the required size to calculate self-attention. The encoded feature representation in the Swin transformer is fed to the CNN decoder at multiple resolutions through skip connections. The final segmentation output consists of 3 output channels corresponding to the three sub-regions ET, WT and TC. Self-attention is calculated in the non-overlapping windows created in the partitioning stage to efficiently establish a tag interaction model. A shift window mechanism is used in subsequent layers to evenly divide a 3D token into multiple regions, and the divided window regions are shifted in the next layer. These token sequences then enter the 4-layer Swin Transformer block for global feature extraction. Each Swin Transformer block contains two parts: a linear embedding block and a patch merging block. The calculation formula of the linear embedding block is:

[0118]

[0119] Among them, W-MSA and SW-MSA refer to the conventional and window segmentation multi-head self-attention modules, respectively; and Represent the output of W-MSA and SW-MSA respectively; MLP and LN represent layer normalization and multi-layer perceptron respectively. The calculation formula of the self-attention mechanism is:

[0120]

[0121] Where Q, K, and V represent query, key, and value, respectively; and d represents the size of the query and key.

[0122] The encoder has 4 levels, with a total of L = 8 layers. In order to maintain the hierarchical structure of the encoder, the patch merging layer is used to reduce the resolution of the feature representation by 1 / 2 in each dimension at the end of each stage, and then the merged features of each layer are extracted and sent to the decoder part. In addition, the patch merging layer groups and concatenates patches with a resolution of 2×2×2 to obtain a 4C-dimensional feature embedding. In order to keep consistent with the resolution of the decoder part, the linear layer is used to reduce the feature dimension of the representation to 2C.

[0123] In the decoder stage, the features of different resolutions extracted by the encoder are connected to the decoder through skip connections. These hidden features first pass through the Residual Block module. In the residual connection module, the network selectively transmits the difference between the input and output, thereby screening and filtering in the process of feature learning, filtering out some feature information other than the spleen contour, retaining some important features of the spleen area, and connecting the output with the input before passing it to the next layer, and then passing through the channel attention mechanism module designed by us, namely CAM. The features that pass through the CAM module are spliced ​​with the previous level after upsampling. The spliced ​​features are reshaped and filtered, and enter the CAM module again, and then upsampled using the deconvolution layer. The upsampled features are spliced ​​with the features of the previous level. Three deep supervision layers are set in the last three upsampling layers, outputting Loss1, Loss2 and Loss3 respectively, and outputting Loss4 at the end of the entire network. They are weighted and summed to realize the deep supervision network. Finally, a feature representation of size H×W×D×48 is obtained. The 3D CT spleen segmentation map is calculated using a 1×1×1 convolution layer and a sigmoid activation function. This is the complete process of a segmentation.

[0124] Adding a channel attention mechanism is to re-screen the extracted features. In the three-dimensional medical image segmentation problem, the feature map includes four dimensions: channel (C), depth (D), height (H), and width (W). This means that each feature map is three-dimensional in space, and there is also a channel dimension to represent different feature channels. In order to add a channel attention mechanism to the network, the transposed convolution operation is first implemented to upsample the input feature map to a larger size. Then the channel attention module accepts the input feature map and the feature map of the encoder layer and outputs the weighted feature map. The channel attention module is then embedded in the upsampling module. The channel attention module (CAM) is an attention mechanism for the channel dimension. Its purpose is to calculate the channel weight matrix based on the dependencies between different channels, so as to focus the model's attention and training on more important channels. Through this mechanism, the model can automatically learn which channels are more critical to the task, thereby improving the model performance. After adding the channel attention mechanism, the model will perform a screening at the feature extraction stage to highlight the important feature channels to enhance the sensitivity and accuracy of the task.

[0125] like Figure 4 As shown in Figure 1, the structure of the channel attention module can be divided into three main parts, each of which includes the steps of matrix reshaping, matrix multiplication and normalization of the input feature map, and finally obtains the channel attention matrix. First, the matrix is ​​reshaped, and the input feature map A∈ΓC×D×H×W is reshaped into ΓC×N, where N=D×H×W. This step merges the spatial dimensions (height, width, depth) in the feature map into one dimension for subsequent matrix operations. Then matrix multiplication is performed, and matrix multiplication is performed between the reshaped feature matrix A and its transpose. Finally, SoftMax normalization is performed, and the matrix obtained by multiplication is SoftMax operated to normalize it to a probability distribution to obtain the channel attention matrix G∈ΓC×C. The SoftMax function converts the value of each element into a probability value between 0 and 1, and ensures that the sum of each row after normalization is 1, thereby indicating the degree of influence of each channel on other channels. In this way, through the above three steps, the CAM structure can calculate the attention weights between channels, that is, the channel attention matrix G, and the calculation formula is as follows:

[0126]

[0127] In this matrix G ji In the example, i,j=1,2,3,...C, each element g jiIt represents the degree of influence of the i-th channel on the j-th channel. After that, matrix multiplication is performed on the channel attention matrix G and the input feature map A to obtain a matrix of size C×N, which represents the weighted sum of each channel at different positions. The result of the matrix multiplication is then reshaped into the same shape as the input feature map ΓC×D×H×W. The purpose of this is to maintain the same spatial structure as the input feature map for subsequent operations. Next, the reshaped result is multiplied by a scale coefficient α. This scale coefficient α is a parameter with an initial value of 0. Through learning, it gradually obtains a larger weight and is used to adjust the scale or range of the output. Finally, the result after multiplying by the scale coefficient is matrix summed with the input feature map A to obtain the final output B∈ΓC×D×H×W:

[0128]

[0129] This summation operation fuses the weighted features with the original features to obtain a more representative feature representation. Through this process, the input feature map A is weighted by the channel attention mechanism to obtain an output B that is more suitable for the task requirements. The introduction of the scale factor α allows the model to adaptively adjust the importance of the output during the learning process to better adapt to different tasks and data.

[0130] Deep supervision (also known as relay supervision) is a technique that introduces additional loss functions at different levels of the network. These additional loss functions are different from the loss function of the final task, but they are still based on the true label (GroundTruth). Compared with multi-task learning, the present invention focuses on the same task, and introduces additional loss functions at the intermediate level of the network to encourage the network to learn richer and more abstract features. These additional loss functions are usually introduced at different depths of the network, and feedback signals are introduced at different layers to guide the feature learning of different layers at the same time. During the training process, the loss functions of these intermediate levels are optimized together with the loss function of the final task so that the network learns more discriminative feature representations. This method can help alleviate the gradient vanishing problem, speed up the convergence of the network, and improve the generalization ability of the network. Then the loss functions at different positions (or different levels) are multiplied by different coefficients and then summed as the total loss function. The purpose of this is to balance the impact of different levels on model training, so that the loss of each level can have a certain impact on the training of the model, without the loss of some levels being too large or too small to affect the overall training effect.

[0131] The last three upsampling layers of the SwinCAS-UNETR network proposed in this invention set three deep supervision layers. Each deep supervision layer introduces a loss function to compare with the true label (Ground Truth) to generate features respectively. Figure 1 ,feature Figure 2 and Features Figure 3 , corresponding to the calculation output Loss1, Loss2 and Loss3, the entire network finally outputs the feature Figure 4 , and calculate the output Loss4 at the same time, and add them up to realize the deep supervision network, such as Figure 5 shown.

[0132] Specifically, these three deep supervision layers upsample the feature maps of the middle layers of the network to the size of the original input data through transposed convolution operations. In the first deep supervision layer, the first transposed convolution layer upsamples the size of the input feature map to a larger size, usually half the size of the original input. The spatial dimension of the spleen 3D CT image is 3, the number of channels of the input feature map is 48, and the number of channels of the output feature map is 2, which meets the needs of subsequent processing. The size of the convolution kernel is 3x3, and the step size of moving in each spatial dimension is 2. In the second deep supervision layer, two transposed convolution layers are combined together, and the second transposed convolution layer upsamples the output feature map of the previous layer to a larger size again. The first transposed convolution layer upsamples the input feature map from 96 channels to 48 channels, and then the second transposed convolution layer upsamples the feature map of these 48 channels to a larger size again. The third deep supervision layer has one more transposed convolution layer than the second deep supervision layer, and the third transposed convolution layer continues to upsample the output feature map of the previous layer to a larger size. While upsampling the feature maps to the size of the original input data, their outputs are used to calculate the loss function for deep supervision, outputting Loss1, Loss2, and Loss3 respectively.

[0133] Loss4 corresponds to the loss of the final output of the entire network, which is directly obtained from the final feature map (feature Figure 4 ) and the true label (GroundTruth), which plays the final supervisory role in the model.

[0134] The loss functions here all use the soft Dice loss function, and the voxel calculation formula is:

[0135]

[0136] Among them, G and Y are the unique hot encoding forms of the prediction results and the true labels, respectively.

[0137] Finally, the weighted sum is used to calculate the loss:

[0138] Loss=0.2×Loss1+0.2×Loss2+0.2×Loss3+Loss4

[0139] This formula is the combined weighted loss formula used in the present invention. Here, the losses of the first three deep supervision layers are given the same weight (0.2) to balance the learning process of the middle layer. The final output loss Loss4 directly affects the final performance of the network, so it is not weighted and the final output loss is directly brought into the calculation, that is, the weight is 1, and the final total loss Loss is obtained.

[0140] Figure 6 The reasoning method of the sliding window principle in the work of the present invention is demonstrated. In order to reduce the consumption of video memory and improve the efficiency of model training, the original three-dimensional data is cropped. The size of the original data is large, and there are images of a size of about (512,512,100) in the data set. Directly inputting such large-sized data into the model will lead to insufficient video memory and difficulty in effective training. Therefore, the original data is cropped into sub-images of a size of (96,96,96), and four cropped samples are generated for each data. Therefore, in the process of model reasoning, it is also necessary to traverse and crop the image to be inferred to this size, and perform reasoning in sequence to obtain the segmentation result of the entire CT image.

[0141] This processing method reduces the size of each training sample, reduces the usage of video memory, and provides more samples for model training, which helps to improve the generalization ability and performance of the model. The image block with a size of 512×512×100 represents the original CT image, and the small block A with a size of 96×96×96 represents the image sent to the model. The sliding window strategy is adopted in the reasoning process, and B, C and D are used to represent the sliding mode along the X-axis, Y-axis and Z-axis respectively. The sliding step is set to 28 pixels (for the X-axis and Y-axis), and for the Z-axis, the sliding step is set to 8 pixels. Since the sliding window may lead to multiple reasonings in overlapping areas, the present invention records and accumulates these reasoning results, and uses a three-dimensional matrix to record their cumulative times. The average evaluation result can be obtained by dividing the reasoning result by the three-dimensional matrix of the cumulative number of times.

[0142] Finally, these results are normalized using the Sigmoid function. The Sigmoid function is a commonly used activation function that has the characteristic of mapping input values ​​between 0 and 1. It is usually used in binary classification tasks to convert the original output of the model into a probability form. This normalization process allows the output to be interpreted as the probability that a sample belongs to the positive category. The formula of the Sigmoid function is as follows:

[0143]

[0144] Where x is the input value. The Sigmoid function maps the input value to the range of (0,1) by performing exponential transformation and normalization on the input value. When x approaches positive infinity, e -As x approaches 0, σ(x) approaches 1. As x approaches negative infinity, e - x approaches positive infinity, so σ(x) approaches 0. When x = 0, σ(x) = 1 / 2.

[0145] In the image segmentation task, we can use a probability map to represent the probability of each pixel belonging to the target category. This probability map can be represented by P(y=1|x), where x is the original image and y=1 represents the target category. Similarly, P(y=0|x)=1-P(y=1|x) represents the probability that the pixel belongs to the background category.

[0146] Therefore, we can regard the image segmentation task as a pixel-level binary classification task, that is, to classify each pixel into two categories to determine whether it belongs to the target category. In this way, we can use the Sigmoid function to map the original output of the model to the (0,1) interval, indicating the probability that each pixel belongs to the target category.

[0147] Specifically, for each pixel’s original output xi, after being processed by the Sigmoid function, the probability that the pixel belongs to the target category is obtained P(yi=1|xi)=σ(xi), where is the Sigmoid function.

[0148] Therefore, in the image segmentation task, we can understand the output of the model as a probability map, where the value of each pixel represents the probability that the pixel belongs to the target category. By setting a threshold, we can convert the probability Figure 2 The final segmentation result is obtained by quantization. Usually, pixels greater than the threshold are classified as target categories, and pixels less than the threshold are classified as background categories.

[0149] In summary, in the image segmentation task, the output of the model can be interpreted as the probability that each pixel belongs to the target category through the Sigmoid function, thus providing an intuitive and explainable processing method.

[0150] The output of the Sigmoid function is converted into a binary prediction result using the threshold method as the final output of the model. The threshold method is a common method for converting a probability map or probability value into a binary prediction result. In the image segmentation task, the threshold method can be used to classify each pixel in the probability map into a target category or a background category. The threshold method compares the probability value of each pixel with a pre-set threshold to determine the final classification result. The formula of the threshold method is as follows:

[0151]

[0152] Among them, P(yi=1|xi) is the probability that the pixel xi belongs to the target category, and τ is the preset threshold. In the threshold method, a threshold τ needs to be set first. The threshold can be adjusted according to the task requirements and data characteristics. Usually, choosing a suitable threshold can enable the model to achieve higher accuracy on the training data. For the probability value P(yi=1|xi) of each pixel, compare it with the preset threshold τ. If the probability value is greater than or equal to the threshold τ, the pixel is classified as the target category, that is, Prediction(xi)=1. If the probability value is less than the threshold τ, the pixel is classified as the background category, that is, Prediction(xi)=0. Repeat the above steps, perform threshold comparison on each pixel in the probability map, and finally obtain a binary prediction result. The prediction result indicates whether each pixel is a target category or a background category, which is the final output of the model.

[0153] Through the threshold method, the probability map can be converted into a binary prediction result, thereby realizing the segmentation of the target category in the image segmentation task.

[0154] The strategy for model training is as follows:

[0155] (1) Use DataLoader to load the training and validation datasets and set an appropriate batch size.

[0156] (2) The model is trained in a loop for a total of 600 epochs. In each epoch, the training dataset is looped through to obtain the input image and label, which are passed into the model for forward propagation, the loss is calculated, and the model parameters are updated through backpropagation.

[0157] (3) Record the average loss of each epoch and use TensorBoard to visualize the loss curve.

[0158] (4) Perform validation at regular intervals to evaluate the performance of the model on the validation set. Use the sliding window inference method for inference and calculate the Dice index as the evaluation criterion. If the Dice index of the current model on the validation set is better than the historical best index, the weight of the model is saved.

[0159] (5) Using the Adam optimizer, it has the ability to adjust the learning rate adaptively and can converge effectively in most cases.

[0160] (6) Using DiceLoss as the loss function is suitable for segmentation tasks and can handle the problem of category imbalance.

[0161] (7) During the training process, the Dice indicator was used as the performance evaluation criterion to guide the direction of model training, and the optimal model weights were selected for preservation during the verification phase.

[0162] (8) Use TensorBoard to record the loss changes and verification indicators during the training process to facilitate monitoring of model training status.

[0163] The present invention selects two commonly used evaluation indicators, Dice coefficient (DSC) and intersection over union (IoU), in the evaluation indicator analysis.

[0164] In order to improve the segmentation quality, the present invention adds a channel attention mechanism to the original structure of Swin UNETR, further screens the extracted features, introduces deep supervision technology, and proposes the SwinCAS-UNETR network.

[0165] In order to verify the effectiveness of the channel attention mechanism and deep supervision technology in improving the network segmentation performance and evaluate the impact of various components or technologies in the model on the overall performance of the model, an ablation experiment was conducted using the Swin UNETR network as a benchmark. First, the channel attention mechanism was introduced into the Swin UNETR network alone to ensure the benefits of the mechanism for network training. Then, the case of only introducing deep supervision technology into the Swin UNETR network was tested. Finally, the two modules of channel attention mechanism and deep supervision technology were added to the network at the same time to verify that the combination of the two can make the network segmentation effect better. The experimental results are shown in Table 2.

[0166] Table 2 Ablation experiment results

[0167] Scenario DSC IoU Swin UNETR 0.930380694 0.871144471 Swin UNETR+CAM 0.940618303 0.890531116 Swin UNETR+Deep Supervision 0.954596632 0.914225168 Swin UNETR+CAM+Deep Supervision(Ours) 0.962073512 0.927276235

[0168] When only the channel attention mechanism is introduced, the average DSC coefficient of the test set increases from 93.04% to 94.06%, and the average IoU increases from 87.11% to 89.05%. This result shows that the introduction of the channel attention mechanism enhances the network's attention to important features, thereby effectively improving the network's segmentation performance.

[0169] The average DSC coefficient of the model combined with only the deep supervision network in the test set reached 95.46%, and the average IoU reached 91.42%. This shows that the introduction of additional auxiliary loss functions in Swin UNETR helps alleviate the gradient vanishing problem and improve the stability and convergence speed of the network.

[0170] Finally, after introducing the channel attention mechanism and deep supervision network into the Swin UNETR network, the final converged model achieved the best results in the test set. The average DSC coefficient of the SwinCAS-UNETR network increased to 96.21%, an increase of about 3.2% compared to the baseline model, and the average IoU reached 92.73%, an increase of about 5.6% compared to the baseline model. This experiment proves that the SwinCAS-UNETR network proposed in the present invention has the best performance in the MSD human spleen CT image three-dimensional segmentation task.

[0171] Other segmentation networks are selected to conduct comparative experiments on the spleen dataset with the proposed SwinCAS-UNETR network. The networks involved in the comparative experiments are V-Net, U-Net and Swin UNETR.

[0172] In this experiment, the parameter settings and experimental methods of each network are completed in the MONAI image segmentation process framework, so the experimental environment can be guaranteed to be consistent. Table 3 shows the segmentation results of each network on the spleen data set in the comparative experiment. Through the results of the comparative experiment, the superiority of the SwinCAS-UNETR network structure proposed in the present invention can be further verified.

[0173] Table 3 Comparative test results

[0174] network DSC IoU V-Net 0.955750346 0.916131649 U-Net 0.922550599 0.857852240 Swin UNETR 0.930380694 0.871144471 The network of the present invention 0.962073512 0.927276235

[0175] On the test set, the average DSC coefficient of the U-Net network was 92.25%, the lowest performance, while the average DSC coefficients of Swin UNETR and V-Net reached 93.03% and 95.58%, respectively, both lower than the results of the SwinCAS-UNETR network. This shows that the SwinCAS-UNETR network has higher segmentation accuracy in the spleen 3D segmentation task.

[0176] Comparing the average IoU values ​​of each network, the average IoU values ​​of V-Net, U-Net and Swin UNETR networks are between 85.79% and 91.61%, while the average IoU value of SwinCAS-UNETR network reaches 92.73%, which is significantly better than other networks. This shows that SwinCAS-UNETR network can more accurately capture the boundary information of spleen images and improve the accuracy of segmentation results.

[0177] The entire comparative experiment shows that the SwinCAS-UNETR network proposed in this invention has the best performance in the MSD spleen CT image three-dimensional segmentation task.

[0178] The present invention proposes the SwinCAS-UNETR network by improving the Swin UNETR network structure. Subsequently, the network proposed by the present invention is evaluated through a large number of experiments and analyses. The final converged model achieves the best effect in the test set, with the average DSC coefficient increased to 96.21% and the average IoU value reaching 92.73%. The effectiveness and superiority of the SwinCAS-UNETR network in the MSD human spleen CT image three-dimensional segmentation task are verified, indicating that it can better cope with medical image information with large foreground size changes.

[0179] The present invention develops a complete 3D medical image segmentation system in a unified framework (MONAI framework), while improving segmentation accuracy and model efficiency. The developed 3D medical image segmentation system can be used for different medical image data sets (NIfTI, DICOM format) to train, segment and evaluate the model. The system model selection module includes V-Net, U-Net, Swin UNETR and the proposed SwinCAS-UNETR network, and the evaluation module includes DSC and IoU coefficients.

[0180] The present invention uses the MONAI (Medical Open Network for AI) framework for software development, and the developed three-dimensional medical image segmentation system has the following advantages:

[0181] (1) Natively supports common medical imaging formats (NIfTI, DICOM formats), suitable for processing 3D medical imaging data such as CT and MRI, and can directly replace data sets.

[0182] (2) Supports preprocessing and data enhancement (such as rotation, flipping, normalization, and cropping) of three-dimensional medical images, and is particularly suitable for processing large-size 3D data such as CT and MRI.

[0183] (3) It provides spatial transformation functions commonly used in the field of medical imaging, supports the alignment and registration of multimodal images, and makes it easier to process data from different scanners or imaging methods.

[0184] (4) It integrates NVIDIA's Apex library, supports automatic mixed-precision training, reduces video memory usage, and improves training speed. It is particularly suitable for medical imaging tasks with large amounts of data. It also has built-in support for multiple GPUs, making it convenient to train large models in high-performance computing environments.

[0185] (5) Integration with the ITK-SNAP medical imaging tool makes it easier for developers to visualize segmentation results directly on these platforms.

[0186] The present invention first provides a complete three-dimensional medical image segmentation process, and then conducts an in-depth analysis of the Swin UN-ETR architecture and improves its network structure. A channel attention mechanism is introduced in the decoder part to further screen the extracted features. At the same time, a deep supervision network is embedded in the upsampling process. By introducing additional loss functions, the network is helped to learn richer and more abstract feature representations, accelerate the convergence speed of the model, and improve the generalization ability. Finally, in order to reduce video memory consumption and improve training efficiency, a sliding window reasoning method is adopted, and the final segmentation output is obtained by accumulating and normalizing the reasoning results. These improvement measures together constitute the SwinCAS-UNETR network structure proposed by the present invention, which improves the performance and efficiency of the model in three-dimensional medical image segmentation tasks.

[0187] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention rather than to limit it. Although the present invention has been described in detail with reference to the above embodiments, ordinary technicians in the relevant field should understand that the specific implementation methods of the present invention can still be modified or replaced by equivalents. Any modification or equivalent replacement that does not depart from the spirit and scope of the present invention should be covered within the scope of protection of the claims of the present invention.

Claims

1. A method for three-dimensional medical image segmentation, characterized in that: The steps include: Step S1, obtaining a human spleen CT image to be segmented and constructing a data set, wherein the data set includes a training set, a test set and a validation set; Step S2: Select the corresponding model and its parameters for instantiation according to the task requirements and data set characteristics; Step S3, performing image preprocessing on the training set and the validation set, and training the selected model to obtain a trained model; Step S4, use the trained model to infer the test set data and output the prediction results; if the trained model is a SwinCAS-UNETR network, use the sliding window inference method to perform inference, perform Sigmoid function normalization, and use the threshold method to output the prediction results; Step S5, performing channel conversion on the prediction results and performing visualization of the three-dimensional medical image segmentation.

2. The method for three-dimensional medical image segmentation according to claim 1, characterized in that: The training set includes training image data and training annotation data.

3. The method for three-dimensional medical image segmentation according to claim 1, characterized in that: The models include V-Net model, U-Net model, Swin UNETR model and SwinCAS-UNETR network model.

4. The method for three-dimensional medical image segmentation according to claim 3, characterized in that: The establishment process of the SwinCAS-UNETR network model includes the following steps: Add channel attention mechanism to each layer of the decoder part of Swin UNETR network; Embed a deep supervision network in the upsampling of the Swin UNETR network; Introducing additional loss functions at the intermediate levels of the deep supervised network.

5. The method for three-dimensional medical image segmentation according to claim 1, characterized in that: The image preprocessing includes loading images and labels, channel priority, intensity scaling, background cropping, orientation adjustment, and pixel spacing adjustment.

6. The method for three-dimensional medical image segmentation according to claim 1, characterized in that: The step S4, using the trained model to infer the test set data and output the prediction result, includes the following steps: The original 3D data is cropped into sub-images of size (96, 96, 96), and four cropped samples are generated for each data. The image to be inferred is cut into this size traversally and inference is performed sequentially to obtain the segmentation result of the entire CT image; If the trained model is the SwinCAS-UNETR network, the sliding window inference method is used for inference, and the segmentation results are normalized by the Sigmoid function: Among them, x is the input value, σ(x) is the normalization result of x; Use the threshold method to convert the output of the Sigmoid function into a binary prediction result: Among them, P(yi=1|xi) is the pixel point The probability of belonging to the target category, τ is the preset threshold.

7. The method for three-dimensional medical image segmentation according to any one of claims 1 to 6, characterized in that: The following steps are also included: Step S6, calculate the Dice coefficient and IoU value of the prediction result, and evaluate the prediction result.

8. The method for three-dimensional medical image segmentation according to claim 7, characterized in that: The step S6, calculating the Dice coefficient and IoU value of the prediction result and evaluating the prediction result, includes the following steps: Calculate the Dice coefficient and use it to measure the overlap between the predicted segmentation and the true segmentation; Calculate the intersection over union (IoU) and use it to measure the ratio of the intersection area to the union area between the predicted segmentation and the true segmentation; The Dice coefficient is calculated as: The calculation formula of the intersection over union (IoU) is: Among them, P represents the predicted binary area, and G represents the real binary area.

9. A device for three-dimensional medical image segmentation, characterized in that: include: A data set construction module is used to obtain a human spleen CT image to be segmented and construct a data set, wherein the data set includes a training set, a test set and a validation set; The model selection module is used to select the corresponding model and its parameters for instantiation according to the task requirements and data set characteristics; The model training module is used to perform image preprocessing on the training set and the validation set, and train the selected model to obtain the trained model; The data inference module is used to use the trained model to infer the test set data and output the prediction results. If the trained model is the SwinCAS-UNETR network, the sliding window inference method is used for inference, and the Sigmoid function is normalized, and the threshold method is used to output the prediction results. The visualization module is used to perform channel conversion on the prediction results and visualize the 3D medical image segmentation.

10. The device for three-dimensional medical image segmentation according to claim 9, characterized in that: Also includes: The result evaluation module is used to calculate the Dice coefficient and IoU value of the prediction results and evaluate the prediction results.

Images