A medical image multi-target domain field self-adaption method based on prompt learning

Abstract

The application discloses a medical image multi-target domain field self-adaptive method based on prompt learning, and comprises the following steps: step 1, a segmentation main module uses a U-net segmentation model to generate a segmentation result from an input medical image; step 2, a prompt generation module uses intermediate features of a domain classifier to generate a prompt corresponding to a domain; and step 3, a feature fusion module fuses the segmentation result and the prompt features, and an adversarial learning module uses an adversarial learning principle to further improve the model performance by performing feature matching on the fused features. The application generates a prompt by using a known input domain category to guide the model to perform specific segmentation processing on different domains, so that a deep learning model with good segmentation performance is realized.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of medical image analysis, and in particular to a medical image multi-target domain field self-adaption method based on prompt learning. BACKGROUND

[0002] Deep learning has brought medical image segmentation into the era of data-driven and made significant progress in this field, i.e., the segmentation accuracy has been significantly improved. Despite the great success, the deployment of the trained segmentation model is usually severely affected by the distribution shift between the training data (or labeled data) and the test data (or unlabeled data), because in this case the segmentation performance will be severely degraded. Domain transfer is usually caused by various factors, including the difference in acquisition protocol (such as parameters, imaging methods, modalities) and the characteristics of the data (such as age, gender, disease severity, etc.).

[0003] In order to deal with the distribution shift problem in medical image segmentation, domain adaptation (DA) methods have been proposed and studied. Many researchers propose to use adversarial learning to solve the distribution shift problem. These methods mainly use the game between the domain classifier and the feature extractor to learn the domain-invariant features. However, they are often prone to balance problems between the feature alignment and the discriminative ability of the model.

[0004] Some recent researchers have begun to explore domain adaptation algorithms based on self-training, which generate pseudo-labels for samples in the "other" domain to complete self-training. However, it is difficult to guarantee the quality of the pseudo-labels of the "other" domain, and it is also difficult to build an efficient model with noisy labels. However, most of these methods cannot handle very different situations between domains well, because it is very challenging to learn domain-invariant features when each domain contains knowledge specific to that domain. In addition, the domain information itself is also difficult to be well utilized in the domain adaptation algorithm. SUMMARY

[0005] The technical problem to be solved by the present application is to provide a medical image multi-target domain field self-adaption method based on prompt learning, which generates prompts using known input domain categories to guide the model to perform specific segmentation processing on different domains, and realizes a deep learning model with good segmentation performance.

[0006] To solve the above technical problems, the present application provides a medical image multi-target domain field self-adaption method based on prompt learning, comprising the following steps:

[0007] Step 1, the segmentation backbone module uses a U-net segmentation model to generate a segmentation result from an input medical image;

[0008] Step 2, the prompt generation module generates a prompt for the corresponding domain using the intermediate features of the domain classifier.

[0009] Step 3, the feature fusion module fuses the segmentation result with the prompt feature, and the adversarial learning module uses the adversarial learning principle to further improve the model performance by performing feature matching on the fused features.

[0010] Preferably, in step 1, the segmentation backbone module uses 3D U-Net as the backbone. 3D U-Net is a popular segmentation model in the field of medical image segmentation. It provides more useful image information to the model by connecting the encoder and decoder, thereby achieving better segmentation results. 3D U-Net is a three-dimensional convolutional neural network based on U-Net, widely used in the field of medical image segmentation. Compared with 2D U-Net, 3D U-Net can segment three-dimensional medical images, with better segmentation accuracy and more accurate shape restoration ability.

[0011] Preferably, in step 2, the prompt generation module includes a classifier and a prompt generator; the classifier is used to classify the input encoder features into their respective domains, and the prompt generator is used to extract and process the intermediate features of the classifier into specific domain prompts.

[0012] Preferably, the classifier h(x) is used to distinguish the domain of the input image, represented as As shown in formula 1:

[0013]

[0014] where x represents the intermediate feature of the model;

[0015] The cross-entropy loss is used to train the classifier, as shown in formula 2.

[0016]

[0017] where represents the model's prediction of the class, and d represents the true class of the input.

[0018] Preferably, in the prompt generator, the second-to-last layer feature z of the classifier is input into the prompt generator g(z), and g(z) is a multi-layer perceptron, as defined in formula 3.

[0019]

[0020] where represents a network layer composed of convolution + normalization + ReLU activation.

[0021] Preferably, in step 3, the feature fusion module is used to fuse the domain-specific prompt obtained by the prompt generation module and the encoder features to construct a feature fusion module A-Fusion, which adapts the image features to a representation form with domain awareness.

[0022] Preferably, the input image is represented as I, the result of the prompt generation module is g(e(I)), where e(I) is the feature from the shallow layer of the encoder, and E(I) is the encoder feature of the image, and then the feature fusion module is used to extract the feature with domain awareness, i.e., F, as shown in equation 4.

[0023] F = ψ(g(e(I)), E(I)) (4)

[0024] where ψ represents the feature fusion module A-Fusion.

[0025] Preferably, two input sources are obtained using the channel and spatial branches, the spatial branch compresses the channel information of the input to obtain spatial weights, and the channel branch compresses the prompt to obtain channel weights; a fully connected layer is used for feature compression and scaling, and f cp and f re are used to represent the weight calculation formulas of the two branches as shown in equation 5:

[0026] W s = f cp (E(I))

[0027] W c = f re (f cp (avgpool(g(e(I))))) (5)

[0028] The weights from the spatial and channel dimensions are combined to obtain a token, learning high-level and low-level features from the encoder features and the prompt to guide the fusion of the two features, and the process is as shown in equation 6:

[0029] W = sigmoid(W c +W s )

[0030] F out = g(e(I))*W + E(I)*(1-W) (6)

[0031] This module only introduces a few training parameters, but can effectively improve the quality of the features after fusion.

[0032] Preferably, the adversarial learning module uses the classic GAN loss to train the discriminator and the generator, and the loss function is as shown in equation 7:

[0033]

[0034]

[0035] where G(z i ) represents other domain data combined with domain-specific cues, and the generator is trained to make the discriminator identify other domain data as source domain data; and x i and represent source domain data and target domain data combined with domain-specific cues when training the discriminator, and the discriminator is trained to separate source domain data and other domain data, thereby achieving adversarial training.

[0036] Preferably, the final loss function is shown in formula 8:

[0037] L total =L seg +λ cls L cls +λ adv L adv (8)

[0038] where λ is a hyperparameter for adjusting the weight of the loss function.

[0039] The beneficial effects of the present application are: (1) the present application learns the cues of a specific domain from the data of multiple domains, and fuses the cues with the encoder features, so that the model can generate features with domain perception ability, thereby improving the domain adaptation ability of the model in multiple domains; (2) it can be compatible with traditional domain adaptation algorithms (such as adversarial learning methods), further improving the performance of the model. BRIEF DESCRIPTION OF DRAWINGS

[0040] Figure 1 is a flowchart of the method of the present application.

[0041] Figure 2 is a visualization schematic diagram of the segmentation graph (details) of all comparison methods of the present application. DETAILED DESCRIPTION

[0042] As shown in formula 1, a medical image multi-target domain domain adaptation method based on cue learning includes the following steps: Figure 1

[0043] Step 1: The segmentation backbone module uses a U-net segmentation model to generate a segmentation result from an input medical image;

[0044] Step 2: The cue generation module uses the intermediate features of the domain classifier to generate the cues of the corresponding domain;

[0045] ​Step 3, the feature fusion module fuses the segmentation results with the prompt features, and the adversarial learning module uses the adversarial learning principle to further enhance the model performance by performing feature matching on the fused features.

[0046] The 3D U-Net is used as the backbone in the segmentation backbone module. The 3D U-Net is a three-dimensional convolutional neural network based on the U-Net, which is widely used in medical image segmentation. Compared with the 2D U-Net, the 3D U-Net can segment three-dimensional medical images, with better segmentation accuracy and more accurate shape restoration ability.

[0047] The prompt generation module mainly consists of two parts: a classifier and a prompt generator. Domain-specific information is incorporated into the prompt to guide the model to adapt to the target domain. To achieve this, a classifier h(x) is introduced to distinguish the domain of the input image (denoted as ), as shown in equation 1:

[0048]

[0049] where x represents the intermediate feature of the model.

[0050] In addition, the cross-entropy loss is used to train the classifier, as shown in equation 2.

[0051]

[0052] where represents the model's prediction of the class, and d represents the true class of the input.

[0053] In the prompt generator part, instead of directly using as the class information, the second-to-last layer feature z of the classifier is input into the prompt generator g(z). Specifically, g(z) is a multi-layer perceptron, as defined in equation 3.

[0054]

[0055] where represents a network layer composed of convolution + normalization + ReLU activation. It is worth noting that this module does not change the size of the feature map z, but converts the extracted class features into a prompt specific to the domain.

[0056] The feature fusion module is used to fuse the domain-specific prompt obtained by the prompt generation module with the encoder features. Since the prompt clearly captures the information of a certain specific domain, while the encoder extracts features that describe the semantic and spatial information of the image, considering their characteristics, the feature fusion module A-Fusion is constructed to combine them, adapting the image features to a domain-aware representation.

[0057] Let the input image be represented as I, the hint generation module result as g(e(I)) (where e(I) is the feature from the shallow layer of the encoder), and E(I) as the encoder feature of the image. Then the feature with domain awareness is extracted by the feature fusion module (i.e., F), as shown in Equation 4.

[0058] F = ψ(g(e(I)), E(I)) (4)

[0059] where ψ represents the feature fusion module A-Fusion. Since the learned hint and the input data feature captured by the encoder represent different aspects of the data, using only addition, multiplication, or convolution as the fusion function ψ does not achieve good results. A-Fusion calculates the channel weight and the spatial weight, respectively, to enhance the input channel and spatial characteristics.

[0060] Two input sources are obtained using the channel and spatial branches. The spatial branch compresses the channel information of the input to obtain the spatial weight; at the same time, the channel branch compresses the hint to obtain the channel weight. A Fully Connected Layer is used for feature compression and scaling, denoted by f cp and f re . The weight calculation formula of the two branches is shown in Equation 5.

[0061] W s = f cp (E(I))

[0062] W c = f re (f cp (avgpool(g(e(I))))) (5)

[0063] After that, the weights from the spatial and channel dimensions are combined to obtain a token, which can learn high-level and low-level features from the encoder feature and the hint to guide the fusion of the two features. The process is shown in Equation 6:

[0064] W = sigmoid(W c +W s )

[0065] F out = g(e(I))*W + E(I)*(1-W) (6)

[0066] This module only introduces a few training parameters, but can effectively improve the quality of the feature after fusion.

[0067] The adversarial learning module is an optional module of the method. The prompt learning method is fully compatible with other domain adaptation algorithms, so the adversarial learning is used as an optional module in our network to enhance the domain adaptation ability of the model. The module uses the classic GAN loss to train the discriminator and the generator (where the adversarial loss for the generator only acts on the prompt generator).

[0068] Based on the above-mentioned various modules, the final loss function of the method is shown in formula 8:

[0069] L total =L seg +λ cls L cls +λ adv L adv (8)

[0070] Wherein λ is a hyperparameter for adjusting the weight of the loss function.

[0071] The present application has done numerous experiments on Iseg infant brain segmentation dataset and BraTs brain tumor segmentation dataset. In the brain segmentation task, the data of 6 months old is taken as the labeled source domain data, and 12 months and 3 months are taken as the unlabeled target domain data; in the brain tumor segmentation task, the high-grade glioma is taken as the labeled source domain data, and the low-grade glioma is taken as the unlabeled target domain data. The domain adaptation effect of the method is evaluated by comparing the performance on the target domain. First, the performance comparison with the existing advanced method is shown in Tables 1, 2 and Figure 2 As can be seen from the experimental results, compared with the existing method, the domain adaptation effect of the method of the present application is better, and the effect is more obvious when there are multiple target domains involved.

[0072] Table 1 Comparison table of the method of the present application and the existing advanced method in the infant brain segmentation task

[0073]

[0074] Table 2 Comparison table of the method of the present application and the existing advanced method in the brain tumor segmentation task

[0075]

[0076] At the same time, the effectiveness of the method of the present application is ablation experiment. The experiment is carried out on 3D 3D-Unet and 2D TransUnet respectively. The experimental results are shown in Tables 3 and 4. In Table 4, rAFusion represents that the encoder feature is placed in the channel branch of the AFusion module, and the prompt is placed in the spatial branch of the AFusion module; and AFusion is the opposite. From the experimental results, we can see that each module of the method of the present application can effectively improve the model performance.

[0077] Table 3 Method module effectiveness verification table

[0078]

[0079] Table 4 Feature fusion module superiority verification table

[0080]

Claims

1. A multi-objective domain adaptive method for medical images based on cue learning, characterized in that, Includes the following steps: Step 1: The segmentation backbone module uses the U-net segmentation model to generate segmentation results from the input medical image; Step 2: The prompt generation module uses the intermediate features of the domain classifier to generate prompts for the corresponding domain. The prompt generation module includes a classifier and a prompt generator. The classifier is used to classify the input encoder features into their respective domains, and the prompt generator is used to extract and process the intermediate features of the classifier into prompts for a specific domain. A classifier h(x) is used to distinguish the domains of the input image, denoted as... As shown in formula (1): Where x represents the intermediate features of the model; The classifier is trained using cross-entropy loss, as shown in Equation (2): in d represents the class predicted by the model, and d represents the true class of the input. In the prompt generator, the penultimate layer feature z of the classifier is input into the prompt generator g(z), which is a multilayer perceptron, as shown in formula (3): in This represents a network layer consisting of convolution, normalization, and ReLU activation. Step 3: The feature fusion module fuses the segmentation results with the prompt features. The adversarial learning module uses the principle of adversarial learning to further improve the model performance by using the fused features as feature pairs. The adversarial learning module uses the classic GAN loss to train the discriminator and generator. The loss function is shown in formula (7). Where G(z) i ) represents data from other domains that incorporate domain-specific cues; during generator training, the goal is to train the discriminator to recognize this other domain data as source domain data; while x i and This represents the source domain data and the target domain data combined with domain-specific cues during discriminator training. The goal of training the discriminator is to separate the source domain data from the other domain data, thereby achieving adversarial training.

2. The medical image multi-objective domain adaptive method based on cue learning as described in claim 1, characterized in that, In step 1, the segmentation backbone module uses 3D U-Net as the backbone, and achieves better segmentation results by making long connections between the encoder and decoder.

3. The multi-objective domain adaptation method for medical images based on cue learning as described in claim 1, characterized in that, In step 3, the feature fusion module is used to fuse the domain-specific cues obtained by the cue generation module and the encoder features to construct the feature fusion module A-Fusion, which adapts the image features into a domain-aware representation.

4. The multi-objective domain adaptation method for medical images based on cue learning as described in claim 3, characterized in that, The input image is represented as I, and the result of its prompt generation module is g(e(I)), where e(I) is the feature from the shallow layer of the encoder, and E(I) is the encoder feature of the image. Then, the feature fusion module extracts the domain-aware feature F, as shown in formula (4): F=ψ(g(e(I)),E(I)) (4) Where ψ represents the feature fusion module A-Fusion.

5. The multi-objective domain adaptation method for medical images based on cue learning as described in claim 3, characterized in that, Two input sources are obtained using channel and spatial branches. The spatial branch compresses the channel information of the input to obtain spatial weights; simultaneously, the channel branch compresses the input cue to obtain channel weights. A fully connected layer is used for feature compression and scaling, using f... cp and f re The weights of the two branches are calculated using formula (5): W s =f cp (E(I)) W c =f re (f cp (avgpool(g(e(I))))) (5) The weights from the spatial and channel dimensions are combined to obtain a token. High-level and low-level features from encoder features and cues are learned to guide the fusion of the two features, as shown in Equation (6): W=sigmoid(W c +W s ) F out =g(e(I))*W+E(I)*(1-W) (6) This module introduces very few training parameters, but it can effectively improve the quality of features after feature fusion.

6. The multi-objective domain adaptation method for medical images based on cue learning as described in claim 1, characterized in that, The final loss function is shown in Equation (8): L total =L seg +λ cls L cls +λ adv L adv (8) Where λ is a hyperparameter used to adjust the weights of the loss function.

Images