Medical image segmentation self-pre-training method and system based on difficult sample mining mask auto-encoder

Abstract

The invention discloses a medical image segmentation self-pre-training method and system based on a difficult sample mining mask auto-encoder, and the method comprises the steps: collecting whole-body CT scanning images of a plurality of subjects, carrying out the preprocessing of the CT scanning images, and obtaining a plurality of three-dimensional sub-images with the same size; carrying out self-pre-training on a deep learning model HPM with an attention mechanism, wherein the deep learning model HPM comprises a block embedding layer, a position embedding layer, a student VIT network and a teacher VIT network; after training is finished, an encoder of the student VIT network is connected with a UNETR decoder, and a segmentation network model is constructed; and training the segmentation network model, inputting a new three-dimensional sub-image into the segmentation network model after training is completed, and finally outputting a 42-class segmentation result. According to the method, the accuracy of medical image segmentation can be improved through difficult sample mining and a self-pre-training mechanism.

Description

Technical Field

[0001] The present invention relates to the technical field of image segmentation, and in particular to a medical image segmentation self-pretraining method and system based on difficult sample mining mask autoencoder. Background Art

[0002] Medical image segmentation is an important technology in medical image processing, which aims to automatically identify and extract regions of interest or targets from medical images to help doctors diagnose and treat diseases more accurately and quickly. With the rapid development of deep learning technology, the accuracy of computer vision and image recognition has been significantly improved, which has greatly promoted the application of deep learning in the field of medical image segmentation. However, in the actual application of medical image processing, it is extremely difficult to obtain a large number of high-quality annotated samples due to the need to protect patient privacy and the fact that medical image classification usually requires the participation of professional radiologists. For deep learning models, especially those with greater depth, limited training samples may cause the model to overfit, thereby reducing its generalization ability and affecting the accuracy of the analysis results.

[0003] As an unsupervised learning method, Masked Autoencoder (MAE) provides effective pre-training for deep learning models, aiming to reduce the dependence of deep learning models on large amounts of labeled data and improve their generalization ability. MAE learns effective representations of data by masking parts of the image and training the deep learning model to reconstruct these masked parts. Although MAE has obvious advantages in learning medical image representations, its training process does not stand in the position of a teacher to determine which data samples are difficult to reconstruct. However, in fact, identifying and learning data samples that are difficult to reconstruct is particularly critical to improving the accuracy of medical image segmentation. In this way, deep learning models can have a deeper understanding of the complexity and subtle differences of medical images, thereby achieving higher accuracy in segmentation tasks. Summary of the invention

[0004] Purpose of the invention: The purpose of the present invention is to provide a medical image segmentation self-pretraining method and system based on difficult sample mining mask autoencoder, which can improve the accuracy of medical image segmentation through difficult sample mining and self-pretraining mechanism.

[0005] Technical solution: To achieve the above-mentioned purpose, the present invention provides a medical image segmentation self-pretraining method based on difficult sample mining mask autoencoder, comprising the following steps:

[0006] Step 1: Collect whole-body CT scan images of multiple subjects. Each CT scan image contains multiple two-dimensional cross-sectional images covering 42 organs in the human body. Divide the CT scan images into a training set and a test set.

[0007] Step 2: Preprocess the CT scan images in the training set and the test set to obtain multiple 3D sub-images of the same size;

[0008] Step 3: Use the training set in step 2 to perform self-pretraining on the deep learning model HPM with attention mechanism. The deep learning model HPM includes a block embedding layer, a position embedding layer, a student VIT network and a teacher VIT network.

[0009] Step 4: After the deep learning model HPM training is completed, the segmentation network model is constructed by connecting the encoder of the student VIT network to the UNETR decoder;

[0010] Step 5: Use the training set in step 2 to train the segmentation network model. After the training is completed, use the test set in step 2 to evaluate the trained segmentation network model.

[0011] Step 6: Input the new 3D sub-image into the segmentation network model, and finally output the 42-category segmentation result.

[0012] Among them, the preprocessing of the CT scan image described in step 2 includes downsampling, redirection, voxel scaling, background clipping, image clipping, random voxel transformation, and random affine transformation operations on each CT scan image in sequence to obtain multiple three-dimensional sub-images with a size of 96×96×96 voxels.

[0013] The student VIT network and the teacher VIT network described in step 3 have the same network structure, both consisting of an encoder, an image reconstruction decoder and a loss prediction decoder, where the parameters of the teacher network are determined by the student network with θ t ←mθ t +(1-m)θ s Update in the way of θ t and θ s Represent the teacher network parameters and the student network parameters respectively.

[0014] The self-pretraining method described in step 3 is:

[0015] Step 301: convert the 3D sub-images in the training set into a 3D block sequence, and then input the 3D block sequence into a block embedding layer, and output a block embedding of shape N*Q, where N represents the length of each block embedding, and Q represents the dimension of each block embedding;

[0016] Step 302: Input the block embedding into the position embedding layer, use the sine-cosine position embedding method to add position information to each block embedding, and obtain an output x with a shape of N*(Q+1);

[0017] Step 303: Input x into the encoder of the teacher VIT network In The output of the teacher VIT network is input to the loss prediction decoder In, after The output is the reconstruction loss L with a block-by-block prediction length of N. pred(t) ;

[0018] Step 304: Set the mask generation strategy from easy to difficult based on the reconstruction loss L pred(t) and the set training rounds, a binary mask M is obtained to indicate that the three-dimensional block sequence is masked during the training process;

[0019] Step 305: Input x and M into the encoder of the student VIT network Will The outputs are respectively input into the image reconstruction decoder of the student VIT network and loss prediction decoder In the reconstruction decoder and loss prediction decoder Output reconstruction loss L respectively rec(s) and the prediction loss L pred(s) ;

[0020] Step 306: Reconstruction loss L rec(s) and the prediction loss L pred(s) The two functions act alternately, gradually urging the student VIT network to learn the complex features in the 3D sub-image.

[0021] The method of converting the 3D sub-image in the training set into a 3D block sequence in step 301 is as follows: the 3D sub-image with a shape of H×W×D×C is reshaped into a series of 3D block sequences N×(P 3 C) = N × Q, Q = P 3 C; where C is the 3D sub-image input channel, (H; W; D) is the 3D sub-image resolution, (P; P; P) is the 3D block sequence resolution, and N = HWD / P 3 is the number of three-dimensional block sequences, that is, the length of the three-dimensional block sequence.

[0022] Among them, the encoder of the teacher VIT network in step 303 It consists of 12 transformer blocks, a loss prediction decoder It consists of 8 Transformer blocks; each Transformer block consists of alternating multi-head self-attention and MLP.

[0023] The reconstruction loss L in step 305 is rec(s) and the prediction loss L pred(s) They are:

[0024]

[0025] is the transformation function, generating the reconstruction target, represents a similarity measure;

[0026]

[0027] Representation loss prediction decoder The output of , σ(.) represents the sigmoid function, represents the relative relationship of reconstruction loss between data block i and data block j, where i,j=1,2...K is the patch index;

[0028]

[0029] M i =M j =0 means that both data block i and data block j are masked during the training process.

[0030] Among them, in the segmentation network model described in step 4, the student VIT network encoder is responsible for converting the input image into feature representation, and the UNETR decoder converts the features into segmentation results.

[0031] The training method of the segmentation network model in step 5 is as follows: the segmentation network model outputs a predicted segmentation result according to the input three-dimensional sub-image, and calculates the loss by comparing the predicted result with the actual segmentation label, including Dice loss and cross entropy loss bcelos;

[0032] Based on the calculated loss, the gradient of the loss function to the segmentation network model parameters is calculated through the back propagation algorithm; the optimization algorithm is used to adjust the model parameters according to the gradient so that the model can produce a smaller loss in the next round of training; the above steps are repeated until the specified round is reached.

[0033] The present invention discloses a medical image segmentation self-pretraining system based on a difficult sample mining mask autoencoder, comprising the following modules:

[0034] Image sample establishment module: Collect whole-body CT scan images of multiple subjects. Each CT scan image contains multiple two-dimensional cross-sectional images covering 42 organs in the human body. The CT scan images are divided into training sets and test sets.

[0035] Image sample preprocessing module: preprocess the CT scan images in the training set and the test set to obtain multiple 3D sub-images of the same size;

[0036] HPM training module: Use the training set in step 2 to perform self-pretraining on the deep learning model HPM with attention mechanism. The deep learning model HPM includes a block embedding layer, a position embedding layer, a student VIT network and a teacher VIT network;

[0037] Segmentation network building module: After the deep learning model HPM training is completed, the segmentation network model is built by connecting the encoder of the student VIT network to the UNETR decoder;

[0038] Segmentation network training module: Use the training set in step 2 to train the segmentation network model. After the training is completed, use the test set in step 2 to evaluate the trained segmentation network model.

[0039] Segmentation network prediction module: The new 3D sub-image is input into the segmentation network model, and the final output is 42 types of segmentation results.

[0040] Beneficial effects: The present invention has the following advantages: 1. In the present invention, the teacher VIT network helps train the student VIT network by mining difficult samples, allowing the student VIT network to focus on learning these difficult samples, effectively improving its feature extraction ability;

[0041] 2. The segmentation network model can combine the general feature extraction capability of the student VIT network encoder and the segmentation capability of the UNETR decoder to form an efficient end-to-end system; thereby significantly improving the segmentation performance of medical images. BRIEF DESCRIPTION OF THE DRAWINGS

[0042] Figure 1 Schematic diagram of the process of this method. DETAILED DESCRIPTION

[0043] The technical solution of the present invention is described in detail below in conjunction with the embodiments and drawings.

[0044] like Figure 1 As shown, the medical image segmentation method of the present invention comprises the following steps:

[0045] Step 1: Collect whole-body CT scan images of multiple subjects. Each CT scan image contains 493 two-dimensional cross-sectional images of 512×512 pixels, covering 42 organs in the human body. Each image has been professionally annotated by doctors, thus formulating the multi-organ segmentation problem as a 42-category segmentation task with a single-channel input; 46 CT scan images are used as training sets and 6 CT scan images are used as test sets.

[0046] Step 2: Preprocess each CT scan image sample in the training set and the test set. The size of the CT scan image is i*j*k, where i, j, and k represent the resolution of the image sample X in the horizontal axis, vertical axis, and vertical axis (or slice direction), respectively. The preprocessing process specifically includes:

[0047] Downsampling: The sampling intervals in the x, y, and z directions were 5.02, 5.02, and 5.01 mm, respectively, to reduce the resolution of the original CT scan image;

[0048] Perform the following operations based on the downsampled image:

[0049] Reorientation: transform the image from its original orientation to a unified coordinate system;

[0050] Voxel scaling: scale the voxel values ​​in the image and perform normalization;

[0051] Background clipping: remove the background part of the image that does not contain the area of ​​interest;

[0052] Image cropping: The image is randomly cropped into two pieces. The cropping split point is determined according to the ratio of negative samples to positive samples. The cropped image size is 96×96×96 voxels.

[0053] Random voxel transformation: Randomly transform the voxel values ​​of the image, such as changing the range of voxel values ​​or applying some noise;

[0054] Random affine transformation: apply random geometric transformations to images, such as rotation, scaling, and translation;

[0055] After the above operations, a training set and a test set consisting of multiple three-dimensional sub-images with a size of 96×96×96 voxels are obtained.

[0056] Step 3: Use the training set in step 2 to perform self-pretraining on the deep learning model HPM with attention mechanism.

[0057] Among them, HPM includes a block embedding layer, a position embedding layer, a student VIT network and a teacher VIT network. The two networks have the same network structure, which consists of an encoder, an image reconstruction decoder and a loss prediction decoder. The parameters of the teacher network are determined by the student network with θ t ←mθ t +(1-m)θ s Update in the way of θ t and θ s Represent the teacher network parameters and the student network parameters respectively.

[0058] In each training round, the 3D sub-image is first converted into a 3D block sequence. The conversion process is as follows: the 3D sub-image with the shape of H×W×D×C is reshaped into a series of 3D block sequences N×(P 3 C) = N × Q, Q = P 3 C; where C is the 3D sub-image input channel, (H; W; D) is the 3D sub-image resolution, (P; P; P) is the 3D block resolution, and N = HWD / P 3 is the number of three-dimensional block sequences, that is, the length of the three-dimensional block sequence.

[0059] The 3D block sequence is input into the block embedding layer, and the output is a block embedding of shape N*Q, where N represents the length of each block embedding and Q represents the dimension of each block embedding. In order to retain the position information of the block embedding in the 3D sub-image, the block embedding is input into the position embedding layer, and the sine-cosine position embedding method is used to add the position information to each block embedding, and the output x of shape N*(Q+1) is obtained. +1 means that a position information dimension is added to each block embedding.

[0060] Feed x into the encoder of the teacher VIT network middle( It consists of 12 transformer blocks, each of which consists of multi-head self-attention and MLP alternating. The output of the teacher VIT network is input to the loss prediction decoder middle( It consists of 8 Transformer blocks, each of which consists of multi-head self-attention and MLP alternating). The output is the reconstruction loss L with a block-by-block prediction length of N. pred(t) .

[0061] Set the mask generation strategy from easy to difficult, based on the reconstruction loss L pred(t) and training rounds, and obtain a binary mask M, which is used to indicate which 3D block sequences are masked during the training process. x⊙M represents the unmasked 3D block sequence, that is, the corresponding 3D block sequence is not masked.

[0062] The mask generation strategy from easy to difficult enables the student VIT network to learn mainly by reconstructing easy data blocks at the beginning of training. As the number of training rounds increases, the mask generation strategy will be gradually adjusted (such as gradually reducing the threshold of reconstruction loss), so that the student VIT network begins to face data blocks that are more difficult to reconstruct. In the later stages of training, the student VIT network will mainly face those data blocks that are most difficult to reconstruct. At this time, the student VIT network must fully explore and learn the complex features of the data. Through this strategy, the student VIT network can gradually adapt to and learn patterns of different difficulty levels in the data, thereby improving its generalization ability and ability to handle complex patterns. This approach helps the student VIT network understand the data more comprehensively and better generalize to unseen data in practical applications.

[0063] Input x and M into the encoder of the student VIT network Learn feature representations on the unmasked 3D patch sequence x⊙M. The output and mask blocks of The missing blocks in the output of are respectively input to the image reconstruction decoder of the student VIT network. and loss prediction decoder Image reconstruction decoder and loss prediction decoder By adding position embedding to all input tokens, reconstruction can be performed at the position of each mask block, and the reconstruction loss L is output respectively. rec(s) and the prediction loss L pred(s) , L (s) =L rec(s) +L pred(s) is the loss function of the student VIT network. Reconstruction loss L rec(s) and the prediction loss L pred(s) These two functions act alternately, gradually urging the student VIT network to learn the complex features in the data more carefully and better reconstruct those data blocks that were previously difficult to reconstruct, thereby gradually improving its reconstruction ability. In this way, HPM is able to gradually improve its understanding and processing capabilities of the data, thereby improving the accuracy of segmentation in medical image segmentation tasks.

[0064]

[0065] is the transformation function, generating the reconstruction target, Represents a similarity measure.

[0066]

[0067] Representation loss prediction decoder The output of , σ(.) represents the sigmoid function, represents the relative relationship between the reconstruction loss of data block i and data block j, where i,j=1,2...K is the patch index.

[0068]

[0069] M i =M j =0 means that both data block i and data block j are masked during the training process.

[0070] Step 4: After HPM training, the segmentation network model is constructed by connecting the encoder of the student VIT network to the UNETR decoder. The student VIT network encoder is responsible for converting the input medical image into feature representations, and the UNETR decoder converts these features into segmentation results. The segmentation network model can combine the general feature extraction capabilities of the student VIT network encoder and the segmentation capabilities of the UNETR decoder to form an efficient end-to-end system for processing 3D medical image segmentation tasks.

[0071] Step 5: Use the training set in step 2 to train the segmentation network model. After the training is completed, use the test set to evaluate the model. During the training process, the segmentation network model will adjust its parameters according to the loss function. Specifically, the segmentation network model outputs the predicted segmentation result based on the input three-dimensional sub-image, and calculates the loss by comparing the predicted result with the actual segmentation label, which usually includes Dice loss and cross entropy loss bcelos. According to the calculated loss, the gradient of the loss function to the segmentation network model parameters is calculated by the back propagation algorithm. The gradient represents the rate of change of the loss function relative to the model parameters, which is used to guide the adjustment of the parameters. Use an optimization algorithm (such as Adam, SGD, etc.) to adjust the model parameters according to the gradient, so that the model can produce a smaller loss in the next round of training. Repeat the above steps until the specified round is reached. In each round, the model will receive a batch of new data for forward feedback, loss calculation and parameter adjustment. The trained segmentation network model is used to predict new data.

[0072] After training, the test set is used to compare the predicted segmentation results of the trained segmentation network model with the actual segmentation labels to evaluate the generalization ability of the model.

[0073] Step 6: Input the new 3D sub-image into the segmentation network model, and finally output the 42-category segmentation result.

Claims

1. A medical image segmentation self-pretraining method based on difficult sample mining mask autoencoder, characterized in that: The following steps are involved: Step 1: Collect whole-body CT scan images of multiple subjects. Each CT scan image contains multiple two-dimensional cross-sectional images covering 42 organs in the human body. Divide the CT scan images into a training set and a test set. Step 2: Preprocess the CT scan images in the training set and the test set to obtain multiple 3D sub-images of the same size; Step 3: Use the training set in step 2 to perform self-pretraining on the deep learning model HPM with attention mechanism. The deep learning model HPM includes a block embedding layer, a position embedding layer, a student VIT network and a teacher VIT network. Step 4: After the deep learning model HPM training is completed, the segmentation network model is constructed by connecting the encoder of the student VIT network to the UNETR decoder; Step 5: Use the training set in step 2 to train the segmentation network model. After the training is completed, use the test set in step 2 to evaluate the trained segmentation network model. Step 6: Input the new 3D sub-image into the segmentation network model, and finally output the 42-category segmentation result.

2. The medical image segmentation self-pretraining method based on difficult sample mining mask autoencoder according to claim 1, characterized in that: The preprocessing of the CT scan image described in step 2 includes downsampling, redirection, voxel scaling, background clipping, image clipping, random voxel transformation, and random affine transformation operations on each CT scan image in sequence to obtain multiple three-dimensional sub-images with a size of 96×96×96 voxels.

3. The medical image segmentation self-pretraining method based on difficult sample mining mask autoencoder according to claim 1, characterized in that: The student VIT network and the teacher VIT network described in step 3 have the same network structure, both consisting of an encoder, an image reconstruction decoder, and a loss prediction decoder, where the parameters of the teacher network are determined by the student network with θ t ←mθ t +(1-m)θ s Update in the way of θ t and θ s Represent the teacher network parameters and the student network parameters respectively.

4. The medical image segmentation self-pretraining method based on difficult sample mining mask autoencoder according to claim 3 is characterized in that: The self-pretraining method described in step 3 is: Step 301: convert the 3D sub-images in the training set into a 3D block sequence, and then input the 3D block sequence into a block embedding layer, and output a block embedding of shape N*Q, where N represents the length of each block embedding, and Q represents the dimension of each block embedding; Step 302: Input the block embedding into the position embedding layer, use the sine-cosine position embedding method to add position information to each block embedding, and obtain an output x with a shape of N*(Q+1); Step 303: Input x into the encoder of the teacher VIT network In The output of the teacher VIT network is input to the loss prediction decoder In, after The output is the reconstruction loss L with a block-by-block prediction length of N. pred(t) ; Step 304: Set the mask generation strategy from easy to difficult based on the reconstruction loss L pred(t) and the set training rounds, a binary mask M is obtained to indicate that the three-dimensional block sequence is masked during the training process; Step 305: Input x and M into the encoder of the student VIT network Will The outputs are respectively input into the image reconstruction decoder of the student VIT network and loss prediction decoder In the reconstruction decoder and loss prediction decoder Output reconstruction loss L respectively rec(s) and the prediction loss L pred(s) ; Step 306: Reconstruction loss L rec(s) and the prediction loss L pred(s) The two functions act alternately, gradually urging the student VIT network to learn the complex features in the 3D sub-image.

5. The medical image segmentation self-pretraining method based on difficult sample mining mask autoencoder according to claim 4, characterized in that: The method of converting the 3D sub-image in the training set into a 3D block sequence in step 301 is as follows: the 3D sub-image with a shape of H×W×D×C is reshaped into a series of 3D block sequences N×(P 3 C) = N × Q, Q = P 3 C; Where C is the 3D sub-image input channel, (H; W; D) is the 3D sub-image resolution, (P; P; P) is the 3D block sequence resolution, and N = HWD / P 3 is the number of three-dimensional block sequences, that is, the length of the three-dimensional block sequence.

6. The medical image segmentation self-pretraining method based on difficult sample mining mask autoencoder according to claim 4, characterized in that: The encoder of the teacher VIT network in step 303 It consists of 12 transformer blocks, a loss prediction decoder It consists of 8 Transformer blocks; each Transformer block consists of alternating multi-head self-attention and MLP.

7. The medical image segmentation self-pretraining method based on difficult sample mining mask autoencoder according to claim 4, characterized in that: Step 305: The reconstruction loss L rec(s) and the prediction loss L pred(s) They are: is the transformation function, generating the reconstruction target, represents a similarity measure; Representation loss prediction decoder The output of , σ(.) represents the sigmoid function, represents the relative relationship of reconstruction loss between data block i and data block j, where i,j=1,2...K is the patch index; M i =M j =0 means that both data block i and data block j are shielded during the training process.

8. The medical image segmentation self-pretraining method based on difficult sample mining mask autoencoder according to claim 1, characterized in that: In the segmentation network model described in step 4, the student VIT network encoder is responsible for converting the input image into feature representation, and the UNETR decoder converts the features into segmentation results.

9. The medical image segmentation self-pretraining method based on difficult sample mining mask autoencoder according to claim 1, characterized in that: The training method of the segmentation network model in step 5 is as follows: the segmentation network model outputs a predicted segmentation result according to the input three-dimensional sub-image, and calculates the loss by comparing the predicted result with the actual segmentation label, including Dice loss and cross entropy loss bcelos; Based on the calculated loss, the gradient of the loss function to the segmentation network model parameters is calculated through the back propagation algorithm; the optimization algorithm is used to adjust the model parameters according to the gradient so that the model can produce a smaller loss in the next round of training; the above steps are repeated until the specified round is reached.

10. A medical image segmentation self-pretraining system based on difficult sample mining mask autoencoder, characterized in that: Includes the following modules: Image sample establishment module: Collect whole-body CT scan images of multiple subjects. Each CT scan image contains multiple two-dimensional cross-sectional images covering 42 organs in the human body. The CT scan images are divided into training sets and test sets. Image sample preprocessing module: preprocess the CT scan images in the training set and the test set to obtain multiple 3D sub-images of the same size; HPM training module: Use the training set in step 2 to perform self-pretraining on the deep learning model HPM with attention mechanism. The deep learning model HPM includes a block embedding layer, a position embedding layer, a student VIT network and a teacher VIT network; Segmentation network building module: After the deep learning model HPM training is completed, the segmentation network model is built by connecting the encoder of the student VIT network to the UNETR decoder; Segmentation network training module: Use the training set in step 2 to train the segmentation network model. After the training is completed, use the test set in step 2 to evaluate the trained segmentation network model. Segmentation network prediction module: The new 3D sub-image is input into the segmentation network model, and the final output is 42 types of segmentation results.

Images