Task-aware pre-training method for 3D medical image reconstruction
By combining high-semantic information masks with dynamic learning strategies, the problems of redundant background and invalid information in 3D medical image reconstruction are solved, achieving more efficient training and higher-precision segmentation effects.
Patent Information
- Application Number
- CN202411619355.9
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Priority Date
- 2024-10-11
- Filing Date
- 2024-11-13
- Publication Date
- 2025-09-30
- Estimated Expiration
- 2044-11-13
AI Technical Summary
Existing three-dimensional medical image reconstruction methods contain redundant background and invalid information, which increases the model training cost and reduces the segmentation accuracy. In addition, random and fixed mask strategies affect the model training effect and efficiency.
A task-aware 3D medical image reconstruction pre-training method was adopted. Through the high semantic information masking strategy, selective masking and dynamic learning strategy, the mask reconstruction main task and segmentation auxiliary task were constructed to jointly optimize the training network.
It improves training efficiency and segmentation accuracy, reduces task difficulty and cost, enhances the model's generalization and feature extraction capabilities, and ensures efficient execution of downstream tasks.
Smart Images

Figure CN119600034B_ABST
Abstract
Description
Technical Field
[0001] The present invention belongs to the technical field of medical image processing, and in particular relates to a three-dimensional medical image reconstruction pre-training method based on task perception. Background Art
[0002] With advances in modern medical technology, 3D medical images have become an essential tool for diagnosing and treating a variety of diseases. They provide detailed anatomical information, assisting physicians in disease assessment, surgical planning, and efficacy monitoring. However, manual annotation of medical image data requires extensive expertise and is time-consuming. Segmentation accuracy is also susceptible to subjective factors. Consequently, computer-assisted automatic segmentation techniques have attracted considerable attention.
[0003] In order to improve the segmentation accuracy of three-dimensional medical images, pre-training methods are widely used. Among them, pre-training methods based on mask reconstruction self-supervised learning have gradually received attention in research. For example, Patent Document 1 (Application No. 202211383092.7) discloses a lung medical image segmentation method based on reconstruction pre-training. This method learns feature representation through unsupervised puzzle reconstruction pre-training, combined with the idea of transfer learning, to improve the robustness and generalization of the model and reduce the model training time. Patent Document 2 (Application No. 202310959502.6) discloses a training method for a medical image processing model and a medical image processing method, which are mainly based on random masking at a given mask rate, and based on this, construct a pre-training agent task for deblurring mask image modeling, so that the latent representation learned in the pre-training stage includes additional attention to the details of the ultrasound image.
[0004] However, during implementation, the method in Patent Document 1 requires that each image be gridded and the order be disrupted, which results in a large amount of redundant background in the image. This causes the model to input too much invalid information, increasing the training cost and reconstruction difficulty, affecting the reconstruction effect, and thus affecting the segmentation accuracy. Furthermore, in subsequent segmentation models, it is necessary to continue adding reconstruction branches, increasing the complexity of downstream tasks. The random masking method used in Patent Document 2 retains a large amount of invalid information, and the mask probability is a fixed value. When the mask probability is too large, it may cause the model to converge to a suboptimal solution. When the mask probability is too small, it increases the learning time or leads to insufficient learning. It also increases the time cost of the mask probability parameter ablation experiment, thereby affecting the segmentation accuracy.
[0005] In summary, the existing methods have a lot of redundant background and invalid information in the image reconstruction process. In addition, the masking strategies of random masking and fixed masking probability will increase the model training cost, affect the reconstruction difficulty and reconstruction effect, and thus affect the segmentation accuracy; and the subsequent segmentation model needs to continue to add reconstruction branches, which increases the complexity of downstream tasks. Summary of the Invention
[0006] To address at least one of the problems in the prior art, the present invention provides a task-aware pre-training method for 3D medical image reconstruction. The technical problem to be solved by the present invention is achieved through the following technical solutions:
[0007] In a first aspect, the present invention proposes a task-aware 3D medical image reconstruction pre-training method, comprising:
[0008] Preprocessing the original three-dimensional medical image to obtain a preprocessed image;
[0009] Construct a mask reconstruction main task and establish a main task network within the mask reconstruction main task; the main task network includes a mask module, a reconstruction main task encoder, and a reconstruction main task decoder connected in series. Input the preprocessed image into the main task network, use the mask module to selectively mask the preprocessed image based on a high semantic information masking strategy to obtain a masked image, reconstruct the masked image using the reconstruction main task encoder and reconstruction main task decoder, and calculate the loss function of the main task network based on the reconstruction result.
[0010] Construct a segmentation auxiliary task and establish an auxiliary task network within the segmentation auxiliary task to perceive the downstream segmentation task; the auxiliary task network includes a segmentation task encoder and a segmentation task decoder; input the original 3D medical image into the auxiliary task network, use the segmentation task encoder and segmentation task decoder to segment the original 3D medical image, and calculate the loss function of the auxiliary task network based on the segmentation results;
[0011] The mask reconstruction main task and the segmentation auxiliary task are jointly learned to simultaneously optimize the training of the main task network and the auxiliary task network, thereby completing the pre-training of 3D medical image reconstruction.
[0012] In a second aspect, the present invention provides a three-dimensional medical image segmentation method, comprising:
[0013] Build a segmentation network model;
[0014] Input the 3D medical image to be segmented into the trained segmentation network model to obtain the 3D medical image segmentation result; wherein the segmentation network model is trained using the following method:
[0015] Randomly initializing the segmentation head of the segmentation network model and loading the weights of the encoder part of the pre-trained network into the remaining layers of the segmentation network model to initialize the parameters of the segmentation network model; the weights of the encoder part of the pre-trained network are obtained using the task-aware 3D medical image reconstruction pre-training method proposed in the first aspect of the present invention;
[0016] The parameters of the initialized segmentation network model are further fine-tuned through optimization training to obtain a trained segmentation network model.
[0017] Beneficial effects of the present invention:
[0018] 1. The present invention proposes a task-aware 3D medical image reconstruction pre-training method. First, the original 3D medical image is pre-processed. Then, a mask reconstruction main task and a segmentation auxiliary task are constructed. A main task network is established in the mask reconstruction main task. The main task network selectively masks the pre-processed image based on a high semantic information masking strategy and realizes image reconstruction. An auxiliary task network is established in the segmentation auxiliary task. The auxiliary task network can perceive the downstream segmentation task, thereby realizing image segmentation. Finally, the joint learning of the reconstruction main task and the segmentation auxiliary task can optimize the two goals at the same time and complete the 3D medical image reconstruction pre-training. On the one hand, this method adopts a high semantic information masking strategy during the pre-training process, so that the mask image pays more attention to the foreground features rich in semantic information, eliminates redundant background, enriches the reconstruction information, reduces the difficulty of the mask reconstruction task, and improves the training efficiency and effect. On the other hand, this method introduces an auxiliary task branch, enhances the task perception ability of the pre-training network, clarifies the goal of network pre-training, and can more efficiently extract 3D medical image features, thereby improving the segmentation accuracy without increasing the complexity of the downstream task network and the training cost.
[0019] 2. During pre-training, this invention employs a dynamically changing learning strategy to control the difficulty of the mask reconstruction task. Compared to a fixed mask probability, this prevents the network from converging to suboptimal solutions early in training, enabling progressive learning from easier to more difficult solutions. Furthermore, this dynamic learning strategy significantly reduces the time cost of parameter tuning and avoids the variability of fixed mask probabilities across different datasets, thereby improving model generalization.
[0020] The present invention will be further described in detail below with reference to the accompanying drawings and embodiments. BRIEF DESCRIPTION OF THE DRAWINGS
[0021] Figure 1 1 is a flow chart of a task-aware 3D medical image reconstruction pre-training method provided by an embodiment of the present invention;
[0022] Figure 2 This is a flowchart of constructing the mask reconstruction main task and the segmentation auxiliary task provided by an embodiment of the present invention. DETAILED DESCRIPTION
[0023] The following will clearly and completely describe the technical solutions in the embodiments of the present invention in conjunction with the accompanying drawings. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without making creative efforts are within the scope of protection of the present invention.
[0024] The first aspect of the present invention proposes a task-aware 3D medical image reconstruction pre-training method. Figure 1 , Figure 1 1 is a flow chart of a task-aware 3D medical image reconstruction pre-training method provided by an embodiment of the present invention. The method mainly includes the following steps:
[0025] Step 1: Preprocess the original three-dimensional medical image to obtain a preprocessed image.
[0026] It is understandable that in the process of implementing the present invention, it is first necessary to construct a data set by collecting a large number of original three-dimensional medical images. Due to the sensitivity, sparsity and difficulty of labeling of medical data, it is difficult to collect a large amount of image data of uniform size. Therefore, it is necessary to preprocess the collected original three-dimensional medical images to solve the above problems. Furthermore, considering that the effective information of medical image data is mainly concentrated in the foreground area, the present invention can randomly crop the central area during preprocessing to ensure that the diversity of the data is increased on the basis of highlighting the target area, thereby improving the generalization ability of the model. Then step 1 can specifically include the following two sub-steps:
[0027] 11) Perform shearing and random flipping operations on the original three-dimensional medical image to obtain a flipped image.
[0028] Specifically, by clipping the original 3D medical image, its values can be limited to a specific range, and these values are rescaled to the range of [0, 1]. Then, these images are randomly flipped to obtain flipped images.
[0029] 12) For the flipped image, a central area is set based on the center of the image, and a point in the central area is randomly selected as the center point for cropping to obtain the preprocessed image.
[0030] Optionally, in this embodiment, the image resolution is set to 96×96×96, that is, the values of the image width W, height H, and depth D are all 96. The above preprocessing operation is performed on each original 3D medical image to obtain the input image of the pre-training stage, wherein a single input image can be represented as x∈R( W×H×D ).
[0031] Step 2: Construct the mask reconstruction main task and establish the main task network in the mask reconstruction main task; the main task network includes a mask module, a reconstruction main task encoder, and a reconstruction main task decoder connected in series; the preprocessed image is input into the main task network, and the mask module is used to selectively mask the preprocessed image based on the high semantic information masking strategy to obtain a masked image. The masked image is reconstructed using the reconstruction main task encoder and the reconstruction main task decoder, and the loss function of the main task network is calculated based on the reconstruction result.
[0032] See Figure 2 , Figure 2 This is a flowchart of constructing the mask reconstruction main task and the segmentation auxiliary task provided by the embodiment of the present invention. Figure 2 As shown, for the mask reconstruction main task, this embodiment designs a mask module and constructs a reconstruction main task encoder and a reconstruction main task decoder. The mask module, the reconstruction main task encoder and the reconstruction main task decoder together form a main task network, which can also be called a reconstruction network.
[0033] Optionally, as an implementation method, in the mask module, this embodiment designs a high-semantic information masking strategy, which mainly analyzes the semantic richness of the data, masks areas with fewer features, and retains areas with rich features for learning.
[0034] Specifically, the mask module is used to perform selective masking on the preprocessed image based on the high semantic information masking strategy, which can be achieved by following the following sub-steps.
[0035] a) Divide the preprocessed image into several non-overlapping image blocks, and divide each image block into several two-dimensional slices.
[0036] For the input image x∈R (W×H×D) , first divide it into N = (W × H × D) / p 3 non-overlapping image patches x p ∈R N×p3 , p is the image block x p Then, for each image block x p Slicing is performed, and after slicing, each image block contains z two-dimensional slices I.
[0037] b) For each image block, the number of key points and the key point contrast response value of each two-dimensional slice are calculated using a scale-invariant feature transformation algorithm to obtain the number of key points and the key point contrast response value of the image block.
[0038] For an image block x p, the Scale-Invariant Feature Transform (SIFT) algorithm is used to calculate the number of key points in each corresponding two-dimensional slice and the contrast response values of these key points.
[0039] It is understood that the key points referred to in this embodiment are points in the image that are significant, invariant, and recognizable, that is, local extreme points in the generated differential Gaussian image, and the contrast response value is the contrast of the local extreme points. i _keypoints represents the number of key points of the i-th two-dimensional slice, with I i _response represents the contrast response value of the key point of the i-th two-dimensional slice. Then the image block x p The number of key points x p _keypoints is the sum of the number of key points of all corresponding two-dimensional slices; correspondingly, the image block x p The contrast response value x of the key point p _response is also the sum of the contrast response values of all its two-dimensional slice key points, and the formula is as follows:
[0040]
[0041] c) Calculate the SIFT score of each image patch based on the number of key points and contrast response value of each image patch.
[0042] Optionally, as an implementation manner, this embodiment calculates the SIFT score of each image block using the following formula.
[0043]
[0044] Where x p _score represents the image block x p SIFT score of x p _keypoints represents the image block x p The number of key points, W represents the weight of the contrast response value, x p _response represents the image block x p The contrast response value, x p _response_batch represents the maximum value of the contrast response value in all image blocks corresponding to the input image x.
[0045] d) Based on the mask probability of the current training round, image blocks with low SIFT scores are selected for masking, while image blocks with high SIFT scores are retained to obtain a mask image.
[0046] Specifically, this embodiment uses the SIFT score of each image block as an indicator of the feature richness of each image block. Obviously, for image blocks with a small number of key points and low contrast response values, their SIFT scores are low, indicating that the features of the image block are not rich enough. Therefore, these image blocks are masked according to the masking probability of the current training round. Conversely, for image blocks with a large number of key points and high contrast response values, their SIFT scores are high, indicating that the features of the image block are richer. These image blocks are retained and entered into the encoder for feature learning.
[0047] The high-semantic information masking strategy designed in this paper prioritizes semantically rich foreground information, focusing on key anatomical areas. It utilizes a scale-invariant feature transformation algorithm to measure feature richness, calculates the number of keypoints in each image patch and their contrast response, and selects feature-rich regions for learning. Compared to jigsaw puzzle reconstruction, this method eliminates redundant background; and compared to random masking, it uniformly samples all regions in the medical image. This improves the accuracy of downstream segmentation tasks while also enhancing pre-training efficiency.
[0048] Furthermore, this embodiment abandons the existing fixed mask probability for determining the mask probability and designs a dynamic learning strategy to adjust the mask probability. Since smaller mask probabilities represent more preserved image features during the masking process, the dynamic learning strategy designed in this embodiment increases the mask probability inversely with the number of training epochs. This allows for a higher mask probability in the early stages of pre-training, which means less information is fed into the network for learning. Subsequently, this information is gradually enriched, implementing a learning strategy that progresses from easy to difficult, thereby controlling the difficulty of the mask reconstruction task.
[0049] Optionally, as an implementation method, the mask probability can be calculated using the following formula:
[0050]
[0051] Where r n represents the mask probability during the nth round of training, n represents the current training round, r init Represents the preset initial mask probability, r final Represents the preset termination mask probability, and num_epoch represents the total number of training rounds.
[0052] During pre-training, this method uses a dynamically changing learning strategy to control the difficulty of the mask reconstruction task. Compared to a fixed mask probability, this method prevents the network from converging to a suboptimal solution early in training, enabling progressive learning from easier to more difficult. Furthermore, this dynamic learning strategy significantly reduces the time cost of parameter tuning and avoids the variability of fixed mask probabilities across different datasets, thereby increasing the model's generalization capabilities.
[0053] Furthermore, for the reconstructed main task encoder in the main task network, this embodiment is mainly designed according to the downstream segmentation task.
[0054] Optionally, as an implementation method, the network used in the downstream segmentation task is the STU-Net architecture, and the reconstruction main task encoder here is also based on the STU-Net architecture. For example, the reconstruction main task encoder can include five residual blocks to downsample the input data five times.
[0055] For the main reconstruction task decoder, this embodiment uses a partial STU-Net decoder and a Vision Transformer decoder that do not include the deep supervision module. Following the spatial divisions used in the masking process, the convolutional network output is reconstructed into small 3D image blocks. This is then transformed using a linear layer and positional encoding is added. This is then input into the decoder for image reconstruction.
[0056] It is understandable that the reconstruction main task decoder is only used for pre-training and not for downstream tasks.
[0057] Furthermore, in the main task of mask reconstruction, for the main task network consisting of the mask module, the reconstruction main task encoder, and the reconstruction main task decoder, its loss function can adopt the mean absolute error (L1) loss. The loss function of this network is only used to evaluate the reconstruction effect of the masked part, and the expression is as follows:
[0058]
[0059] Where, L res represents the loss function of the main task network, also known as the reconstruction task loss, x mask Represents the original image of the masked part, x' mask Represents the reconstruction result of the masked part, N mask Indicates the number of representative mask blocks.
[0060] Step 3: Construct a segmentation auxiliary task and establish an auxiliary task network in the segmentation auxiliary task to perceive the downstream segmentation task; the auxiliary task network includes a segmentation task encoder and a segmentation task decoder; input the original 3D medical image into the auxiliary task network, use the segmentation task encoder and segmentation task decoder to segment the original 3D medical image, and calculate the loss function of the auxiliary task network based on the segmentation results.
[0061] Generally speaking, the functions of auxiliary tasks include enhancing feature representation, promoting multi-task learning, and improving model robustness. This embodiment adds a segmentation auxiliary task during the training process. When applied in pre-training, it can provide downstream task goals, thereby achieving higher accuracy and efficiency in medical image segmentation tasks.
[0062] For details, please refer to Figure 2 In the auxiliary task network, the segmentation task encoder is shared with the reconstruction main task encoder. The segmentation head in the segmentation task decoder consists of a convolutional layer with a kernel size of 1, so that the auxiliary task shares the same task attributes with the downstream task, but does not share the label space as a reference for the downstream task.
[0063] Furthermore, for the auxiliary task network formed by the segmentation task encoder and the segmentation task decoder in series, its loss function can be the sum of the Dice coefficient loss and the cross entropy loss, as shown below:
[0064]
[0065] Where, L seg Represents the loss function of the auxiliary task network, also known as the auxiliary task loss, L Dice represents the Dice coefficient loss, L CE represents the cross entropy loss, m represents the number of voxels in the sample, and t i represents the target value of the i-th voxel, p i It represents the probability that the i-th voxel output by the network belongs to a certain category, M represents the number of samples, C represents the number of categories, and t j,c Indicates the target value that the jth sample belongs to category c, p j,c represents the probability that the jth sample belongs to category c.
[0066] Step 4: Jointly learn the mask reconstruction main task and the segmentation auxiliary task to simultaneously optimize the training of the main task network and the auxiliary task network, thereby completing the pre-training of 3D medical image reconstruction.
[0067] Specifically, the mask reconstruction main task and the segmentation auxiliary task are integrated to perform multi-task learning. By weighting the losses of the two branch tasks and optimizing the two objectives at the same time, the total loss L is:
[0068]
[0069] Where, L res Represents the loss function of the main task network, L seg represents the loss function of the auxiliary task network, λ res and λ seg Represent the weights of the loss functions of the main task network and the auxiliary task network respectively.
[0070] As can be seen, this embodiment gives greater weight to the reconstruction loss of the main task, ensuring that the model learns its core capabilities and acquires more general and representative features. A smaller weight is given to the auxiliary task loss, which prevents overfitting while maintaining the model's sensitivity to local information. In short, this strategy enriches the model's feature expression while ensuring the learning of core features, improving the model's generalization ability, training stability, and adaptability to downstream tasks.
[0071] This method constructs a primary reconstruction task and an auxiliary segmentation task within the pretrained network. By weighting the loss functions of the two branches, joint learning is achieved, optimizing both objectives simultaneously during training. The auxiliary segmentation task incorporates the implicit expertise of the downstream task, further clarifying the goals of network pretraining and subsequently contributing to fine-tuning and improving the segmentation accuracy of three-dimensional medical images. Furthermore, by increasing the task-awareness capabilities of the model in the pretrained task, segmentation accuracy can be improved without increasing the complexity of the downstream task network or the training cost.
[0072] From the above description, it can be understood that the main task network and the auxiliary task network together constitute the pre-trained network. After multiple rounds of training, the weights of the pre-trained network can be obtained, which can be used for downstream tasks.
[0073] Based on the same inventive concept, the second aspect of the present invention further provides a three-dimensional medical image segmentation method, which mainly includes:
[0074] S1: Build a segmentation network model.
[0075] Optionally, the segmentation network model may adopt a STU-Net model.
[0076] S2: Input the 3D medical image to be segmented into the trained segmentation network model to obtain the 3D medical image segmentation result. The segmentation network model is trained using the following method:
[0077] a) randomly initializing the segmentation head of the segmentation network model and loading the weights of the encoder part in the pre-trained network into the remaining layers of the segmentation network model to initialize the parameters of the segmentation network model; wherein the weights of the encoder part in the pre-trained network are obtained using the task-aware three-dimensional medical image reconstruction pre-training method proposed in the first aspect of the present invention.
[0078] b) Further fine-tune the parameters of the initialized segmentation network model through optimization training to obtain a trained segmentation network model.
[0079] Since the segmentation auxiliary task already contains the implicit expertise of the downstream task in the pre-training stage, the goal of network pre-training is clarified. Therefore, after applying the obtained encoder weights to the segmentation network model, the training of the segmentation network model can be completed through fine-tuning, and the trained network has high segmentation accuracy.
[0080] It is understandable that, during the training phase of the segmentation network model, the training set used is still the pre-processed original three-dimensional medical image.
[0081] In addition, the preprocessed images can be divided into a training set and a validation set according to a certain ratio (for example, 4:1). After completing the training using the training set, the validation set can be used to evaluate the performance of the network.
[0082] The proposed task-aware pre-training method for 3D medical image reconstruction first preprocesses the original 3D medical image. It then constructs a main task (mask reconstruction) and an auxiliary task (segmentation). A main task network is established for the main task, which selectively masks the preprocessed image based on a high-semantic-information masking strategy and achieves image reconstruction. An auxiliary task network is established for the auxiliary task (segmentation), which can perceive the downstream segmentation task and thus achieve image segmentation. Finally, joint learning of the main reconstruction and auxiliary tasks simultaneously optimizes both objectives, completing pre-training for 3D medical image reconstruction. This method employs a high-semantic-information masking strategy during pre-training, allowing the mask image to focus more on semantically rich foreground features, eliminating redundant background, enriching reconstruction information, reducing the difficulty of mask reconstruction, and improving training efficiency and effectiveness. Furthermore, this method introduces an auxiliary task branch, enhancing the task-awareness of the pre-trained network and clarifying the network's pre-training objectives. This allows for more efficient feature extraction of 3D medical images, thereby improving segmentation accuracy without increasing the complexity or training cost of the downstream task network.
[0083] The method provided in the embodiments of the present invention can be applied to electronic devices. Specifically, the electronic devices can be desktop computers, portable computers, smart mobile terminals, servers, etc. This is not limited here; any electronic device that can implement the present invention falls within the scope of protection of the present invention.
[0084] As for the electronic device embodiment, since it is basically similar to the method embodiment, the description is relatively simple, and the relevant parts can be referred to the partial description of the method embodiment.
[0085] Finally, it should be noted that the above description further illustrates the present invention in conjunction with specific preferred embodiments, and the present invention should not be considered to be limited to these descriptions. Those skilled in the art will be able to make several simple deductions or substitutions without departing from the scope of the present invention, all of which should be considered to fall within the scope of protection of the present invention.
Claims
1. A task-aware 3D medical image reconstruction pre-training method, characterized in that: include: Preprocessing the original three-dimensional medical image to obtain a preprocessed image; Constructing a mask reconstruction main task, and establishing a main task network within the mask reconstruction main task; the main task network includes a mask module, a reconstruction main task encoder, and a reconstruction main task decoder connected in series; inputting the preprocessed image into the main task network, using the mask module to selectively mask the preprocessed image based on a high semantic information masking strategy to obtain a masked image, reconstructing the masked image using the reconstruction main task encoder and the reconstruction main task decoder, and calculating a loss function of the main task network based on the reconstruction result; The mask module is used to selectively mask the preprocessed image based on a high semantic information masking strategy to obtain a masked image, specifically comprising: Dividing the preprocessed image into a plurality of non-overlapping image blocks, and dividing each image block into a plurality of two-dimensional slices; For each image block, the number of key points and contrast response value of each two-dimensional slice are calculated using the scale-invariant feature transformation algorithm to obtain the number of key points and contrast response value of the key points of the image block; Based on the number of key points and the key point contrast response value of each image block, the SIFT score of each image block is calculated; the calculation formula of the SIFT score is: ; Where, Represents an image block SIFT score, Represents an image block The number of key points, represents the weight of the contrast response value, Represents an image block The contrast response value, Represents the input image The maximum value of the contrast response value in all corresponding image blocks; Based on the mask probability of the current training round, image blocks with low SIFT scores are selected for masking, while image blocks with high SIFT scores are retained to obtain a mask image. The mask probability is determined using a dynamic learning strategy and is calculated as follows: ; Where, Indicates the The mask probability during round training, Indicates the current training round, represents the preset initial mask probability, represents the preset termination mask probability, represents the total number of training rounds; Constructing a segmentation auxiliary task, establishing an auxiliary task network within the segmentation auxiliary task to perceive the downstream segmentation task; the auxiliary task network includes a segmentation task encoder and a segmentation task decoder; inputting the original 3D medical image into the auxiliary task network, performing segmentation processing on the original 3D medical image using the segmentation task encoder and the segmentation task decoder, and calculating a loss function of the auxiliary task network based on the segmentation results; the network used in the downstream segmentation task is a STU-Net architecture, and the reconstruction main task encoder here is also based on the STU-Net architecture, and the reconstruction main task encoder includes five residual blocks to downsample the input data five times; The mask reconstruction main task and the segmentation auxiliary task are jointly learned to simultaneously optimize and train the main task network and the auxiliary task network, thereby completing the three-dimensional medical image reconstruction pre-training.
2. The task-aware 3D medical image reconstruction pre-training method according to claim 1, characterized in that: The reconstruction main task encoder adopts the STU-Net structure designed based on the downstream segmentation task, specifically including multiple residual blocks; The reconstruction main task decoder uses a partial STU-Net decoder and a VisionTransformer decoder that do not include a deep supervision module.
3. The task-aware 3D medical image reconstruction pre-training method according to claim 2, characterized in that: The segmentation task encoder is shared with the reconstruction main task encoder; the segmentation head in the segmentation task decoder consists of a convolutional layer with a kernel size of 1, so that the auxiliary task shares the same task properties with the downstream task.
4. The task-aware 3D medical image reconstruction pre-training method according to claim 1, characterized in that: In the mask reconstruction main task, the loss function of the main task network adopts the mean absolute error loss, which is expressed as follows: ; Where, represents the loss function of the main task network, represents the original image of the masked part, represents the reconstruction result of the masked part, Indicates the number of representative mask blocks.
5. The task-aware 3D medical image reconstruction pre-training method according to claim 1, wherein: In the segmentation auxiliary task, the loss function of the auxiliary task network is the sum of the Dice coefficient loss and the cross entropy loss, and the expression is as follows: ; Where, represents the loss function of the auxiliary task network, represents the Dice coefficient loss, represents the cross entropy loss, represents the number of voxels of the sample, Indicates the The target value of the voxel, The network output The probability that a voxel belongs to a certain class, represents the number of samples, represents the number of categories, Indicates the samples belong to the category The target value, Indicates the samples belong to the category probability.
6. The task-aware 3D medical image reconstruction pre-training method according to claim 1, characterized in that: During the joint learning process of the mask reconstruction main task and the segmentation auxiliary task, the total loss for: ; Where, represents the loss function of the main task network, represents the loss function of the auxiliary task network, and Represent the weights of the loss functions of the main task network and the auxiliary task network respectively.
7. A three-dimensional medical image segmentation method, characterized in that: include: Build a segmentation network model; Input the 3D medical image to be segmented into a trained segmentation network model to obtain a 3D medical image segmentation result; wherein the segmentation network model is trained using the following method: Randomly initializing the segmentation head of the segmentation network model, and loading the weights of the encoder part in the pre-trained network into the remaining layers of the segmentation network model to initialize the parameters of the segmentation network model; the weights of the encoder part in the pre-trained network are obtained using the task-aware three-dimensional medical image reconstruction pre-training method according to any one of claims 1 to 6; The parameters of the initialized segmentation network model are further fine-tuned through optimization training to obtain a trained segmentation network model.