Three-dimensional medical image reconstruction pre-training method based on task perception
By adopting a joint learning method of high semantic information masking strategy and task-awareness in three-dimensional medical image reconstruction, the problems of redundant background and invalid information are solved, and more efficient training and more precise segmentation effects are achieved.
Patent Information
- Application Number
- CN202411619355.9
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Priority Date
- 2024-10-11
- Filing Date
- 2024-11-13
- Publication Date
- 2025-03-11
- Estimated Expiration
- 2044-11-13
AI Technical Summary
The existing three-dimensional medical image reconstruction methods have redundant background and invalid information, which leads to an increase in model training costs, affecting reconstruction effect and segmentation accuracy, and the random mask and fixed mask probability strategies will cause the model to converge to suboptimal solutions or insufficient learning.
The pre-training method of three-dimensional medical image reconstruction based on task perception is adopted, and selective masking is performed through high semantic information masking strategy to construct the main task of mask reconstruction and segmentation auxiliary task, realize joint learning, and dynamically adjust the mask probability to control task difficulty.
It effectively reduces the task difficulty of mask reconstruction, improves training efficiency and effect, improves segmentation accuracy, and enhances the generalization ability and task perception ability of the model.
Smart Images

Figure CN119600034A_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 the advancement of modern medical technology, 3D medical images have become an important tool for diagnosing and treating various diseases. They can provide detailed anatomical information and assist doctors in disease assessment, surgical planning, and efficacy monitoring. However, manual annotation of medical image data requires a lot of expertise and is very time-consuming. The segmentation accuracy is also easily affected by subjective factors. Therefore, computer-assisted automatic segmentation technology has attracted much 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 Number 202211383092.7) discloses a lung medical image segmentation method based on reconstruction pre-training. This method learns feature representation through unsupervised jigsaw reconstruction pre-training, combined with the idea of transfer learning, improves the robustness and generalization of the model, and reduces the model training time. Patent Document 2 (Application Number 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 constructs a pre-training agent task for deblurred mask image modeling based on this, so that the latent representation learned in the pre-training stage includes additional attention to the details of the ultrasound image.
[0004] However, in the implementation process of the method in Patent Document 1, each image needs to be gridded and the order needs to be disrupted, which results in a lot of redundant background in the image, causing the model to input too much invalid information, increasing the training cost and reconstruction difficulty, affecting the reconstruction effect, and further affecting the segmentation accuracy; and in the subsequent segmentation model, it is necessary to continue to add reconstruction branches, increasing the complexity of downstream tasks. The random masking method used in the method in Patent Document 2 will retain a lot 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. At the same time, 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] In order to solve at least one problem existing in the prior art, the present invention provides a task-aware three-dimensional medical image reconstruction pre-training method. The technical problem to be solved by the present invention is achieved by 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 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 in sequence; 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, use the reconstruction main task encoder and the reconstruction main task decoder to reconstruct the masked image, and calculate the loss function of the main task network based on the reconstruction result;
[0010] Constructing a segmentation auxiliary task, and establishing 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; inputting the original 3D medical image into the auxiliary task network, using the segmentation task encoder and the segmentation task decoder to segment the original 3D medical image, and calculating the loss function of the auxiliary task network according to the segmentation result;
[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 proposes a three-dimensional medical image segmentation method, comprising:
[0013] Build a segmentation network model;
[0014] The three-dimensional medical image to be segmented is input into the trained segmentation network model to obtain the three-dimensional medical image segmentation result; wherein the segmentation network model is trained using the following method:
[0015] The segmentation head of the segmentation network model is randomly initialized, and the weights of the encoder part in the pre-trained network are loaded 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 by using the task-aware three-dimensional 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 task-aware 3D medical image reconstruction pre-training method proposed in the present invention first pre-processes the original 3D medical image, then constructs the mask reconstruction main task and the segmentation auxiliary task, establishes a main task network in the mask reconstruction main task, and the main task network selectively masks the pre-processed image based on the high semantic information mask strategy and realizes image reconstruction; establishes an auxiliary task network in the segmentation auxiliary task, and 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, the method adopts a high semantic information mask strategy in 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, the 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 the pre-training process, the present invention adopts a dynamically changing learning strategy to control the difficulty of the mask reconstruction task. Compared with a fixed mask probability, the present invention can prevent the network from converging to a suboptimal solution in the early stage of training, and achieve progressive learning from easy to difficult. In addition, the dynamic learning strategy not only greatly reduces the time cost of parameter tuning, but also avoids the differences in fixed mask probabilities on different data sets, thereby increasing the generalization ability of the model.
[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 It is a flowchart of a task-aware three-dimensional medical image reconstruction pre-training method provided in an embodiment of the present invention;
[0022] Figure 2 It is a flowchart of constructing a mask reconstruction main task and a segmentation auxiliary task provided by an embodiment of the present invention. DETAILED DESCRIPTION
[0023] The following will be combined with the drawings in the embodiments of the present invention to clearly and completely describe the technical solutions in the embodiments of the present invention. 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 creative work 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 : is a flowchart 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, so 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 three-dimensional 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 an input image for the pre-training stage, where 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 in sequence; 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, and 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 also Figure 2 , Figure 2 1 is a flowchart of constructing the mask reconstruction main task and the segmentation auxiliary task provided by an 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 out 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 according to 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 an 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 contrast response value of the key points of each two-dimensional slice are calculated using a scale-invariant feature transformation algorithm to obtain the number of key points and the contrast response value of the key points 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 can be 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 point. 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, 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 by using the following formula.
[0043]
[0044] In the formula, x p _score represents the image block x p SIFT score, 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 to measure the feature richness of each image block. Obviously, for an image block with a small number of key points and a low contrast response value, its SIFT score is low, indicating that the features of the image block are not rich enough. Therefore, this part of the image block is masked according to the mask probability of the current training round. Correspondingly, for an image block with a large number of key points and a high contrast response value, its SIFT score is high, indicating that the features of the image block are relatively rich, and this part of the image block is retained and entered into the encoder for feature learning.
[0047] The high semantic information masking strategy designed by the present invention pays more attention to the foreground with rich semantic information, focuses on key anatomical parts, uses a scale-invariant feature transformation algorithm to measure the feature richness, calculates the number of key points in each image block and the contrast response value of these key points, and selects feature-rich areas for learning. Compared with jigsaw reconstruction, the present invention excludes redundant background; compared with random masking, all areas in medical images are uniformly sampled. While improving the segmentation accuracy of downstream tasks, the present invention also improves the efficiency of pre-training.
[0048] Furthermore, for the determination of the mask probability, this embodiment abandons the existing fixed mask probability and designs a dynamic learning strategy to adjust the mask probability. Since in the masking process, the smaller the mask probability, the more image features are retained. Therefore, in the dynamic learning strategy designed in this embodiment, the mask probability increases inversely with the training epoch, so that a larger mask probability is used in the early stage of pre-training, that is, less information is input into the network learning, and then the information is gradually enriched to implement a learning strategy 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] In the formula, 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 the pre-training process, the present invention uses a dynamically changing learning strategy to control the difficulty of the mask reconstruction task. Compared with a fixed mask probability, the present invention can prevent the network from converging to a suboptimal solution in the early stages of training, and achieve progressive learning from easy to difficult. In addition, the dynamic learning strategy not only greatly reduces the time cost of parameter tuning, but also avoids the differences in fixed mask probabilities on different data sets, thereby increasing the generalization ability of the model.
[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 by 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 may include five residual blocks to downsample the input data five times.
[0055] For the reconstruction main task decoder, this embodiment can use a partial STU-Net decoder and a Vision Transformer decoder that do not include a deep supervision module. According to the spatial division in the masking process, the output after the convolution network is reconstructed into small three-dimensional image blocks, the dimensions are converted through a linear layer and position encoding is added, and the blocks are 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 mask reconstruction main task, for the main task network composed 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 Represents the number of 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 three-dimensional medical image into the auxiliary task network, use the segmentation task encoder and the segmentation task decoder to segment the original three-dimensional medical image, and calculate the loss function of the auxiliary task network based on the segmentation result.
[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 convolutional layers with kernel size 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 connected in series, its loss function can be the sum of the Dice coefficient loss and the cross entropy loss, expressed as follows:
[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 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 represents 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 main task of mask reconstruction and the auxiliary task of segmentation are integrated for 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 function of the main task network and the auxiliary task network respectively.
[0070] It can be seen that this embodiment gives a larger weight to the reconstruction main task loss, which can ensure the learning of the core capabilities of the model and learn more general and representative features. A smaller weight is given to the auxiliary task loss, which can avoid overfitting the model while keeping it sensitive to local information. In short, this strategy can enrich the feature expression of the model while ensuring the learning of core features, and improve the model's generalization ability, training stability, and adaptability to downstream tasks.
[0071] The present invention constructs the reconstruction main task and the segmentation auxiliary task in the pre-training network, and realizes joint learning by weighting the loss functions of the two branches, and optimizes the two goals simultaneously during the training process. The segmentation auxiliary task contains the invisible professional knowledge of the downstream task, which further clarifies the goal of network pre-training, and then acts on fine-tuning to improve the segmentation accuracy of three-dimensional medical images. In addition, increasing the task perception ability of the model in the pre-training task can improve the segmentation accuracy without increasing the complexity of the downstream task network and 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 and clarifies the goal of network pre-training, the training of the segmentation network model can be completed by fine-tuning after the obtained encoder weights are applied to the segmentation network model, and the trained network has a higher 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 the training is completed using the training set, the validation set is used to evaluate the performance of the network.
[0082] The present invention proposes a task-aware 3D medical image reconstruction pre-training method. First, the original 3D medical image is pre-processed, and then a mask reconstruction main task and a segmentation auxiliary task are constructed. In the mask reconstruction main task, a main task network is established. The main task network selectively masks the pre-processed image based on a high semantic information masking strategy and realizes image reconstruction. In the segmentation auxiliary task, an auxiliary task network is established. 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, the method adopts a high semantic information masking strategy in 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 task difficulty of mask reconstruction, and improves the training efficiency and effect. On the other hand, the 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.
[0083] The method provided in the embodiment of the present invention can be applied to electronic devices. Specifically, the electronic device can be: a desktop computer, a portable computer, an intelligent mobile terminal, a server, etc. This is not limited here, and any electronic device that can implement the present invention belongs to the protection scope 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 contents are further detailed descriptions of the present invention in combination with specific preferred embodiments, and it cannot be determined that the specific implementation of the present invention is limited to these descriptions. For ordinary technicians in the technical field to which the present invention belongs, several simple deductions or substitutions can be made without departing from the concept of the present invention, which should be regarded as falling 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 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 in sequence; 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, using the reconstruction main task encoder and the reconstruction main task decoder to reconstruct the masked image, and calculating the loss function of the main task network according to the reconstruction result; Constructing a segmentation auxiliary task, establishing 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; inputting the original three-dimensional medical image into the auxiliary task network, using the segmentation task encoder and the segmentation task decoder to segment the original three-dimensional medical image, and calculating the loss function of the auxiliary task network according to the segmentation result; 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 pre-training of three-dimensional medical image reconstruction.
2. The task-aware 3D medical image reconstruction pre-training method according to claim 1, characterized in that: 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 including: 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 a scale-invariant feature transformation algorithm to obtain the number of key points and contrast response value of the key points of the image block; Calculate the SIFT score of each image block based on the number of key points and the key point contrast response value of each image block; Based on the mask probability of the current training round, image patches with low SIFT scores are selected for masking, while image patches with high SIFT scores are retained to obtain a mask image.
3. The task-aware 3D medical image reconstruction pre-training method according to claim 2, characterized in that: The calculation formula of the SIFT score is: In the formula, 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.
4. The task-aware 3D medical image reconstruction pre-training method according to claim 2, characterized in that: The mask probability is determined using a dynamic learning strategy, and the calculation formula is as follows: In the formula, 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.
5. The task-aware 3D medical image reconstruction pre-training method according to claim 1, characterized in that: The reconstruction main task encoder adopts a 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.
6. The task-aware 3D medical image reconstruction pre-training method according to claim 5, 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.
7. 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, L res represents the loss function of the main task network, x mask represents the original image of the masked part, x' mask Represents the reconstruction result of the masked part, N mask Represents the number of mask blocks.
8. The task-aware 3D medical image reconstruction pre-training method according to claim 1, characterized in that: 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, L seg represents the loss function of the auxiliary task network, 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 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 represents 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.
9. The task-aware 3D medical image reconstruction pre-training method according to claim 1, characterized in that: In the joint learning process of the mask reconstruction main task and the segmentation auxiliary task, the total loss L is: 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 function of the main task network and the auxiliary task network respectively.
10. A three-dimensional medical image segmentation method, characterized in that: include: Build a segmentation network model; The three-dimensional medical image to be segmented is input into a trained segmentation network model to obtain a three-dimensional medical image segmentation result; wherein the segmentation network model is trained using the following method: Randomly initialize the segmentation head of the segmentation network model, and load the weights of the encoder part in the pre-trained network to 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 by using the task-aware three-dimensional medical image reconstruction pre-training method according to any one of claims 1 to 9; The parameters of the initialized segmentation network model are further fine-tuned through optimization training to obtain a trained segmentation network model.
Citation Information
Patent Citations
Lung medical image segmentation method based on reconstruction pre-training
CN115690115A
Training method of medical image processing model and medical image processing method
CN116883288A
3D medical image segmentation model establishment method based on mask modeling and application thereof
CN116664588A
3D self-supervised pre-training method based on large-scale language-image model guidance
CN116681107A
Medical image segmentation method and device based on multi-task self-supervised learning, and medium
CN117576033A
Cited By
Image data reconstruction method and system based on spatial region perception and mask guidance
CN121639959A
Image Data Reconstruction Method and System Based on Spatial Region Awareness and Mask-Guided Design
CN121639959B