A three-dimensional medical image segmentation method and medium

Abstract

The application relates to a three-dimensional medical image segmentation method, comprising the following steps: (1) acquiring a non-labeled three-dimensional medical image sample, and performing data enhancement on the non-labeled three-dimensional medical image sample to obtain two groups of medical images; (2) initializing a self-supervised learning pre-training framework, wherein the self-supervised learning pre-training framework comprises a student network and a teacher network; (3) inputting the two groups of medical images into the student network and the teacher network in the self-supervised learning pre-training framework, entering training iteration, and adjusting network parameters of the student network and the teacher network in each training iteration based on a preset composite loss function; and (4) completing a preset number of training iterations, obtaining a final student network, and using the final student network to perform a three-dimensional medical image segmentation task on a medical image to be segmented. The scheme improves the image segmentation precision of a downstream task.

Description

Technical Field

[0001] This invention relates to the field of computer vision technology in deep learning, and in particular to a three-dimensional medical image segmentation method and medium. Background Technology

[0002] In the field of medical image analysis, the accuracy of medical image segmentation is crucial for the construction and development of intelligent medical systems, especially in disease diagnosis, prognosis and treatment, and the formulation of postoperative rehabilitation plans.

[0003] In recent years, many medical image segmentation algorithms have emerged, significantly improving the segmentation performance of models. However, the excellent performance of these algorithms usually requires a large amount of carefully labeled data. Yet, in the field of medical image analysis, obtaining a sufficient amount of labeled 3D medical data is far more challenging than obtaining labeled natural images. This is due to the complexity and time-consuming nature of labeling 3D medical data, further demanding the broad involvement of medical professionals who contribute their deep domain expertise.

[0004] Therefore, researching self-supervised pre-training methods that can effectively utilize a large number of unlabeled 3D medical images and are tailored to the characteristics of 3D medical images is particularly important.

[0005] Therefore, there is an urgent need for a self-supervised pre-training method model tailored to the characteristics of 3D medical images to further improve the accuracy of automatic 3D medical image segmentation. Summary of the Invention

[0006] To address the aforementioned problems, the present invention aims to provide a three-dimensional medical image segmentation method. This method utilizes a pre-defined self-supervised learning pre-training framework to train a network model based on unlabeled three-dimensional medical image samples. This trained network model can capture higher-level semantic information without altering the network structure, thereby improving the image segmentation accuracy of downstream tasks.

[0007] To achieve the above objectives, the present invention adopts the following technical solution: In a first aspect, this application provides a three-dimensional medical image segmentation method, the method comprising: (1) Obtain unlabeled three-dimensional medical image samples and perform data augmentation on the unlabeled three-dimensional medical image samples to obtain two sets of medical images; (2) Initialize the self-supervised learning pre-training framework, which includes a student network and a teacher network; (3) Input the two sets of medical images into the student network and teacher network in the self-supervised learning pre-training framework respectively, enter the training iteration, and adjust the network parameters of the student network and teacher network in each training iteration based on the preset composite loss function. (4) Complete the preset number of training iterations to obtain the final student network, and use it for the task of three-dimensional medical image segmentation of the medical images to be segmented.

[0008] In one implementation of the present invention, the self-supervised learning pre-training framework includes a student network. Teacher Network ;in Student Network Includes encoder Projector head and decoder ; Teacher Network Includes encoder and projector head Its structure is the same as that of the encoder and projector head of the student network.

[0009] In one implementation method of the present invention, the characteristic is that the input unlabeled three-dimensional medical image sample is... Two different views are generated by applying different combinations of two data augmentations, denoted as... , ; Will and Sent to the student network respectively and teacher networks In the middle, both networks generate A probability distribution of dimension is denoted as . and ;

[0010] ; in It's a temperature parameter, within a fixed teacher network parameter. In this case, we learn to match these distributions by minimizing the cross-entropy loss; Parameters of student network Update based on the loss from the self-distillation module: , in .

[0011] In one implementation of the present invention, smoothing and sharpening operations are applied to the output of the teacher network. The sharpening operation adds a bias term to the output of the teacher network. ; Cluster center item Update using a moving exponential average strategy; The update rule is expressed as:

[0012] in It is a rate parameter. It refers to the batch size.

[0013] In one implementation of the present invention, during the training iteration, the input view is... A preset portion is masked, and the original image is predicted within the masked area; and a multi-scale mask image modeling loss is calculated based on the predicted image and the original image.

[0014] In one implementation of the present invention, the preset composite loss function is the self-distillation module loss and the multi-scale mask image modeling loss:

[0015] in The hyperparameters are used to adjust the weights of the two components.

[0016] In one implementation of this invention, after each training iteration, the teacher network updates its parameters based on the student network using a moving exponential average strategy. The specific calculation method is as follows: , in, It is a hyperparameter.

[0017] In a second aspect, a computer-readable storage medium is provided, wherein a computer program is stored in the computer-readable storage medium, and the computer program, when running, controls the device where the computer-readable storage medium is located to perform the three-dimensional medical image segmentation method described in the first aspect.

[0018] The present invention has the following advantages due to the adoption of the above technical solutions: 1. This invention ingeniously integrates the widely adopted self-supervised learning paradigm, significantly improving the model's feature extraction efficiency when processing unlabeled data. Through this innovative combination, the model can more effectively learn useful information from large amounts of unlabeled data, thus capturing key data features and improving overall learning performance even in the absence of explicit labels.

[0019] 2. This invention innovatively introduces a multi-scale module for the field of 3D medical image segmentation. This strategy enables the model to capture more detailed and richer semantic information when processing image data. By analyzing and integrating features at different scales, the model can better understand the complex structures in the image, thereby improving the segmentation accuracy during the fine-tuning stage.

[0020] 3. This invention is of great significance for improving the accuracy of medical image analysis, especially when processing highly heterogeneous and complex three-dimensional medical images, and can provide more reliable technical support for clinical diagnosis and treatment. Attached Figure Description

[0021] Figure 1 This is a flowchart illustrating the self-distillation contrastive learning algorithm upon which the three-dimensional medical image segmentation method proposed in this invention is based; Figure 2 This is a schematic diagram of the medical image segmentation model framework based on self-distillation contrastive learning in an embodiment of the present invention. Detailed Implementation

[0022] To make the objectives, technical solutions, and advantages of the embodiments of the present invention clearer, the technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some, not all, of the embodiments of the present invention. All other embodiments obtained by those skilled in the art based on the described embodiments of the present invention are within the scope of protection of the present invention.

[0023] To facilitate understanding of this invention, some self-supervised frameworks in the prior art will be explained and described first.

[0024] Typical self-supervised learning schemes include using agent tasks, contrastive learning, DINO, and MIM.

[0025] Agent tasks include designing specific tasks, such as puzzle and rotation prediction, for image context recovery.

[0026] Contrastive learning aims to gather positive samples and push away negative samples in the latent hidden space.

[0027] DINO and MIM are recently developed techniques. The former dynamically constructs a teacher network during training, thereby redefining knowledge distillation as an intrinsic self-supervised pre-training objective rather than a post-processing step. The latter shows significant potential by learning hidden embedding features by predicting masked parts of the input data. However, many self-supervised learning techniques, especially DINO, are primarily designed for natural images. Therefore, its direct application in the field of medical image analysis may affect model performance. This difference is caused by (1) the complexity of 3D medical image data, due to its higher dimensionality and diverse modalities; and (2) the fact that some modalities, such as magnetic resonance imaging and computed tomography, use imaging methods that differ significantly from those used in natural images. This inconsistency limits the direct application of these methods to medical images.

[0028] To address the limitation of existing technologies, which are primarily applied to natural images rather than 3D medical images, this invention researches a self-supervised pre-training method that can effectively utilize a large number of unlabeled 3D medical images while also taking into account the characteristics of 3D medical images.

[0029] This application provides a three-dimensional medical image segmentation method, the method comprising: (1) Obtain unlabeled three-dimensional medical image samples and perform data augmentation on the unlabeled three-dimensional medical image samples to obtain two sets of medical images; (2) Initialize the self-supervised learning pre-training framework, which includes a student network and a teacher network; (3) Input the two sets of medical images into the student network and teacher network in the self-supervised learning pre-training framework respectively, enter the training iteration, and adjust the network parameters of the student network and teacher network in each training iteration based on the preset composite loss function. (4) Complete the preset number of training iterations to obtain the final student network, and use it for the task of three-dimensional medical image segmentation of the medical images to be segmented.

[0030] This invention employs a pre-defined self-supervised learning pre-training framework to train a network model based on unlabeled 3D medical image samples. This trained network model can capture higher-level semantic information and improve the image segmentation accuracy of downstream tasks without altering the network structure.

[0031] The following is combined Figure 1 and Figure 2 In a more detailed embodiment of this application, the above method is described, and its technical advantages or effects are explained.

[0032] In a more detailed embodiment of the present invention, the processing steps of the method include: 1) The entire pre-training framework introduces two networks: a student network and a pre-training network. Teacher Network The student network includes an encoder. Projector head and decoder The teacher network includes an encoder. and projector head Its structure is the same as that of the encoder and projector head of the student network.

[0033] 2) Initialize the student network and teacher network parameters. Input unlabeled image data... Two different views are generated by applying different combinations of two data augmentations, denoted as... , .

[0034] 3) and Send in separately and In the middle, both networks generate A probability distribution of dimension is denoted as . and These probability distributions It is obtained by normalizing the network output using the softmax function. Formally, it can be expressed as:

[0035]

[0036] in It is a temperature parameter that controls the sharpness of the output distribution. With fixed teacher network parameters... In this case, matching these distributions is learned by minimizing the cross-entropy loss. The parameters of the student network... The update is as follows: , in .

[0037] 4) To prevent model learning collapse, smoothing and sharpening operations are applied to the output of the teacher network. Smoothing helps avoid a single dimension dominating the output distribution, promoting a more uniform distribution and preventing model collapse, while sharpening has the opposite effect. By applying these two operations simultaneously, model instability is effectively prevented when dynamically constructing the teacher network. The sharpening operation relies on first-order batch statistics and can be conceptualized as adding a bias term to the output of the teacher network. .

[0038] Cluster center item The method uses an Exponential Moving Average (EMA) strategy for updates to ensure stable performance across different training batch sizes. The update rule is expressed as:

[0039] in It is a rate parameter. This refers to the batch size. A lower temperature coefficient is used in the teacher softmax normalization. The output sharpening is achieved by using a certain value.

[0040] 5) Input View A portion of the image is masked, and the original image is predicted within the masked area. For view The feature map after passing through a Swin Transformer layer of depth 2 has a size of [size missing]. Half of the input occluded region. To transform the occluded region, this invention replaces each occluded patch with a learnable mask token vector. The dimensions of these token vectors match the dimensions of the other visible patch representations after patch embedding. To determine the occluded region, this invention employs a random masking strategy aligned with the patch. The processed image is denoted as... and sent into Reconstructed images of different scales are obtained from this process, denoted as... and The mean absolute error between the original view and the actual view is calculated and expressed as:

[0041] 6) Subsequently, based on the symmetry of the operation, the present invention provides another view. The same processing procedure is used. The formula for the entire multi-scale mask image modeling process is as follows: 7) The loss function used in the pre-training phase is a combination of the self-distillation module loss and the multi-scale mask image modeling loss, specifically expressed as:

[0042] in The hyperparameters are used to adjust the weights of the two components.

[0043] 8) After each training iteration, the teacher network updates its parameters based on the student network using the EMA method. The specific calculation method is as follows: , in, It is a hyperparameter that controls how much information teachers obtain from the student network.

[0044] 9) Finally, the pre-trained student encoder weights are used to initialize the backbone network in the downstream fine-tuning task, thereby completing the segmentation task fine-tuning on the specific dataset.

[0045] In a specific example, during the execution process 1) Download the data for the pre-training phase. These images are from four public CT image datasets.

[0046] These datasets include: HNSCC, with 1297 head and neck cases; LUNA16, with 888 lung cases; COVID-19, with 250 lung cases; and FLARE21, with 411 abdominal cases. During the pre-training phase, this invention uses only unlabeled data and does not use its actual labeled content. 2) For the entire pre-training dataset, this invention applies a max-min normalization operation to each 3D image, then removes the background by cropping the image based on the foreground, and finally, for each input data, ... Eight samples were randomly cut from the given dimensions; 3) The experiments in this example were conducted on a machine equipped with an Intel Core i9-10900X CPU and an NVIDIA GeForce RTX 3090 GPU, and the code was implemented using PyTorch. This example selected the Swin UNTER encoder as the pre-trained encoder, and the network depth was set to... The feature dimension is set to Window size set to The AdamW optimizer was used throughout the pre-training experiments, and the learning rate was set to [value missing] during pre-training. The maximum training cycle during the pre-training period is set to EMA update parameters ; 4) Initialize the weights of the student network and the teacher network to be equal. For each input image, apply a random combination of different data augmentations to the input sample to generate two views, and simultaneously generate masked images corresponding to the two views, where the mask size is... The occlusion rate is set to ; 5) Feed the two views into the teacher network and student network respectively to obtain the corresponding feature distributions. The output of the teacher network should be sharpened before softmax smoothing, and then the cross-entropy loss between the two networks should be calculated. ; 6) The masked images corresponding to the two views are sequentially fed into the image reconstruction network. In the process, reconstructed images at different scales are obtained, and the mean absolute loss between the reconstructed images and the input view is calculated respectively. ; 7) Calculate the composite loss function After each training iteration, the teacher network updates its parameters based on the student network using a moving exponential averaging strategy. 8) After completing the maximum training cycle, only the weights of the student network encoder are saved, which are used to initialize the backbone network encoder during the fine-tuning phase.

[0047] 9) Download the three datasets BTCV, MSD-Spleen and BraTS 2023 for the fine-tuning phase, which correspond to multi-organ segmentation, spleen segmentation and brain tumor segmentation, respectively; 10) Regarding data preprocessing in the fine-tuning stage, for the BTCV dataset, this invention first resamples the images and labels to... The image was then scaled in intensity range to reduce the range. Signals within the range are normalized to The background was removed through cropping. Finally, based on the label distribution characteristics, for each input data point, the negative-to-positive ratio was calculated. Randomly cut out sizes according to proportions Four samples. For the MSD-Spleen dataset, except for the intensity scaling range... Apart from that, the remaining preprocessing steps are consistent with the BTCV dataset. For the BraTS 2023 dataset, since each sample in this dataset consists of four modalities that have been resampled and aligned, this paper first combines the four modalities of each sample into a single four-channel 3D dataset. Next, the images and corresponding labels are cropped to remove the background. Then, the samples and corresponding labels are further cropped or spatially augmented to a uniform size. .

[0048] 11) Since these datasets do not publicly disclose the true labels of the test sets, this invention divides the data into two parts: the training set (…). ) and validation set ( For the three datasets, the following is specified: During the fine-tuning phase, for the BTCV dataset... Training cycles; for the MSD-Spleen dataset, it is... Training cycle; for the BraTS2023 dataset: Training cycle; 12) Complete the fine-tuning experiment and save the model weights.

[0049] In step 5) above, the sharpening operation can be conceptualized as adding a bias term and a cluster center term to the output of the teacher network. The EMA method is also used for updates to ensure that the method has stable performance under different training batch sizes.

[0050] In step 6) above, one of the input views For example, its feature map at another scale is denoted as The size is Half of it. After applying a random masking strategy, the partially occluded view is denoted as... , Send in Two reconstructed images and features at different scales are generated, denoted as... and .calculate and , and The average absolute error between them, and then based on symmetry, the view Using the same operation, the mean absolute error is finally summed to obtain... .

[0051] In step 11) above, the BTCV dataset and the MSD-Spleen dataset are used during fine-tuning. Loss function. The BraTS 2023 dataset is used during fine-tuning. Loss function:

[0052]

[0053]

[0054]

[0055] In the above equation, It represents the total number of voxels, that is, the set of all voxels in a 3D image; This represents the total number of categories, indicating the different tissue or structure types that may appear in the image. For each voxel location... and categories , This indicates that the model predicts the category at this location. The probability output, and This indicates whether the voxel at the corresponding position in the actual label belongs to the category. Information. It is an indicator function, representing the position. Is the true category? .

[0056] To verify the effectiveness and practicality of the proposed method, fine-tuning experiments on the BTCV, MSD-Spleen, and BraTS2023 datasets are presented below. Table 1 shows the segmentation results on the validation set, with the Dice similarity coefficient as the metric. These results are compared with other mainstream self-supervised methods in this field, including DINO, SimMIM, HybridMIM, and SwinSSL. As shown in Table 1, although all self-supervised algorithms show improvement compared to the baseline, the fine-tuning results on all three datasets achieved the highest segmentation accuracy after using the pre-training algorithm proposed in this invention.

[0057]

[0058] Table 1 In one aspect of the embodiments of this application, a computer-readable storage medium is also provided, which stores a computer program that, when executed by a computer device, implements the method described in the embodiments of this application.

[0059] Those skilled in the art will clearly understand that, for the sake of convenience and brevity, the specific working processes of the systems, devices, and units described above can be referred to the corresponding processes in the foregoing method embodiments, and will not be repeated here. In the embodiments provided by this invention, it should be understood that the disclosed systems, apparatuses, and methods can be implemented in other ways. For example, the apparatus embodiments described above are merely illustrative; for instance, the division of the units described above is only a logical functional division, and in actual implementation, there may be other division methods. For example, multiple units or components may be combined or integrated into another system, or some features may be ignored or not executed. Furthermore, the coupling or direct coupling or communication connection shown or discussed may be an indirect coupling or communication connection between devices or units through some interfaces, and may be electrical, mechanical, or other forms.

[0060] The integrated units implemented as software functional units described above can be stored in a computer-readable storage medium. These software functional units, stored in a storage medium, include several instructions to cause a computer device (which may be a personal computer, server, or network device, etc.) or processor to execute some steps of the methods described in the various embodiments of the present invention. The aforementioned storage medium includes various media capable of storing program code, such as USB flash drives, portable hard drives, read-only memory (ROM), random access memory (RAM), magnetic disks, or optical disks.

[0061] The above description is merely a preferred embodiment of the present invention and is not intended to limit the present invention. Any modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of the present invention should be included within the scope of protection of the present invention.

Claims

1. A three-dimensional medical image segmentation method, characterized in that, The method includes: (1) Obtain unlabeled three-dimensional medical image samples and perform data augmentation on the unlabeled three-dimensional medical image samples to obtain two sets of medical images; (2) Initialize the self-supervised learning pre-training framework, which includes a student network. Teacher Network Among them, student networks Includes encoder Projector head and decoder Teacher Network Includes encoder and projector head Its structure is the same as the encoder and projector head of the student network; (3) Input the two sets of medical images into the student network and teacher network in the self-supervised learning pre-training framework respectively, enter the training iteration, and adjust the network parameters of the student network and teacher network in each training iteration based on the preset composite loss function. (4) Complete the preset number of training iterations to obtain the final student network, and use it for the task of three-dimensional medical image segmentation of the medical images to be segmented; Among them, the input unlabeled 3D medical image samples Two different views are generated by applying different combinations of two data augmentations, denoted as... , ; Will and Sent to the student network respectively and teacher networks In the middle, both networks generate A probability distribution of dimension is denoted as . and ; ； in It's a temperature parameter, within a fixed teacher network parameter. In this case, we learn to match these distributions by minimizing the cross-entropy loss; Parameters of student network Update based on the loss from the self-distillation module: ， in ; Smoothing and sharpening operations are applied to the output of the teacher network, where the sharpening operation adds a bias term to the output of the teacher network: ； Cluster center item Update using a moving exponential average strategy; The update rule is expressed as: in It is a rate parameter. It refers to the batch size; During the training iteration, the input view A portion of the preset image is masked, and the original image is predicted within the masked area. For view The feature map after the Swin Transformer layer has a size of [size missing]. Half of the original view is calculated, and the mean absolute error between the reconstructed images at different scales and the corresponding original views is calculated, and then compared with another view. Using the same processing, the multi-scale mask image modeling loss is obtained. ; The preset composite loss function is the self-distillation module loss and the multi-scale mask image modeling loss: in The hyperparameters used to adjust the weights of the two; After each training iteration, the teacher network updates its parameters based on the student network using a moving exponential average strategy. The specific calculation method is as follows: ， in, It is a hyperparameter.

2. A computer-readable storage medium, characterized in that, The computer-readable storage medium stores a computer program, which, when executed, controls the device containing the computer-readable storage medium to perform the three-dimensional medical image segmentation method of claim 1.

Images