Weak supervision parameter efficient fine tuning whole heart segmentation method based on uncertainty guidance

Abstract

According to the weak supervision parameter efficient fine tuning whole heart segmentation method based on uncertainty guidance, different variants of an input image are generated through random disturbance, the consistency loss is calculated in combination with prediction results of a teacher model and a student model, meanwhile, uncertainty estimation of the prediction results is introduced, and the models are guided to pay attention to difficult areas in a segmentation task. In the fine tuning process, a small number of pixel-level labeling samples are utilized to directly optimize model parameters, and an efficient parameter fine tuning technology is adopted to adapt to the pre-trained SAM, so that general knowledge learned in the field of natural images is fully reserved. By designing a lightweight fine tuning mechanism, the method can complete the efficient migration of the model only by training a small number of newly added parameters, significantly reduces the demand of computing resources, and improves the segmentation precision and robustness of the model in a weak annotation scene at the same time. The method shows excellent transfer learning ability, and an efficient and practical solution is provided for medical image segmentation in a resource limited environment.

Description

Technical Field

[0001] The present invention belongs to the technical field of medical image segmentation, and specifically is a whole heart segmentation method based on uncertainty-guided weakly supervised parameter efficient fine-tuning, which combines weakly supervised learning with efficient parameter fine-tuning strategy to improve the segmentation accuracy of the whole heart structure. Background Art

[0002] In recent years, the mortality rate of cardiovascular diseases has continued to rise, highlighting the urgent need for the prevention and treatment of heart diseases. With the rapid development of imaging technology, especially non-invasive imaging methods such as computed tomography (CT), cardiac ultrasound (US), and magnetic resonance imaging (MR), they have become the key support for the diagnosis of heart diseases. By performing high-precision segmentation of cardiac medical images, not only can more accurate diagnosis of heart diseases be achieved, but also intelligent surgery and other treatment plans can be implemented.

[0003] In the field of medical image segmentation, deep learning models have shown great potential due to their ability to identify complex imaging features, but these methods require sufficient data to fit a huge parameter space. However, it is difficult to create accurate annotated datasets on a large scale in the medical field because detailed pixel-level annotations are usually only available through highly specialized personnel. In order to solve the problem of scarce pixel-level annotations, many semi-supervised methods that use a small amount of annotated data and a large amount of unlabeled data have been proposed to improve the performance and generalization ability of the model. The scarcity of high-quality pixel-level annotated data still leads to some semi-supervised algorithms over-relying on unlabeled data, which limits their effectiveness and applicability in actual clinical settings. In current research, how to efficiently utilize limited pixel-level annotated data remains a key issue that needs to be solved, which seriously restricts the actual application and effect of existing methods in clinical settings. As an effective strategy, weakly supervised learning can achieve high-precision segmentation results under the condition of limited annotated data, showing broad application prospects.

[0004] The application of transfer learning in the field of medical image segmentation is developing rapidly and has become an important strategy to deal with data scarcity and difficulty in labeling in medical image analysis. Since medical images usually face challenges such as limited sample size, high labeling costs, and modality diversity, transfer learning can effectively improve the initial performance and learning efficiency on specific medical image tasks by borrowing model parameters trained on large-scale datasets. Specifically, transfer learning fine-tunes the existing basic model, thereby utilizing the rich features learned by the pre-trained model to handle specific medical image segmentation tasks. Therefore, the present invention proposes a whole-heart segmentation method with efficient fine-tuning of weakly supervised parameters based on uncertainty guidance. Summary of the invention

[0005] In view of the deficiencies in the prior art, the technical problem that the present invention intends to solve is to propose a whole heart segmentation method based on uncertainty-guided weakly supervised parameters with efficient fine-tuning.

[0006] The technical solution adopted by the present invention to solve the technical problem is as follows:

[0007] In a first aspect, the present invention provides a whole heart segmentation method based on uncertainty-guided weakly supervised parameter efficient fine-tuning, the segmentation method comprising the following steps:

[0008] Step 1: Obtain several 3D whole heart structure CT images and MR images, cut the 3D images into 2D slices, crop and fill them to a uniform size, and preprocess the images to form a data set. Each 2D slice has a pixel-level label. Extract the pixel-level labels in the 2D slices, lock the region of interest, set the perturbation range, and randomly simulate and generate bounding boxes for the 2D slices within the perturbation range, which are used as weak annotations of the bounding boxes of the 2D slices.

[0009] Step 2: Build a segmentation model for the whole heart medical image segmentation task;

[0010] The segmentation model is based on the SAM model and includes an encoder, a mask decoder and a prompt encoder. The encoder is composed of α convolution fine-tuning modules and G-α frequency domain fine-tuning modules connected in series.

[0011] The convolution fine-tuning module includes layer normalization, multi-head attention, layer normalization and multi-layer perceptron connected in sequence, and the input of the first layer normalization is added to the output of the multi-head attention to obtain feature A, and feature A is processed by the second layer normalization and multi-layer perceptron to output feature B, and feature A is processed by the convolutional spatial local perception adapter to obtain feature C, and features A, B, and C are added and input into the next level module; the convolutional spatial local perception adapter includes a linear layer, a ReLu activation function, a convolutional layer, a ReLu activation function and a linear layer connected in sequence;

[0012] The frequency domain fine-tuning module includes layer normalization, frequency domain feature extraction layer, layer normalization, and multi-layer perceptron connected in sequence, and the input of the frequency domain fine-tuning module is added to the output of the multi-layer perceptron of the module as the output of the frequency domain fine-tuning module; the frequency domain feature extraction layer includes Fourier transform and inverse Fourier transform, and the result of Fourier transform is multiplied by the learnable weight and then subjected to inverse Fourier transform to obtain the output of the frequency domain feature extraction layer;

[0013] Step 3: Segmentation model training

[0014] The SAM without convolutional spatial local perception adapter and frequency domain feature extraction layer is pre-trained using natural images, and the weight parameters of the pre-trained SAM are used as the initial parameters of the segmentation model. The segmentation model is then fine-tuned using the dataset to achieve migration from the field of natural image segmentation to the field of medical image segmentation to adapt to the data distribution and characteristics of the whole heart segmentation task; during the fine-tuning training process, all parameters in the encoder except the convolutional spatial local perception adapter and frequency domain feature extraction layer are frozen, and only these two specific layers are trained, while the prompt encoder and mask decoder are fine-tuned;

[0015] The dataset is divided into two parts: one part uses complete pixel-level annotations, and the other part is used as unlabeled data. For all input images, whether they have pixel-level annotations or not, their corresponding bounding box weak annotations are fed into the hint encoder as additional hint information to guide the model to make more accurate predictions of the target area.

[0016] For image data with pixel-level annotations, the student model calculates supervised loss by comparing the predicted results with the actual pixel-level annotations to directly optimize the segmentation accuracy. For unlabeled image data, two variants are generated by adding random perturbations to them and input into the teacher model and student model respectively. The output results of the teacher model and the student model are used to calculate the consistency loss guided by uncertainty.

[0017] During fine-tuning training, the supervised loss L S The uncertainty-guided consistency loss is used as a joint loss to optimize the entire network; both the student model and the teacher model adopt a segmentation model structure;

[0018] Step 4: Use the trained student model as the final segmentation model for the whole heart medical image segmentation task.

[0019] Furthermore, based on the uncertainty-guided consistency loss L C It is expressed as formula (8):

[0020]

[0021] Among them, I is the sigmoid function, X U(i) represents unlabeled data, θ, θ′ represent the weights of the student model and the teacher model respectively, P represents the bounding box weak label used to prompt the location information, and f θ (X U(i) ,P) and f θ′ (X U(i) ,P) are the prediction results of the student model and the teacher model for the i-th sample; u i is the uncertainty of the prediction results of the student model, obtained by Monte Carlo estimation; H is the uncertainty threshold.

[0022] Furthermore, the formula for the supervised loss is:

[0023]

[0024] Among them, FL is FocalLoss loss, w FL Represents the weight, which is used to adjust the importance of FocalLoss loss; L Iou represents IOU loss, scale Iou Indicates L Iou Scaling factor of L Dice represents Dice loss; p t Represents the probability of the model correctly predicting each category; α t represents the weight factor of each sample in the loss function, which is used to adjust the importance of different categories; γ is the focusing parameter, which is used to reduce the contribution of easy-to-classify samples to the total loss; f θ (X L(i) ,P) represents the prediction result of the student model for labeled data; μ represents a constant to prevent the denominator from being 0; Y L(i) ∈{0,1} H×W×C Represents X L(i) Tags, X L(i) It is pixel-level annotation data, and P represents the weak label of the bounding box used to provide location information.

[0025] Furthermore, the whole heart structure refers to the main anatomical structure of the heart, including the myocardium, left atrium, left ventricle, right atrium, right ventricle, ascending aorta and pulmonary artery.

[0026] Furthermore, the convolutional layer in the convolutional spatial local perception adapter has 768 channels and the convolution kernel size is 3×3.

[0027] Furthermore, the frequency domain feature extraction layer processes the input data in the following steps: first, the input image feature data is converted from the spatial domain to the frequency domain through a two-dimensional real fast Fourier transform to obtain spectral features; then, the spectral features are gated and the feature representation capability is enhanced using the learned complex weights; after spectral gating, the data is converted back to the spatial domain using an inverse fast Fourier transform to restore it to the shape of the original feature map.

[0028] Furthermore, the value of G is 10-15, and the value of α is 4-10.

[0029] In a second aspect, the present invention provides a computer-readable storage medium having a computer program stored thereon, which can implement the steps of the method when executed by a processor.

[0030] Compared with the prior art, the present invention has the following beneficial effects:

[0031] (1) The present invention proposes an efficient fine-tuning method for weakly supervised parameters, which significantly reduces the reliance on high-precision pixel-level annotations in the segmentation of whole-heart medical images due to insufficient annotations. This method generates different variants of the input image through random perturbations, calculates the consistency loss based on the prediction results of the teacher model and the student model, and introduces uncertainty estimates of the prediction results to guide the model to focus on the difficult areas in the segmentation task. During the fine-tuning process, a small number of pixel-level annotated samples are used to directly optimize the model parameters, and an efficient parameter fine-tuning technique is used to adapt the pre-trained SAM, fully retaining the general knowledge learned in the field of natural images. By designing a lightweight fine-tuning mechanism, this method only needs to train a small number of new parameters to complete the efficient migration of the model, significantly reducing the demand for computing resources, while improving the segmentation accuracy and robustness of the model in weakly labeled scenarios. This method demonstrates excellent transfer learning capabilities and provides an efficient and practical solution for medical image segmentation in resource-constrained environments.

[0032] (2) Aiming at the shortcomings of the traditional Transformer model in local feature modeling, the present invention designs an innovative convolutional spatial local perception adapter, which improves the model's ability to capture local spatial features by integrating the spatial inductive bias of the convolution operation and the global self-attention mechanism of the Transformer. The adapter is embedded in the encoder of SAM in an efficient fine-tuning manner, and only a small number of new parameters need to be trained to optimize the model, avoiding the consumption of computing resources caused by large-scale parameter updates. In the design, the adapter fully utilizes the advantages of the convolution operation in capturing local patterns, and complements the Transformer's long-distance dependency modeling ability, thereby achieving collaborative modeling of the relationship between global and local features. Compared with the existing technology, the adapter improves the model performance without significantly increasing the complexity and computational overhead, has high flexibility and versatility, and can be widely used in a variety of segmentation tasks. It significantly enhances the accuracy and robustness of cardiac image segmentation, and provides an efficient solution for the segmentation of complex medical images.

[0033] (3) The present invention innovatively proposes a frequency domain feature extraction layer to replace the traditional multi-head self-attention mechanism, and realizes the optimized design of the encoder by combining the convolution fine-tuning module and the frequency domain fine-tuning module. The frequency domain feature extraction layer extracts the global frequency domain information of the image through Fourier transform, replacing the calculation of the attention layer. This design utilizes the global representation characteristics of Fourier transform, which can effectively reduce the computational complexity of the multi-head self-attention mechanism when modeling long-distance dependencies, thereby significantly improving the model efficiency and reducing memory usage. In addition, the convolution fine-tuning module introduces a convolutional spatial local perception adapter to enhance the model's ability to capture local features, thereby making up for the shortcomings of the attention layer in processing spatial local information. By applying the convolution fine-tuning module and the frequency domain fine-tuning module to different parts of the encoder respectively, the model realizes the collaborative modeling of global and local features, significantly enhancing the model's ability to capture fine-grained features, and effectively improving the accuracy and efficiency of the segmentation task. BRIEF DESCRIPTION OF THE DRAWINGS

[0034] Figure 1 It is a schematic diagram of the overall process of the uncertainty-guided weakly supervised parameter efficient fine-tuning whole heart segmentation method of the present invention;

[0035] Figure 2 It is a structural schematic diagram of the segmentation model in the present invention;

[0036] Figure 3 Schematic diagram of the structure of the convolution fine-tuning block in the present invention;

[0037] Figure 4 It is a structural schematic diagram of the spectrum fine-tuning block in the present invention;

[0038] Figure 5 This is a structural diagram of the uncertainty-guided weakly supervised training framework of the present invention. DETAILED DESCRIPTION

[0039] The technical solution of the present invention is described in detail below in conjunction with the accompanying drawings and specific implementation methods, but the protection scope of the present application is not limited thereto.

[0040] The present invention is based on uncertainty-guided weakly supervised parameter efficient fine-tuning whole heart segmentation method (see Figure 1 ), including the following steps:

[0041] Step 1: Obtain several 3D whole heart structure CT images and MR images, cut the 3D images into 2D slices, crop and fill to a uniform size, and preprocess the images to form a data set. Each 2D slice has a pixel-level label. Extract the pixel-level labels in the 2D slices, lock the region of interest, set the perturbation range, and randomly simulate and generate bounding boxes for the 2D slices within the perturbation range as the bounding box weak annotations of the 2D slices. Extract the existing pixel-level labels in the data set to define the bounding box of the region of interest in the image, so as to simulate and generate bounding box weak annotations and form bounding box weak labels.

[0042] Step 2: Build a segmentation model for the whole heart medical image segmentation task.

[0043] Fine-tuning training is performed using the Mean Teacher (MT) architecture that includes a student model and a teacher model. The student model and the teacher model have the same structure, both of which use the Segment Anything Model (SAM) as the backbone, including an encoder, a decoder, and a prompt encoder. The segmentation model is based on the SAM model, and the encoder part is improved. The improved encoder consists of α convolutional fine-tuning modules and G-α frequency domain fine-tuning modules in series;

[0044] The convolution fine-tuning module inserts a convolutional spatial local perception adapter based on the SAM encoder. As an effective way to achieve fine-tuning, the adapter combines the spatial inductive bias of the convolutional layer and the self-attention mechanism of the Transformer to enhance the encoder's perception of spatial local information, effectively capture and integrate the relationship between the image spatial local information and global features, thereby improving the accuracy and performance of the heart segmentation task.

[0045] The frequency domain fine-tuning module replaces the multi-head self-attention with the frequency domain feature extraction layer, making full use of the properties of Fourier transform to reduce the spatiotemporal complexity of the multi-head self-attention to O(NlogN), which reduces the complexity and memory usage of the model to a certain extent and improves the model training and reasoning speed. All attention layers have the disadvantage of not being able to accurately capture local features, and all spectrum layers have the disadvantage of not being able to accurately process global image attributes or semantic features. Therefore, a spectrum fine-tuning module that uses a frequency domain feature extraction layer to replace the multi-head self-attention mechanism is proposed to improve these shortcomings.

[0046] The convolution fine-tuning module includes layer normalization, multi-head attention, layer normalization and multi-layer perceptron connected in sequence, and the input of the first layer normalization is added to the output of the multi-head attention to obtain feature A, and feature A is processed by the second layer normalization and multi-layer perceptron to output feature B, and feature A is processed by the convolutional spatial local perception adapter to obtain feature C, and features A, B, and C are added and input into the next level module; the convolutional spatial local perception adapter includes a linear layer, a ReLu activation function, a convolutional layer, a ReLu activation function and a linear layer connected in sequence;

[0047] The frequency domain fine-tuning module includes layer normalization, frequency domain feature extraction layer, layer normalization, and multi-layer perceptron connected in sequence. The input of the frequency domain fine-tuning module is added to the output of the multi-layer perceptron of the module as the output of the frequency domain fine-tuning module; the frequency domain feature extraction layer includes Fourier transform and inverse Fourier transform. The result of Fourier transform is multiplied by the learnable weight and then subjected to inverse Fourier transform to obtain the output of the frequency domain feature extraction layer.

[0048] Step 3: Combine the pixel-level annotation data and the bounding box weak annotation data to perform weakly supervised fine-tuning training on the segmentation model. The specific process is as follows:

[0049] After the model is built, the pre-trained weights of SAM are loaded to utilize its general feature extraction capabilities and initial parameters obtained on natural images to perform transfer learning on whole-heart medical images. The present invention achieves migration from the field of natural image segmentation to the field of medical image segmentation through fine-tuning. It mainly uses an innovative encoder composed of a convolution fine-tuning module and a spectrum fine-tuning module to adjust the segmentation model to adapt to the data distribution and characteristics of the whole-heart segmentation task, thereby improving the performance of the segmentation model in the new field. During the fine-tuning training process, all parameters in the encoder part except the convolutional spatial local perception adapter and the frequency domain feature extraction layer are frozen, and only these two specific layers are trained. At the same time, the prompt encoder and decoder are fine-tuned to ensure segmentation accuracy.

[0050] First, the dataset is divided into two parts: one with complete pixel-level annotations and the other as unlabeled data. For all input images, whether or not they have pixel-level annotations, their corresponding bounding box weak annotations are fed into the hint encoder as additional hint information to guide the model to make more accurate predictions of the target area.

[0051] For image data with pixel-level annotations, the student model calculates the supervised loss by predicting the results and the actual pixel-level annotations to directly optimize the segmentation accuracy. For unlabeled image data, two variants are generated by adding random perturbations and sent to the teacher model and student model respectively. The output results of the teacher model and the student model are used to calculate the uncertainty-guided consistency loss. The average teacher strategy based on uncertainty estimation is adopted, that is, the uncertainty of the prediction results is used to guide the model to pay more attention to the areas with higher prediction uncertainty, so as to achieve a more robust optimization effect on unlabeled data.

[0052] Specifically, the unlabeled data with different noise perturbations and their corresponding bounding box weak annotations are input into the teacher model and the student model, the Monte Carlo estimation is enabled and T forward propagations are performed to generate the unlabeled data X′ after the perturbation is added U The prediction results are obtained by dividing the uncertainty through the prediction entropy approximation model and guiding the calculation of the consistency loss L based on the uncertainty. C . The predicted entropy is used to approximate the model segmentation uncertainty u, as shown in formula (7):

[0053]

[0054] In this embodiment, T=8, which means it is executed 8 times, and each forward propagation will produce different outputs based on different noise disturbances and dropout randomness.

[0055] Uncertainty-guided consistency loss L C It can be expressed as formula (8):

[0056]

[0057] Among them, I is the sigmoid function, X U(i) represents unlabeled data, θ, θ′ represent the weights of the student model and the teacher model respectively, P represents the bounding box weak label used to prompt the location information, and f θ (X U(i) ,P) and f θ′ (X U(i) ,P) are the prediction results of the student model and the teacher model for the i-th sample, u i is the uncertainty of the prediction results of the student model, obtained by Monte Carlo estimation, and H is the uncertainty threshold, which is used to screen the high uncertainty area.

[0058] During the training process, a certain proportion of pixel-level annotated data is used to calculate the supervised loss L S, and combined with the consistency loss calculated by unlabeled data, the whole network is optimized in a joint loss manner. Through weakly supervised fine-tuning training, the student model fully utilizes the potential of weakly labeled data under the condition of limited pixel-level labeled data and completes the effective segmentation of whole heart medical images.

[0059] In this embodiment, the value of G is 10-15, and the value of α is 4-10, which can be determined according to training data.

[0060] Example 1

[0061] The present invention provides a whole heart segmentation method based on uncertainty-guided weakly supervised parameter efficient fine-tuning, which specifically includes the following steps:

[0062] Step 1: Obtain several 3D CT and MR image data of the whole heart structure, including the main anatomical structures of the heart, such as the myocardium (Myo), left atrium (LA), left ventricle (LV), right atrium (RA), right ventricle (RV), ascending aorta (Aorta), and pulmonary artery (PA). To adapt the network structure, the 3D image data is cut into 2D slices, and the slices are uniformly processed, including cropping, filling, and size normalization operations. All images are adjusted to a fixed size of 256×256 by the bilinear interpolation method to ensure the consistency of data input. At the same time, the images are standardized and preprocessed to improve the training efficiency and segmentation accuracy of the model; in this process, the existing pixel-level labels in the dataset are extracted, and the regions of interest in the image are randomly defined with bounding boxes to simulate the weak annotation of the generated bounding boxes.

[0063] Step 2: Build a segmentation model for the whole heart medical image segmentation task, using the Segment Anything Model (SAM) as the backbone network, including an encoder, a decoder, and a prompt encoder.

[0064] The main task of the encoder of the segmentation model is to extract effective multi-level feature information from the input image. It contains α convolution fine-tuning modules and G-α frequency domain fine-tuning modules. Figure 2 As shown, the segmentation accuracy and parameter efficiency of the model can be significantly improved. In this embodiment, G=12, α=8.

[0065] Convolutional fine-tuning module such as Figure 3As shown on the left, this is achieved by integrating convolutional spatial locality-aware adapters in parallel within the Transformer block. Non-adapter layers are frozen during fine-tuning, and only adapters with a small number of parameters and a lightweight mask decoder are open for fine-tuning to minimize parameter updates, thereby significantly accelerating the model adaptation process without affecting segmentation accuracy.

[0066] like Figure 3 As shown on the right, the convolutional spatial local-aware adapter introduces a spatial inductive bias through convolution operations, thereby optimizing the base model in a fine-tuning manner. The core component of the adapter is a 768-channel 2D convolutional layer that uses a 3×3 convolution kernel to keep the spatial dimension unchanged. The convolutional spatial local-aware adapter can enhance the network's sensitivity to local detail features while maintaining low computational complexity. The convolutional spatial local-aware adapter processes the input image feature data in the following steps: First, a linear dimensionality reduction layer is used to flatten the input data for subsequent processing, and a ReLU activation function is used to introduce nonlinear transformations to enhance the network's expressiveness. Next, a 2D convolutional layer is applied, using a 3×3 convolution kernel to extract spatial features. The convolution operation effectively captures the spatial relationship in the image through local perception and weight sharing, and further increases nonlinearity through the ReLU activation function, improving the network's ability to handle complex features. The convolution layer keeps the spatial dimension unchanged to ensure the integrity of the spatial information. Finally, a linear dimensionality increase layer is used to restore the data, and the number of channels is flexibly adjusted to ensure that the spatial information is effectively retained and the network has better adaptability in different tasks. The whole process optimizes the performance of the model and enhances the sensitivity to local detail features through the combination of convolution and nonlinear transformation.

[0067] The structure of the frequency domain fine-tuning module is as follows: Figure 4 As shown in the figure, this module replaces the traditional multi-head self-attention mechanism by using a frequency domain feature extraction layer. The frequency domain feature extraction layer extracts frequency domain information from the image, which enhances the model's sensitivity to local details while improving its global feature modeling ability. In the specific implementation, except for the frequency domain feature extraction layer, other layers remain frozen, and only a small number of weight parameters are open for learning, thereby minimizing parameter updates. The frequency domain feature extraction layer, according to the steps of input data processing, is as follows: First, the input image feature data is converted from the spatial domain to the frequency domain through a two-dimensional real fast Fourier transform to obtain the spectral features. Then, the spectral features are gated, and the learned complex weights are used to enhance the feature representation capability. After spectral gating, the inverse fast Fourier transform is used to convert the data back to the spatial domain and restore it to the shape of the original feature map. This structure effectively combines frequency domain feature extraction and processing, and improves the model's ability to capture complex data patterns.

[0068] The specific steps are as follows: First, the time domain features Perform a two-dimensional Fourier fast transform on the spatial dimension to obtain the representation of the feature in the frequency domain, as shown in formula (1):

[0069]

[0070] Then, a learnable weight parameter is introduced to weight the feature representation in the frequency domain to obtain the weighted frequency domain feature. Finally, the weighted feature representation is converted back to the time domain using the inverse fast Fourier transform, which can be expressed as follows:

[0071]

[0072] Where w(u,v) represents the learnable frequency domain weight, ⊙ represents element-by-element multiplication, x(u,v) represents the feature converted from the time domain to the frequency domain through FFT, and u and v represent amplitude and phase respectively.

[0073] The hint encoder converts the input bounding box weak annotations as weak supervision annotations into feature information, thereby guiding the model to focus on the region of interest. By fine-tuning the hint encoder and mask decoder, the process of feature extraction and segmentation image reconstruction is further optimized. The mask decoder reconstructs the feature map fused from the encoder and hint encoder into the final segmentation result through upsampling and convolution operations, thereby achieving efficient and accurate segmentation.

[0074] After completing the construction of the segmentation model, the pre-trained weights of SAM are loaded to utilize its general feature extraction capabilities and initial parameters obtained on natural images for transfer learning on whole heart medical images. The SAM pre-training process is an existing technology. In the specific fine-tuning training process, all parameters in the encoder part except the proposed convolutional spatial local perception adapter and the frequency domain feature extraction layer that replaces the multi-head self-attention are frozen. Only these two specific layers are trained, and the prompt encoder and decoder are fine-tuned at the same time.

[0075] Step 3: Combine pixel-level annotation data and bounding box weak annotation data to perform weakly supervised fine-tuning training on the segmentation model. The weakly supervised training framework is as follows: Figure 5 As shown, the average teacher architecture is adopted.

[0076] First, the preprocessed image data is divided into two parts: one part uses complete pixel-level annotations, and the other part is used as unlabeled data. For all input images, whether they have pixel-level annotations or not, their corresponding bounding box weak annotations are fed into the hint encoder as additional hint information to guide the model to make more accurate predictions of the target area. The pixel-level data is represented as The unlabeled data is represented as Where X L(i) , represents a CT image or MR image with height H, width W and channel C, Y L(i) ∈{0,1} H×W×C Represents X L(i) The goal of the whole heart image segmentation task can be expressed as formula (3):

[0077]

[0078] Where θ, θ′ represent the weights of the student model and the teacher model respectively, P represents the weak label of the bounding box used to indicate the location information, and L S is the supervised loss, L C is the consistency loss. λ is a time-dependent Gaussian function used to dynamically adjust the weight of the consistency loss, which can be expressed as formula (4):

[0079]

[0080] Where t represents the current iteration number, t max Represents the maximum number of iterations. The weight θ of the teacher model t ′ through the weight θ of the student model t Its exponential moving average is updated and can be expressed as formula (5):

[0081] θ t ′=βθ t ' -1 +(1-β)θ t (5)

[0082] Where β is the decay coefficient of EMA, which is used to control the retention of the teacher model weights from the previous step, and β is set to 0.99. t ′ is the weight of the teacher model at the current step t, θ t is the weight of the student model at the current step t.

[0083] Supervised loss L S The calculation method is: pixel-level annotation data and bounding box weak label data are input into the student model or teacher model to obtain the prediction result f of pixel-level annotation data θ (X L(i) ,P) and f θ′ (X L(i),P), by comparing the deviation between the prediction results and the pixel-level annotations, the designed supervised loss function is optimized to improve the segmentation accuracy of the model. The supervised loss function introduces Focal Loss, which dynamically adjusts the loss weight of easy-to-classify samples, so that the model can pay more attention to those categories that are difficult to classify or have a small number of samples, effectively alleviating the problem of category imbalance and improving the model's learning ability for small targets and difficult-to-segment areas. Combined with Dice Loss and Iou Loss, the segmentation accuracy is further optimized, especially in the accuracy of the boundary area. This combined loss function can significantly improve the performance of the model in multi-category medical image segmentation tasks, especially when dealing with unbalanced data sets. As shown in formula (6), it can solve the problem of category imbalance and achieve better segmentation performance.

[0084]

[0085] Among them, FL is FocalLoss loss, w FL Represents the weight, which is used to adjust the importance of FocalLoss loss; L Iou represents IOU loss, scale Iou Indicates L Iou The scaling factor of Dice represents Dice loss; p t Represents the probability of the model correctly predicting each category; α t represents the weight factor of each sample in the loss function, which is used to adjust the importance of different categories; γ is the focusing parameter, which is used to reduce the contribution of easy-to-classify samples to the total loss; f θ (X L(i) ,P) represents the prediction result of the student model for labeled data; μ represents a constant to prevent the denominator from being 0; Y L(i) ∈{0,1} H×W×C Represents X L(i) Tags, X L(i) It is pixel-level annotation data, and P represents the weak label of the bounding box used to provide location information.

[0086] Uncertainty-guided consistency loss L C The calculation method is as follows: First, add different noise perturbations to the unlabeled data X′ U The bounding box weak label data P is input into the student model to simulate the diversity and uncertainty of the input data. The dropout layer is turned on to introduce randomness so that the result of each forward propagation is different. Then, T forward propagations are performed to obtain a set of unlabeled data X′ after adding perturbations. U The prediction result f θ (X′ U,P). By analyzing T different prediction results, the Monte Carlo estimation method can be used to approximate the uncertainty of the model's output for the input data. The core idea of ​​Monte Carlo estimation is to approximate the characteristics of the prediction distribution through multiple random sampling. Finally, for each forward propagation prediction result, the prediction entropy is calculated to quantify the segmentation uncertainty u of the model, as shown in formula (7):

[0087]

[0088] Uncertainty-guided consistency loss L C It can be expressed as formula (8):

[0089]

[0090] Among them, I is the sigmoid function, X U(i) represents unlabeled data, θ, θ′ represent the weights of the student model and the teacher model respectively, P represents the bounding box weak label used to prompt the location information, and f θ (X U(i) ,P) and f θ′ (X U(i) ,P) are the prediction results of the student model and the teacher model for the i-th sample; u i is the uncertainty of the prediction results of the student model, obtained by Monte Carlo estimation; H is the uncertainty threshold.

[0091] Step 4: Use the trained student model for the whole heart medical image segmentation task.

[0092] The experimental part of the present invention is based on the MMWHS public dataset, which aims to verify the performance of the innovative method proposed in the present invention in the whole heart medical image segmentation task. The MMWHS dataset contains whole heart images (with 7 parts) from different modalities (CT and MR), with high-quality pixel-level annotations, which are suitable for evaluating the segmentation accuracy and robustness of the model. Through experimental verification on the MMWHS dataset, combined with weakly supervised learning, convolutional adapters and frequency domain feature extraction layers, the segmentation accuracy of the model is improved, which has strong practical application value. The experimental results are shown in Tables 1-3.

[0093] Table 1 Comparison of segmentation performance of the proposed method and the fully supervised method of SAM variant under bounding box prompts on the MMWHS dataset

[0094]

[0095] Table 2 Comparison of segmentation performance of the method of the present invention and other MT series methods on the MMWHS dataset under the setting of 10% labeled data

[0096]

[0097] Table 3 Comparison of segmentation performance of the method of the present invention and other MT series methods on the MMWHS dataset under the setting of 20% labeled data

[0098]

[0099] In the table, ours-SCT is the result of single-category training, ours-NP-SCT is the result of single-category training without weak label prompts. Each category is trained separately. When fine-tuning a single category under single category experiments, no prompt box weak label is used as a position prompt. Ours is the result of multi-category training. When training multiple categories together, the prompt box weak label is used. The MT series is the test result of single-category training.

[0100] Any matters not described in the present invention are applicable to the prior art.

[0101] [1]Kirillov,A.,Mintun,E.,Ravi,N.,Mao,H.,Rolland,C.,Gustafson,L.,Xiao,T.,Whitehead,S.,Berg,AC,Lo,W.-Y.et al.(2023).Segment anything.InProceedings of the IEEE / CVF International Conference on Computer Vision(pp.4015–4026).

[0102] [2]Cheng J,Ye J,Deng Z,et al.Sam-med2d[J].arXiv preprint arXiv:2308.16184,2023.

[0103] [3]Ma,J.,He,Y.,Li,F.,Han,L.,You,C.,&Wang,B.(2024).Segment anything inmedical images.Nature Communications,15,654.URL:https: / / doi.org / 10.1038 / s41467-024-44824-z.doi:10.1038 / s41467-024-44824-z.

[0104] [4]Chai S,Jain R K,Teng S,et al.Ladder fine-tuning approach forsamintegrating complementary network[J].arXiv preprint arXiv:2306.12737,2023.

[0105] [5]Tarvainen A,Valpola H.Mean teachers are better role models:Weight-averaged consistency targets improve semi-supervised deep learning results[J].Advances in neural information processing systems,2017,30.

[0106] [6]Yu L,Wang S,Li X,et al.Uncertainty-aware self-ensembling model forsemi-supervised 3D left atrium segmentation[C].Medical image computing andcomputer assisted intervention–MICCAI 2019:22nd international conference,Shenzhen,China,October 13–17,2019,proceedings,part II 22.SpringerInternational Publishing,2019:605-613.

[0107] [7]Basak H,Bhattacharya R,Hussain R,et al.An exceedingly simpleconsistency regularization method for semi-supervised medical imagesegmentation[C] / / 2022IEEE 19th International Symposium on Biomedical Imaging(ISBI).IEEE,2022:1-4.

[0108] [8]Dosovitskiy A,Beyer L,Kolesnikov A,et al.An image is worth 16x16words:Transformers for image recognition at scale[J].arXiv preprint arXiv:2010.11929,2020.

Claims

1. A whole heart segmentation method based on uncertainty-guided weakly supervised parameter efficient fine-tuning, characterized in that: The segmentation method comprises the following steps: Step 1: Obtain several 3D whole heart structure CT images and MR images, cut the 3D images into 2D slices, crop and fill them to a uniform size, and preprocess the images to form a data set. Each 2D slice has a pixel-level label. Extract the pixel-level labels in the 2D slices, lock the region of interest, set the perturbation range, and randomly simulate and generate bounding boxes for the 2D slices within the perturbation range, which are used as weak annotations of the bounding boxes of the 2D slices. Step 2: Build a segmentation model for the whole heart medical image segmentation task; The segmentation model is based on the SAM model and includes an encoder, a mask decoder and a prompt encoder. The encoder is composed of α convolution fine-tuning modules and G-α frequency domain fine-tuning modules connected in series. The convolution fine-tuning module includes layer normalization, multi-head attention, layer normalization and multi-layer perceptron connected in sequence, and the input of the first layer normalization is added to the output of the multi-head attention to obtain feature A, and feature A is processed by the second layer normalization and multi-layer perceptron to output feature B, and feature A is processed by the convolutional spatial local perception adapter to obtain feature C, and features A, B, and C are added and input into the next level module; the convolutional spatial local perception adapter includes a linear layer, a ReLu activation function, a convolutional layer, a ReLu activation function and a linear layer connected in sequence; The frequency domain fine-tuning module includes layer normalization, frequency domain feature extraction layer, layer normalization, and multi-layer perceptron connected in sequence, and the input of the frequency domain fine-tuning module is added to the output of the multi-layer perceptron of the module as the output of the frequency domain fine-tuning module; the frequency domain feature extraction layer includes Fourier transform and inverse Fourier transform, and the result of Fourier transform is multiplied by the learnable weight and then subjected to inverse Fourier transform to obtain the output of the frequency domain feature extraction layer; Step 3: Segmentation model training The SAM without convolutional spatial local perception adapter and frequency domain feature extraction layer is pre-trained using natural images, and the weight parameters of the pre-trained SAM are used as the initial parameters of the segmentation model. The segmentation model is then fine-tuned using the dataset to achieve migration from the field of natural image segmentation to the field of medical image segmentation to adapt to the data distribution and characteristics of the whole heart segmentation task; during the fine-tuning training process, all parameters in the encoder except the convolutional spatial local perception adapter and frequency domain feature extraction layer are frozen, and only these two specific layers are trained, while the prompt encoder and mask decoder are fine-tuned; The dataset is divided into two parts: one part uses complete pixel-level annotations, and the other part is used as unlabeled data. For all input images, whether they have pixel-level annotations or not, their corresponding bounding box weak annotations are fed into the hint encoder as additional hint information to guide the model to make more accurate predictions of the target area. For image data with pixel-level annotations, the student model calculates supervised loss by comparing the predicted results with the actual pixel-level annotations to directly optimize the segmentation accuracy. For unlabeled image data, two variants are generated by adding random perturbations to them and input into the teacher model and student model respectively. The output results of the teacher model and the student model are used to calculate the consistency loss guided by uncertainty. During fine-tuning training, the supervised loss L S The uncertainty-guided consistency loss is used as a joint loss to optimize the entire network; both the student model and the teacher model adopt a segmentation model structure; Step 4: Use the trained student model as the final segmentation model for the whole heart medical image segmentation task.

2. The method according to claim 1, characterized in that Uncertainty-guided consistency loss L C It is expressed as formula (8): Among them, I is the sigmoid function, X U(i) represents unlabeled data, θ, θ′ represent the weights of the student model and the teacher model respectively, P represents the bounding box weak label used to prompt the location information, and f θ (X U(i) ,P) and f θ′ (X U(i) ,P) are the prediction results of the student model and the teacher model for the i-th sample; u i is the uncertainty of the prediction results of the student model, obtained by Monte Carlo estimation; H is the uncertainty threshold.

3. The method according to claim 1, characterized in that The formula for the supervised loss is: Among them, FL is FocalLoss loss, w FL Represents the weight, which is used to adjust the importance of FocalLoss loss; L Iou represents IOU loss, scale Iou Indicates L Iou Scaling factor of L Dice represents Dice loss; p t Represents the probability of the model correctly predicting each category; α t represents the weight factor of each sample in the loss function, which is used to adjust the importance of different categories; γ is the focusing parameter, which is used to reduce the contribution of easy-to-classify samples to the total loss; f θ (X L(i) ,P) represents the prediction result of the student model for labeled data; μ represents a constant to prevent the denominator from being 0; Y L(i) ∈{0,1} H×W×C Represents X L(i) Tags, X L(i) It is pixel-level annotation data, and P represents the weak label of the bounding box used to provide location information.

4. The method according to claim 1, characterized in that: The whole heart structure refers to the main anatomical structure of the heart, including the myocardium, left atrium, left ventricle, right atrium, right ventricle, ascending aorta and pulmonary artery.

5. The method according to claim 1, characterized in that The convolutional layer in the convolutional spatial local perception adapter has 768 channels and the convolution kernel size is 3×3.

6. The method according to claim 1, characterized in that The frequency domain feature extraction layer processes the input data in the following steps: first, the input image feature data is converted from the spatial domain to the frequency domain through a two-dimensional real fast Fourier transform to obtain the spectrum feature; then, the spectrum feature is gated and the feature representation capability is enhanced using the learned complex weights; after spectrum gating, the data is converted back to the spatial domain using an inverse fast Fourier transform to restore it to the shape of the original feature map.

7. The method according to claim 1, characterized in that The value of G is 10-15, and the value of α is 4-10.

8. A computer-readable storage medium having a computer program stored thereon, characterized in that: When the program is executed by a processor, the steps of any method described in claims 1-7 can be implemented.

Images