Masked modeling-based method for constructing 3D medical image segmentation model and application thereof

Abstract

Disclosed are a masked modeling-based method for constructing a three-dimensional (3D) medical image segmentation model and an application thereof, belonging to the field of medical image segmentation. The method comprises: establishing a semi-supervised learning network comprising a student network and a teacher network, wherein the student network comprises a coding module for extracting a potential feature and a segmentation decoder for predicting a segmentation result, and the teacher network comprises a coding module and a segmentation decoder which have the same structure as the student network; and training the semi-supervised learning network, during the training, two random mask operations being performed on each image, then images being inputted into two networks, the weight of the student network being optimized and updated according to a training loss function, and then the weight being migrated to the teacher network, wherein the training loss function comprises a prototype representation loss for representing the difference between prototypes extracted and generated by the two networks, and the student network may further comprise a reconstruction decoder and an auxiliary segmentation decoder. The present method can improve the robustness, generalization performance and accuracy of a 3D medical image segmentation model when there are scarce samples.

Description

Method and application of 3D medical image segmentation model establishment based on mask modeling

Technical field

[0001] The present invention belongs to the field of medical image segmentation, and more specifically, relates to a method for establishing a 3D medical image segmentation model based on mask modeling and its application. [Background Technology]

[0002] Accurate and robust segmentation of organs or lesions from 3D medical images plays a vital role in clinical applications such as diagnosis and treatment. Compared with traditional segmentation algorithms, deep learning-based methods, such as convolutional neural networks, have been widely used in various vision tasks due to their excellent ability to learn high-level semantic image features. However, most 3D medical image segmentation methods rely on a large amount of labeled data for fully supervised training. Manual labeling is labor-intensive and costly, which has a negative impact on the performance of fully supervised segmentation methods. In contrast, semi-supervised learning (SSL) methods learn directly from labeled and unlabeled images, achieving better segmentation accuracy than using labeled data alone.

[0003] SSL typically uses labeled images to provide traditional supervised losses such as cross-entropy and Dice loss, while unlabeled images provide unsupervised losses such as consistency regularization and entropy minimization. In essence, the unsupervised losses all serve a common goal, which is to absorb additional knowledge from unlabeled images by training two learners and maximizing their decision consistency on the same task under certain perturbations (e.g., perturbations to the image, model, etc.). In practice, SSL methods typically involve a dual-model architecture, with the two models playing the roles of teacher and student respectively. The student model is trained normally, and its weights are merged into the teacher model through a synchronization mechanism (e.g., exponential moving average) at each update. The teacher model's predictions are relatively stable and reliable, and can therefore be used as pseudo-labels for the student model, that is, if the true labels are not available, the student's predictions can be kept consistent with the teacher.

[0004] While SSL methods have achieved great success, they face two major challenges for 3D medical image segmentation. First, even unlabeled 3D medical data is extremely scarce, and the segmentation objective adheres to clear anatomical priors. This results in a lack of diversity in tasks trained with consistency constraints, leading to a rapid agreement between the teacher and student, and reduced guidance gained from the consistency constraints. Second, the synchronization between the two models tightly couples their knowledge. This coupling inevitably becomes increasingly strict as training progresses, and the teacher network is unaware of errors in the knowledge learned by the student network. This exacerbates the aforementioned problems with consistency constraints, preventing the student model from providing constructive guidance. These two challenges result in 3D medical image segmentation models trained using SSL methods having poor robustness and generalization, and segmentation accuracy cannot be guaranteed.

[0005] [Summary of the invention]

[0006] In response to the defects of the existing technology and the need for improvement, the present invention provides a method for establishing a 3D medical image segmentation model based on mask modeling and its application. The purpose is to use random masks to allow the student network and the teacher network in the semi-supervised network to input different three-dimensional images with missing information, thereby improving the diversity of tasks, so that the two networks can robustly learn related and complementary features, so that the consistency constraints at the feature level can produce effective unsupervised guidance in the entire training process, and ultimately improve the robustness, generalization and accuracy of the three-dimensional medical image segmentation model.

[0007] To achieve the above objectives, according to one aspect of the present invention, a method for establishing a 3D medical image segmentation model based on mask modeling is provided, comprising:

[0008] (S1) Establish a semi-supervised learning network consisting of a student network and a teacher network;

[0009] The student network includes: a first encoding module and a decoding module; the first encoding module includes an encoder for extracting different scale features of the three-dimensional input image to obtain latent features; the decoding module includes a first segmentation decoder, and the first segmentation decoder is used to extract and upsample the latent features to obtain segmentation results;

[0010] The teacher network includes: a second encoding module having a structure consistent with the first encoding module, and a second segmentation decoder having a structure consistent with the first segmentation decoder;

[0011] (S2) A semi-supervised learning network is trained using a 3D medical image segmentation dataset that includes both labeled and unlabeled images. The training method is as follows: the weights of the teacher network are frozen, each image is subjected to two random masking operations, and then input into the student network and the teacher network respectively. The weights of the student network are optimized and updated according to a preset training loss function, and the updated weights are transferred to the teacher network. The training loss function includes the prototype representation loss L p1 , used to characterize the potential features V extracted from the student network and the teacher network s and V t The difference between the features of the corresponding regions of the segmented target;

[0012] (S3) extracting a first encoding module and a first decoder to connect and form a 3D medical image segmentation model.

[0013] Furthermore, for the student network or the teacher network, the feature p of the segmentation target corresponding area in its potential features is fg is calculated as follows:

[0014] Among them, V represents the potential feature, P represents the segmentation result; C represents the number of channels of the potential feature, V j represents the jth channel of V, P j represents the j-th channel of P; UP() represents the upsampling operation;

[0015] And, the prototype representation loss L p1 The expression is:

[0016] in, and Represent the potential features V s and V t The features of the corresponding area of ​​the segmented target, L mse represents the root mean square error, N and M represent the number of labeled and unlabeled images in the 3D medical image segmentation dataset, respectively.

[0017] Furthermore, the training loss function also includes: potential feature loss L fea ; Latent feature loss L fea The expression used to characterize the difference between the latent features extracted by the student network and the teacher network is:

[0018] Among them, L mse represents the root mean square error, N and M represent the number of labeled images and unlabeled images in the 3D medical image segmentation dataset, respectively; V i s and V it Represents the i-th image X in the 3D medical image segmentation dataset i After input, the latent features extracted by the student network and the teacher network.

[0019] Furthermore, in the student network, the decoding module further includes K auxiliary segmentation decoders; the auxiliary segmentation decoders are used to extract and upsample the latent features to obtain segmentation results; the upsampling methods of the K auxiliary segmentation decoders are different from each other and are all different from the first segmentation decoder;

[0020] In addition, the training loss function also includes: segmentation consistency loss L mc The segmentation consistency constraint is used to characterize the difference between the segmentation results of the K auxiliary segmentation decoders and the first segmentation decoder, and its expression is:

[0021] Where K is a positive integer, and Represents image X respectively i After input, the segmentation results predicted by the mth segmentation decoder and the nth segmentation decoder, express The result after sharpening processing; the segmentation decoder is an auxiliary segmentation decoder or a first segmentation decoder.

[0022] Furthermore, in the student network, the decoding module further includes: a reconstruction decoder; the reconstruction decoder is used to extract and upsample the latent features to restore the original image information to obtain a reconstructed image;

[0023] In addition, the training loss function also includes: reconstruction loss L sup1 ; The reconstruction loss is used to characterize the difference between the reconstructed image obtained by the student network and the original image, and its expression is:

[0024] in, Represents image X i The reconstructed image obtained by the student network after input, α represents the balance parameter.

[0025] Furthermore, in the first encoding module, the encoder further includes: F HybridFormer modules connected in sequence; in the second encoding module, the encoder further includes: F HybridFormer modules connected in sequence;

[0026] The HybridFormer module is used to calculate self-attention in pixel space and sample dimension;

[0027] Furthermore, the latent features extracted by the student network are the feature images extracted by the encoder in the first encoding module and processed by F HybridFormer modules, and the latent features extracted by the teacher network are the feature images extracted by the encoder in the second encoding module and processed by F HybridFormer modules;

[0028] Wherein, F is a positive integer.

[0029] Furthermore, the training loss function also includes: segmentation loss L sup2 , which is used to characterize the difference between the segmentation result predicted by the first segmentation decoder and the gold standard, and its expression is:

[0030] Where N represents the number of labeled images in the 3D medical image segmentation dataset, and Y i represents the i-th image X in the 3D medical image segmentation dataset i The gold standard of the corresponding segmented image, P i s Represents image X i After input, the segmentation result predicted by the first segmentation decoder; L seg Represents the sum of DICE loss and cross entropy loss.

[0031] Furthermore, the random mask operation includes:

[0032] The 3D medical image is divided into equal-sized and non-overlapping cubic blocks, a certain proportion of the cubic blocks are randomly selected, and the pixels in the corresponding areas are set to zero.

[0033] According to another aspect of the present invention, a 3D medical image segmentation method is provided, comprising:

[0034] The 3D medical image to be segmented is input into the 3D medical image segmentation model established by the 3D medical image segmentation model establishment method based on mask modeling provided by the present invention, and the segmentation result is obtained from the output of the 3D medical image segmentation model.

[0035] According to another aspect of the present invention, a computer-readable storage medium is provided, comprising a stored computer program. When the computer program is executed by a processor, the device containing the computer-readable storage medium is controlled to execute the method for establishing a 3D medical image segmentation model based on mask modeling provided by the present invention, and / or the 3D medical image segmentation method provided by the present invention.

[0036] In general, the above technical solutions conceived by the present invention can achieve the following beneficial effects:

[0037] (1) When training a 3D medical image segmentation model, the present invention performs a random masking operation on the input original image to generate two different 3D mask images, which are input into the student network and the teacher network respectively, so that the input images of the two models are incomplete, but they both contain complementary and global information. Due to the randomness of the masking strategy, the image segmentation tasks of the two networks also have huge diversity, so that relevant and complementary features can be learned robustly, so that the consistency constraint at the feature level produces effective unsupervised guidance in the entire training process, thereby achieving the effect of improving robustness, generalization and segmentation accuracy. At the same time, the training loss function includes the prototype representation loss L p1 , which can constrain the feature space and prototype representation to achieve better segmentation performance.

[0038] (2) In the preferred embodiment of the present invention, the features of the latent space (i.e., the latent features extracted by the encoding module) and the prediction map output by the segmentation decoder are used to extract the prototype representation of its foreground, which can effectively explore the connection between the feature map and the probability map and further improve the training effect of the model.

[0039] (3) In the preferred embodiment of the present invention, a latent feature loss is designed in the training loss function to characterize the difference between the latent features extracted by the student network and the teacher network. The introduction of this loss can enhance the similarity of the latent features extracted by the student network and the teacher network, thereby effectively realizing information complementarity when the inputs of the two networks use independent masks.

[0040] (4) In the preferred embodiment of the present invention, an auxiliary segmentation decoder is also introduced into the student network, and consistency constraints are set between the auxiliary segmentation decoders and the original first segmentation decoders, so that these segmentation decoders learn from each other. Since the upsampling methods used by the auxiliary segmentation decoders are different from each other and different from the original first segmentation decoder, the mutual learning between students is a decoupled student-to-student learning, which enables the first segmentation decoder to receive additional decoupled knowledge and obtain more constructive unsupervised guidance, providing the teacher model with error suspicion monitoring and correction capabilities, and can effectively avoid the impact of prediction uncertainty on network robustness, generalization and accuracy.

[0041] (5) In the preferred embodiment of the present invention, a reconstruction decoder for restoring the original image information is also introduced into the student network, so that the student network can complete the segmentation task and the reconstruction task at the same time, thereby realizing a joint learning task. At the same time, the reconstruction loss is correspondingly introduced into the training loss function. Since the two tasks in the joint learning task share the same encoder structure, the introduction of the reconstruction decoder and the reconstruction loss will enable the encoder structure to accurately extract potential features even when part of the image information is missing, thereby ensuring the prediction accuracy of the segmentation decoder.

[0042] (6) In the preferred embodiment of the present invention, one or more HybridFormer modules are introduced after the encoder in the encoding modules of the student network and the teacher network to calculate the self-attention in the pixel space and sample dimension, thereby further improving the feature representation capability in the latent space.

Brief Description of the Drawings

[0043] FIG1 is a schematic diagram of a framework of a semi-supervised learning model provided by an embodiment of the present invention;

[0044] FIG2 is a schematic diagram of a method for establishing a 3D medical image segmentation model based on mask modeling provided by an embodiment of the present invention;

[0045] Figure 3 is a schematic diagram comparing the segmentation results of the 3D medical image segmentation method provided by an embodiment of the present invention and the existing method on the same left atrial dataset; wherein, (a) is the segmentation gold standard, (b) is the segmentation result of the existing MT method, (c) is the segmentation result of the existing UA-MT method, (d) is the segmentation result of the existing SSASNet method, (e) is the segmentation result of the existing DTC method, (f) is the segmentation result of the existing URPC method, (g) is the segmentation result of the existing MCNet+ method, and (h) is the segmentation result of the 3D medical image segmentation method provided by an embodiment of the present invention. [Specific implementation method]

[0046] In order to make the objectives, technical solutions and advantages of the present invention more clearly understood, the present invention is further described in detail below with reference to the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are only intended to illustrate the present invention and are not intended to limit the present invention. In addition, the technical features involved in the various embodiments of the present invention described below may be combined with each other as long as they do not conflict with each other.

[0047] In the present invention, the terms "first", "second", etc. (if any) in the present invention and the drawings are used to distinguish similar objects and are not necessarily used to describe a specific order or sequence.

[0048] In order to solve the technical problems of poor robustness, generalization and low segmentation accuracy of existing SSL methods based on dual-model architecture when applied to three-dimensional medical image segmentation due to the scarcity of samples, the present invention provides a 3D medical image segmentation model establishment method based on mask modeling and its application. The overall idea is: by performing a random masking operation on the input original image, two different mask images are generated, which are input into the student network and the teacher network respectively, so that the input images of the two models are incomplete, but they jointly contain overall information. Due to the randomness of the masking strategy, the image segmentation tasks of the two networks also have huge diversity, so that relevant and complementary features can be robustly learned, so that the consistency constraints at the feature level produce effective unsupervised guidance in the entire training process, thereby achieving the effect of improving robustness, generalization and segmentation accuracy.

[0049] Based on the above masking strategy, in order to further improve the performance of 3D medical image segmentation, this paper also proposes two new learning methods, namely diverse joint task learning (DJL) and decoupled student learning (DIL), and instantiates them into an enhanced teacher-student architecture for robust 3D medical image segmentation. In diverse joint task learning, the student network completes the 3D image segmentation task while also completing the task of restoring the original image information. These two tasks will share the same encoder structure. Based on different mask images generated by the masking strategy, the teacher and student models learn to jointly segment the same target but restore different image contents. Due to the randomness of the masking strategy, this joint task also has great diversity. To facilitate the above-mentioned DJL, the student model in DIL is also attached with one or more auxiliary decoding branches. By using different upsampling designs, they can be equivalently regarded as other "students". By optimizing the pairwise consistency constraints of the output level, the auxiliary branches learn from the original branches, but their weights are independent of the teacher-student synchronization. The student thus receives additional decoupled knowledge and obtains more constructive unsupervised guidance, avoiding errors in the knowledge learned by the original branch that affect the teacher network's learning of correct knowledge, and providing the teacher model with error suspicion monitoring and correction capabilities.

[0050] The following are examples.

[0051] Example 1:

[0052] A method for establishing a 3D medical image segmentation model based on mask modeling, as shown in Figures 1 and 2, includes:

[0053] (S1) establishing a semi-supervised learning network including a student network and a teacher network. The structure of the semi-supervised learning network is shown in FIG1 ;

[0054] Referring to Figure 1, the student network includes: a first encoding module and a decoding module; the first encoding module includes an encoder for extracting different scale features of the 3D input image to obtain latent features; the decoding module includes a first segmentation decoder for extracting and upsampling the latent features to obtain segmentation results;

[0055] Optionally, in this embodiment, the encoder in the student network is composed of multiple convolutional layers and downsampling layers; the first segmentation decoder is correspondingly composed of multiple convolutional layers and upsampling layers, and the upsampling is specifically implemented by transposed convolution;

[0056] In this embodiment, the inputs to the student network and the teacher network are images that have undergone random masking operations, and some information may be missing from the images. To address this issue, in this embodiment, the first encoding module further includes, after the encoder, F sequentially connected HybridFormer modules, where F is a positive integer; the HybridFormer module is used to calculate self-attention in pixel space and sample dimensions; and the potential features extracted by the student network are the feature images extracted by the encoder in the first encoding module and processed by the F HybridFormer modules.

[0057] Optionally, in the HybridFormer module, the computation of pixel-space self-attention and sample-dimensional self-attention is performed by two separate parts, each of which includes a multi-layer convolutional layer and a downsampling layer. In practical applications, the number F of HybridFormer modules can be flexibly set according to actual needs.

[0058] The potential features output by the first encoding module are input into the first segmentation decoder. Through the cross connection method, the low-order output feature map and the feature map extracted by the corresponding large-scale convolution kernel of the previous level are spliced ​​along the channel direction. The large-scale features are further extracted using 3×3 convolution, and finally the segmentation prediction map is obtained.

[0059] Referring to Figure 1, the teacher network includes: a second encoding module and a second segmentation decoder; the structure of the second encoding module is consistent with the structure of the first encoding module in the student network, including an encoder and F HybridFormer modules after the encoder, which are used to extract the potential features of the input mask image; the structure of the second segmentation decoder is consistent with the structure of the first segmentation decoder in the student network, which is used to extract and upsample the potential features to obtain the segmentation result.

[0060] The introduction of the HybridFormer module in the student network and the teacher network effectively improves the feature representation ability in the latent space.

[0061] Referring to Figure 1, to achieve diverse joint task learning (DJL) and better learn the complementary information in mask image pairs, the decoding module in the student network includes two types of branches: a reconstruction decoder and a segmentation decoder. The reconstruction decoder is used to extract and upsample the latent features to restore the original image information and obtain a reconstructed image; the segmentation decoder predicts the segmentation result. At the same time, to reduce the uncertainty of the segmentation result prediction, in this embodiment, the segmentation decoders specifically include three, one of which is a first segmentation decoder and the other two are auxiliary segmentation decoders. These two auxiliary segmentation decoders are also used to extract and upsample the latent features to obtain segmentation results. However, the upsampling methods used by these two auxiliary segmentation decoders are different from each other and differ from the upsampling method of the first segmentation decoder, namely linear interpolation and neighboring interpolation. The auxiliary segmentation decoders cooperate with the first segmentation decoder to achieve decoupled student-to-student learning (DIL).

[0062] It should be noted that the number of auxiliary segmentation decoders and the upsampling method adopted here are merely exemplary descriptions and should not be understood as the only limitation to the present invention. In practical applications, they can be flexibly adjusted according to actual needs.

[0063] For the established semi-supervised learning model, this embodiment proposes a corresponding model training method, specifically:

[0064] (S2) training a semi-supervised learning network using a 3D medical image segmentation dataset including labeled and unlabeled images;

[0065] The training method is as follows: freeze the weights of the teacher network, perform two random mask operations on each image, and then input them into the student network and the teacher network respectively, optimize and update the weights of the student network according to the preset training loss function, and migrate the updated weights to the teacher network; optionally, in this embodiment, the student network is migrated to the teacher network, specifically through the exponential moving average (EMA).

[0066] In this embodiment, X i ∈□ H×W×D and Y i ∈{0,1} H×W×D To represent the input image and segmentation label, where H, W and D represent the height, width and depth of the 3D image. The labeled dataset and the unlabeled dataset are represented as Da L ={X i ,Y i |i=1,...N} and Da U ={X i |i=N+1,...N+M}, where N and M represent the number of labeled images and unlabeled images respectively.pred |X;θ) is the probability map of the input image, which is used to represent the probability of each pixel belonging to each type of segmentation target.

[0067] Optionally, in this embodiment, the specific manner of performing the mask operation on the three-dimensional image includes:

[0068] The 3D medical image is divided into cube blocks of equal size (for example, 4×4×4) and non-overlapping. A certain proportion of the cube blocks are randomly masked from a uniform distribution, and the pixels in the masked area are set to zero. Through the masking operation, the two mask images input to the student network and the teacher network are missing some information. Due to the randomness of the random masking operation, the mask areas of the two mask images are different, and the features in the latent space will be destroyed by the random mask, which is crucial for the medical image segmentation task. At the same time, the two mask images together contain complete image information.

[0069] Based on the semi-supervised learning network established in step (S1), the mask image input to the student network is processed by its encoding module to obtain the potential feature V s , and then the first segmentation decoder predicts the segmentation result P s_s_Tconv At the same time, the segmentation result P is obtained through two auxiliary segmentation decoders s_s_Tri , P s_s_Nea , after the reconstruction decoder restores the original image information, the reconstructed image is obtained The mask image input to the teacher network is processed by its encoding module to obtain the potential feature V t , and then the first segmentation decoder predicts the segmentation result P s_t .

[0070] In the existing dual-model architecture, pixel space information is extracted only by forcing the probability map of the student network to be aligned with the probability map of the teacher network, while ignoring the constraints of the latent feature space and the mining of information in unlabeled data. Although consistency learning plays a very important role in semi-supervised segmentation tasks, the mask operation adopted in this embodiment reduces redundant information and creates a feature representation task that is more challenging than noise or transformation. The same image with different random masks may lead to different predictions, especially when the target area to be segmented is masked. Directly aligning the prediction results of the teacher and student networks is too strict and may lead to prediction collapse. Inspired by prototype learning, this embodiment uses the features and prediction maps of the latent space to extract the prototype representation of its foreground to explore the connection between the feature map and the probability map. Accordingly, in this embodiment, the training loss function includes the prototype representation loss L p1 , used to characterize the potential features V extracted from the student network and the teacher network s and V tThe prototype, that is, the feature of the segmentation target corresponding area in the potential feature, is jointly determined by the potential feature extracted by the encoding module and the segmentation result predicted by the segmentation decoder. Therefore, in this embodiment, the design of the training loss function can constrain the feature space and prototype representation to train the established model to obtain better segmentation performance.

[0071] In order to extract multi-channel feature information, this embodiment averages the features along the channel dimension when calculating the prototype. Accordingly, for the student network or teacher network, the features of the region corresponding to the segmentation target in its latent features are calculated as follows:

[0072] Among them, V represents the potential feature, P represents the segmentation result; C represents the number of channels of the potential feature, V j represents the jth channel of V, P j represents the jth channel of P; UP() represents the upsampling operation; P fg ∈[0,1] H×W×D That is, it represents the foreground prototype. Based on the prototype generation method, the prototype representation loss L p1 The expression is:

[0073] in, and Represent the potential features V s and V t The features of the corresponding area of ​​the segmented target, L mse represents the root mean square error, N and M represent the number of labeled and unlabeled images in the 3D medical image segmentation dataset, respectively.

[0074] This embodiment further considers that the masks of the student network and the teacher network are independent, so the training loss function is further improved to enhance the potential features (V s ,V t ) to achieve information complementarity. Accordingly, in this embodiment, the training loss function also includes: potential feature loss L fea ; Latent feature loss L fea The expression used to characterize the difference between the latent features extracted by the student network and the teacher network is:

[0075] Among them, L mse represents the root mean square error, N and M represent the number of labeled images and unlabeled images in the 3D medical image segmentation dataset, respectively; V i s and V i tRepresents the i-th image X in the 3D medical image segmentation dataset i After input, the latent features extracted by the student network and the teacher network.

[0076] In order to effectively avoid the impact of prediction uncertainty on network robustness, generalization and accuracy, this embodiment introduces an auxiliary segmentation decoder, which makes the auxiliary segmentation decoder and the original first segmentation decoder pass consistency constraints, and uses a sharpening function to reduce the impact of pixels that are easily misclassified. For any segmentation decoder, P i s represents the i-th image X i The predicted segmentation result after input is expressed as follows:

[0077] Where T represents the hyperparameter used to control the degree of sharpening; P i s_sharp That is, the result after the sharpening function is processed. It will be used as a pseudo label in the consistency constraint between the segmentation decoders, so that the segmentation result predicted by the segmentation decoder can be close to the result after the sharpening process;

[0078] In this embodiment, the consistency constraint between the two segmentation decoders is achieved by the segmentation consistency loss L mc To characterize, accordingly, the training loss function also includes: segmentation consistency loss L mc ; The segmentation consistency constraint is used to characterize the difference between the segmentation results of the two auxiliary segmentation decoders and the first segmentation decoder, and its expression is:

[0079] in, and Represents image X respectively i After input, the segmentation results predicted by the mth segmentation decoder and the nth segmentation decoder, express The result after sharpening processing; the segmentation decoder is an auxiliary segmentation decoder or a first segmentation decoder; this embodiment can constrain the segmentation results of each segmentation decoder to remain consistent by introducing a segmentation consistency constraint in the training loss function.

[0080] Since this embodiment implements joint task learning, the joint task includes not only the segmentation task but also the reconstruction task, and the two tasks share the same encoding structure; when the mask image information is accurate, the reconstruction task can reconstruct the original, undamaged voxel information, so that the encoder can extract features more effectively. In order to ensure the accuracy of the reconstruction task, in this embodiment, the training loss function also includes: reconstruction loss Lsup1 ; The reconstruction loss is used to characterize the difference between the reconstructed image obtained by the student network and the original image, and its expression is:

[0081] in, Represents image X i The reconstructed image obtained by the student network after input, α represents the balance parameter.

[0082] At the same time, in order to ensure the accuracy of the segmentation task, in this embodiment, the training loss function also includes: segmentation loss L sup2 , which is used to characterize the difference between the segmentation result predicted by the first segmentation decoder and the gold standard, and its expression is:

[0083] Where N represents the number of labeled images in the 3D medical image segmentation dataset, and Y i represents the i-th image X in the 3D medical image segmentation dataset i The gold standard of the corresponding segmented image, P i s Represents image X i After input, the segmentation result predicted by the first segmentation decoder; L seg Represents the sum of DICE loss and cross entropy loss.

[0084] Based on the above analysis, in this embodiment, the overall training loss function can be expressed as follows: L total =L p1 +L fea +L mc +L sup1 +L sup2 .

[0085] On the basis of random masking, based on the above semi-supervised learning network structure and the corresponding training loss function, this embodiment can obtain a student network with excellent segmentation performance.

[0086] Based on the above steps (S1) and (S2), this embodiment further includes: (S3) extracting a first encoding module and a first decoder to connect and form a 3D medical image segmentation model.

[0087] In general, this embodiment ensures task diversity through a random masking strategy. On this basis, through diverse joint task learning (DJL) and decoupled student-to-student learning (DIL), the student network and teacher network in the semi-supervised learning network can robustly learn relevant and complementary features, so that the feature-level consistency constraints produce effective unsupervised guidance throughout the training process, and the student network receives decoupled knowledge to obtain more constructive unsupervised guidance, providing the teacher network with error suspicion monitoring and correction capabilities. Ultimately, the 3D medical image segmentation model established in this embodiment has stronger robustness and generalization, as well as higher segmentation accuracy.

[0088] Example 2:

[0089] A 3D medical image segmentation method, comprising:

[0090] The 3D medical image to be segmented is input into the 3D medical image segmentation model established by the 3D medical image segmentation model establishment method based on mask modeling provided in the above embodiment 1, and the segmentation result is obtained from the output of the 3D medical image segmentation model.

[0091] Since the 3D medical image segmentation model established in this embodiment 1 has stronger robustness and generalization, as well as higher segmentation accuracy, based on this 3D medical image segmentation model, this embodiment can obtain high-precision segmentation results in various three-dimensional medical image segmentation scenarios.

[0092] Example 3:

[0093] A computer-readable storage medium includes a stored computer program. When the computer program is executed by a processor, the device where the computer-readable storage medium is located is controlled to execute the method for establishing a 3D medical image segmentation model based on mask modeling provided in the above-mentioned embodiment 1, and / or the 3D medical image segmentation method provided in the above-mentioned embodiment 2.

[0094] The following uses the segmentation results of the 3D GE-MRI dataset from the Left Atrial Segmentation Challenge to further verify the beneficial effects of the method provided by the present invention. The dataset was acquired using a clinical whole-body MRI scanner with a data resolution of 0.625×0.625×0.625mm. 3 , of which 154 scans have labels given by experts. 123 scans were randomly selected for training and 31 data samples were used for testing. Before network training, all scans were preprocessed, pixel intensities were normalized to unit variance and zero mean, and samples were randomly cropped to 112×112×80mm. 3. On this dataset, the image segmentation visualization results of the 3D medical image segmentation method provided by the present invention and other existing segmentation methods on the left atrium dataset are shown in Figure 3; in Figure 3, (a) is the segmentation gold standard, (b) to (g) are the segmentation results of the existing MT, UA-MT, SSASNet, DTC, URPC, and MCNet+ methods respectively, and (h) is the segmentation result of the 3D medical image segmentation method provided by an embodiment of the present invention. The quantitative segmentation results at 10% and 20% label amounts are shown in Table 1. The evaluation indicators include segmentation accuracy (Dice, Jaccard), 95% Hausdorff distance (95HD), and average surface distance (ASD). Among them, the larger the Dice and Jaccard, the better the network performance, and the smaller the 95HD and ASD, the better the network performance. In Table 1, "Ours" represents the segmentation method provided by the present invention.

[0095] As shown in Table 1, the segmentation results of the supervised training V-Net at 10% and 20% labels are used as the baseline. It can be observed from the table that, compared with the baseline, all semi-supervised methods can provide more effective guidance on the segmentation results at 10% label amount, which reveals a large amount of diverse and useful information for model training contained in the unlabeled data. In particular, as shown in Figure 3, the supervised V-Net with 100% labeled data outperforms all semi-supervised methods. However, there is only a slight difference between the method of the present invention and the supervised V-Net with 100% labeled data. In addition, all the quantitative results also comprehensively reflect that the segmentation method provided by the present invention has more substantial improvements than other semi-supervised methods.

[0096] Specifically, mainstream semi-supervised methods achieved significant improvements over their supervised V-Net counterparts on both 20% and 10% labeled data. In particular, our method outperformed the compared methods on all four quantitative metrics. Among the compared methods, MC-Net+ slightly lagged behind our method. Segmentation results for all semi-supervised methods improved with increasing amounts of labeled data. Even with extremely limited labeled data, our method still delivered significantly improved segmentation metrics compared to other state-of-the-art semi-supervised algorithms. To visually demonstrate the superiority of our method, Figure 3 shows 2D and 3D visualizations of the segmentation results generated by all compared methods using 20% ​​labeled data. Clearly, our method preserves more detail and sharpens the most isolated regions. While our method does not achieve optimal results in terms of parameter count and computational complexity, the differences are not significant compared to the other compared algorithms, and therefore, given the powerful computing power of modern computers, it does not impose a significant computational burden.

[0097] Table 1 Quantitative comparison of various evaluation methods in the left atrium segmentation task

[0098] It will be easily understood by those skilled in the art that 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, and improvements made within the spirit and principles of the present invention should be included in the scope of protection of the present invention.

Claims

1. A method for establishing a 3D medical image segmentation model based on mask modeling, characterized in that: include: (S1) Establish a semi-supervised learning network including a student network and a teacher network; The student network comprises: a first encoding module and a decoding module; the first encoding module comprises an encoder, which is used to extract different scale features of the three-dimensional input image to obtain latent features; the decoding module comprises a first segmentation decoder, which is used to extract features and upsample the latent features to obtain segmentation results; The teacher network includes: a second encoding module having a structure consistent with the first encoding module, and a second segmentation decoder having a structure consistent with the first segmentation decoder; (S2) The semi-supervised learning network is trained using a 3D medical image segmentation dataset including annotated images and unannotated images, wherein the training method is as follows: the weight of the teacher network is frozen, each image is subjected to two random mask operations, and then input into the student network and the teacher network respectively, the weight of the student network is optimized and updated according to a preset training loss function, and the updated weight is transferred to the teacher network; the training loss function includes a prototype representation loss L p1 , used to characterize the potential features V extracted by the student network and the teacher network s and V t The difference between the features of the corresponding regions of the segmented target; (S3) Extract the first encoding module and the first decoder to form a 3D medical image segmentation model.

2. The method for establishing a 3D medical image segmentation model based on mask modeling according to claim 1, characterized in that: For the student network or the teacher network, the feature p of the segmentation target corresponding area in its potential features fg The calculation method is as follows: Among them, V represents the potential feature, P represents the segmentation result; C represents the number of channels of the potential feature, V j represents the jth channel of V, P j represents the j-th channel of P; UP() represents the upsampling operation; And, the prototype characterization loss L p1 The expression is: in, and Represent the potential features V s and V t The features of the corresponding area of ​​the segmented target, L mse represents the root mean square error, N and M represent the number of labeled images and unlabeled images in the 3D medical image segmentation dataset, respectively.

3. The method for establishing a 3D medical image segmentation model based on mask modeling according to claim 1, characterized in that: The training loss function also includes: potential feature loss L fea ; The potential feature loss L fea The expression used to characterize the difference between the potential features extracted by the student network and the teacher network is: Among them, L mse represents the root mean square error, N and M represent the number of annotated images and unannotated images in the 3D medical image segmentation dataset, respectively; V i s and V i t Respectively represent the i-th image X in the 3D medical image segmentation dataset i After input, the student network and the teacher network extract latent features.

4. The method for establishing a 3D medical image segmentation model based on mask modeling according to claim 3, characterized in that: In the student network, the decoding module further includes K auxiliary segmentation decoders; the auxiliary segmentation decoders are used to extract and upsample the potential features to obtain segmentation results; the upsampling methods of the K auxiliary segmentation decoders are different from each other and are all different from the first segmentation decoder; In addition, the training loss function also includes: segmentation consistency loss L mc The segmentation consistency constraint is used to characterize the difference between the segmentation results of the K auxiliary segmentation decoders and the first segmentation decoder, and its expression is: Where K is a positive integer, and Represents image X i After input, the segmentation results predicted by the mth segmentation decoder and the nth segmentation decoder are, express The result after sharpening processing; the segmentation decoder is an auxiliary segmentation decoder or the first segmentation decoder.

5. The method for establishing a 3D medical image segmentation model based on mask modeling according to claim 4, characterized in that: In the student network, the decoding module further includes: a reconstruction decoder; the reconstruction decoder is used to extract and upsample the latent features to restore the original image information and obtain a reconstructed image; In addition, the training loss function also includes: reconstruction loss L sup1 ; The reconstruction loss is used to characterize the difference between the reconstructed image obtained by the student network and the original image, and its expression is: in, Represents image X i The reconstructed image is obtained by reconstructing the student network after input, and α represents the balance parameter.

6. The method for establishing a 3D medical image segmentation model based on mask modeling according to claim 5, characterized in that: In the first encoding module, the encoder further includes: F HybridFormer modules connected in sequence; in the second encoding module, the encoder further includes: F HybridFormer modules connected in sequence; The HybridFormer module is used to calculate the self-attention in pixel space and sample dimension; Furthermore, the potential features extracted by the student network are the feature images extracted by the encoder in the first encoding module and processed by F HybridFormer modules, and the potential features extracted by the teacher network are the feature images extracted by the encoder in the first encoding module and processed by F HybridFormer modules. The potential features extracted by the network are the feature images extracted by the encoder in the second encoding module and processed by F HybridFormer modules; Wherein, F is a positive integer.

7. The method for establishing a 3D medical image segmentation model based on mask modeling according to any one of claims 1 to 6, characterized in that: The training loss function also includes: segmentation loss L sup2 , which is used to characterize the difference between the segmentation result predicted by the first segmentation decoder and the gold standard, and its expression is: Where N represents the number of annotated images in the 3D medical image segmentation dataset, and Y i represents the i-th image X in the 3D medical image segmentation dataset i The gold standard of the corresponding segmented image, P i s Represents image X i After input, the first segmentation decoder predicts the segmentation result; L seg Represents the sum of DICE loss and cross entropy loss.

8. The method for establishing a 3D medical image segmentation model based on mask modeling according to any one of claims 1 to 6, characterized in that: The random mask operation includes: The 3D medical image is divided into equal-sized and non-overlapping cubic blocks, a certain proportion of the cubic blocks are randomly selected, and the pixels in the corresponding area are set to zero.

9. A 3D medical image segmentation method, characterized in that: include: The 3D medical image to be segmented is input into the 3D medical image segmentation model established by the 3D medical image segmentation model establishment method based on mask modeling according to any one of claims 1 to 8, and the segmentation result is obtained from the output of the 3D medical image segmentation model.

10. A computer-readable storage medium, characterized in that: It includes a stored computer program, which, when executed by a processor, controls the device where the computer-readable storage medium is located to execute the method for establishing a 3D medical image segmentation model based on mask modeling as described in any one of claims 1 to 8, and / or the 3D medical image segmentation method as described in claim 9.

Images