Unsupervised domain adaptation medical image segmentation method based on teacher fusion model

Abstract

The present invention discloses an unsupervised domain adaptation medical image segmentation method based on a teacher fusion model, and relates to the technical field of medical image processing. The unsupervised domain adaptation medical image segmentation method based on a teacher fusion model provided by the present invention includes a frequency domain alignment module based on 3D wavelet transform, which performs appearance alignment by exchanging low frequency spectra to obtain a class target domain image one; a feature alignment module based on 2.5D spatial direction continuous smoothness, which performs feature domain alignment by constructing 2.5D image blocks to obtain a class target domain image two; the class target domain image one and the class target domain image two are respectively input into the first teacher model and the second teacher model for supervised training, and the student model training is fused and guided; the weight parameters are subjected to singular value decomposition, and training is performed by updating the singular value parameters. Therefore, the above method can effectively obtain the three-dimensional feature information of the image, while reducing the amount of calculation and improving the accuracy of image segmentation.

Description

Technical Field

[0001] The present invention relates to the technical field of medical image processing, and in particular to an unsupervised domain adaptation medical image segmentation method based on a teacher fusion model. Background Art

[0002] In recent years, unsupervised domain adaptation (UDA) methods based on deep learning (DL) have made progress in medical image segmentation tasks. This algorithm abandons the iterative idea of ​​traditional algorithms and learns the parameter representation of a function at a higher level. It only needs to use the learned parameters to calculate the target domain image whose labels are invisible, and can directly output the segmentation results of the target domain based on the knowledge learned from the source domain. At present, most medical image domain adaptation methods based on UDA mainly focus on two aspects: one is image alignment, which uses unpaired image-to-image conversion to convert the image style from the source domain to the target domain; the other is feature alignment, which extracts domain-invariant features through adversarial training, combines the domain-varying features of the source domain with the domain-invariant features of the target domain, and obtains images with the style of the target domain and the content of the source domain.

[0003] With the continuous in-depth research on U-net and Transformer, the existing UDA-based medical image segmentation methods have made great progress accordingly, but there are still some shortcomings in the method: (1) Many current UDA methods use slices to process 2D data, ignoring the inherent three-dimensional information in medical images. (2) Pure 3D UDA methods have defects such as high complexity and excessive use of computing resources. (3) There is a lack of combining frequency domain features with spatial domain features to adapt to different domains. (4) There is a lack of effective learning and judgment methods for images that are difficult to accurately segment. Therefore, there is an urgent need for an unsupervised domain adaptation medical image segmentation method that can obtain features from different domains, thereby improving the adaptation performance to the target domain. Summary of the invention

[0004] The purpose of the present invention is to provide an unsupervised domain adaptation medical image segmentation method based on a teacher fusion model, which can effectively obtain the three-dimensional feature information of the image, reduce the amount of calculation and improve the image segmentation accuracy.

[0005] To achieve the above object, the present invention provides an unsupervised domain adaptation medical image segmentation method based on a teacher fusion model, comprising the following steps:

[0006] S1, frequency domain alignment module based on 3D wavelet transform, aligns the source domain image and the target domain image by exchanging the low frequency spectrum to obtain the target domain image 1;

[0007] S2, a feature alignment module based on continuous smoothing in 2.5D spatial direction, which constructs 2.5D image blocks to perform feature domain alignment and obtains the second target domain image;

[0008] S3. Input the target domain image 1 and the target domain image 2 into the first teacher model and the second teacher model respectively, and perform supervised training according to the source domain label, fuse the trained outputs, and guide the student model training.

[0009] Preferably, step S2 includes:

[0010] S21, obtain continuous slices of the source domain image, extract features through the encoder, cut the extracted feature map into blocks sequentially, and perform feature enhancement through the orthogonal attention module;

[0011] S22. Decode the enhanced features through a decoder, and iteratively update the same target slice using a weighted moving average strategy to obtain a 2.5D image block after feature alignment, merge the target slices in the obtained 2.5D image block in the spatial direction, and obtain the target domain image 2.

[0012] Preferably, the feature alignment module based on continuous smoothness in 2.5D spatial direction uses adversarial loss and reconstruction loss to train the network, as follows:

[0013] ;

[0014] ;

[0015] In the formula, , are adversarial loss and reconstruction loss respectively, , The source domain image block and the target domain image block are respectively, Generator for feature alignment module with continuous and smooth feature in 2.5D space direction, The discriminator of the feature alignment module is a continuous and smooth feature alignment module in 2.5D space direction. , They are and The expected sample distribution of .

[0016] Preferably, in step S3, the first teacher model and the second teacher model are both trained using Dice loss and cross entropy loss, as follows:

[0017] ;

[0018] ;

[0019] In the formula, , are the losses of the first teacher model and the second teacher model, respectively. , are the weights of Dice loss and cross entropy loss, respectively. , are the predicted values ​​of the first teacher model and the second teacher model, respectively. is the source domain label.

[0020] Preferably, the student model is trained using an entropy-based multi-teacher model segmentation loss function, as follows:

[0021] ;

[0022] ;

[0023] ;

[0024] In the formula, is the segmentation loss function for the multi-teacher model, is the loss function calculated based on the three-dimensional KL divergence, For the The weights of the teacher models, , , are the height, width and depth of the 3D image, respectively. is the temperature coefficient, is the predicted value of the student model, represents the Shannon entropy, express Divergence, is the softmax function, is a hyperparameter.

[0025] Preferably, in step S3, the student model training includes mapping the convolution layer weight parameters to singular value parameters using singular value decomposition when performing knowledge distillation based on the trained first teacher model and the second teacher model, and training by updating the singular value parameters, as follows:

[0026] ;

[0027] In the formula, , They are the pre-trained weight layer and the weight layer updated after mapping, , are the left and right singular vectors, respectively, are the singular value parameters mapped after pre-training of the first teacher model and the second teacher model, No. The singular value parameters updated at the iteration.

[0028] Therefore, the present invention adopts the above-mentioned unsupervised domain adaptation medical image segmentation method based on the teacher fusion model, which has the following technical effects:

[0029] (1) Adaptive training is performed in both the frequency domain and feature domain to the target domain through the 3D wavelet transform frequency domain alignment module and the 2.5D spatial direction continuous smoothing feature alignment module. This can preserve the inherent three-dimensional information in medical images and complete the conversion of image style and features from the source domain to the target domain.

[0030] (2) The first teacher model and the second teacher model are trained respectively through the class-target domain images in the frequency domain and feature domain. The output pseudo-labels can amplify the missing data. At the same time, the two teacher models are fused to guide the student model, which can improve the segmentation accuracy.

[0031] (3) The singular value decomposition method is used to map the convolution weight layer parameters into singular value parameters, which can be updated. This not only makes the generated output match the distribution of the target domain, but also reduces the number of parameters required for learning.

[0032] The technical solution of the present invention is further described in detail below through the accompanying drawings and embodiments. BRIEF DESCRIPTION OF THE DRAWINGS

[0033] Figure 1 is a schematic diagram of an overall network framework in an embodiment of an unsupervised domain adaptation medical image segmentation method based on a teacher fusion model;

[0034] Figure 2 A frequency domain alignment module based on 3D wavelet transform in an embodiment of an unsupervised domain adaptation medical image segmentation method based on a teacher fusion model;

[0035] Figure 3 is a schematic diagram of a feature alignment module based on 2.5D spatial direction continuous smoothing in an embodiment of an unsupervised domain adaptation medical image segmentation method based on a teacher fusion model;

[0036] Figure 4 2 is a schematic diagram of an entropy-based multi-teacher guided segmentation network in an embodiment of an unsupervised domain adaptation medical image segmentation method based on a teacher fusion model. DETAILED DESCRIPTION

[0037] The present invention can be explained in more detail by the following examples. The purpose of disclosing the present invention is to protect all changes and improvements within the scope of the present invention. The present invention is not limited to the following examples.

[0038] like Figure 1 As shown, the unsupervised domain adaptation medical image segmentation method based on the teacher fusion model provided by the present invention includes an image translation process and an image segmentation process:

[0039] In the image translation process, image translation methods for transform domain and feature domain are proposed respectively. By training the translation network, stable and realistic target domain images can be generated. Through the segmentation labels of the corresponding source domain images, two teacher models are used to perform fully supervised training on the images generated in the transform domain and feature domain respectively. At the same time, the two teacher models output preliminary pseudo labels of the learned knowledge.

[0040] During the image segmentation process, the preliminary pseudo-labels are adjusted according to the entropy-based dynamic multi-teacher co-learning strategy to obtain the final pseudo-labels; then, the lighter-weight student segmentation model is fully supervised trained, and the parameter learning is updated according to the loss function. The trained student segmentation model is the final cross-modal adaptive medical image segmentation model required.

[0041] Specifically, start from the following steps:

[0042] S1. Please see Figure 2 , taking a 3D medical image from a certain type of scan as the source domain image (source image) and a 3D medical image from another type of scan as the target domain image (target image), both are input into the frequency domain alignment module based on 3D wavelet transform to obtain a target-like image whose content follows the source domain but has the style of the target domain.

[0043] The input 3D source domain image is decomposed into eight blocks in order of frequency through 3D wavelet transform, which are defined as “LLL, LLH, LHL, LHH, HLL, HLH, HHL and HHH” respectively, with “LLL” indicating that the block is low-frequency information in the X, Y and Z directions.

[0044] Considering that exchanging low-frequency spectra between images can exchange their styles / colors / brightness without changing the semantic content, while exchanging high-frequency spectra will introduce unnecessary artifacts. Therefore, this embodiment aligns the appearance only by exchanging low-frequency spectra, and adopts a fusion method of combining the "LLL" block of the target domain image with other blocks of the source domain image and performing inverse wavelet transform to obtain a target domain image in the frequency domain, that is, target domain image 1.

[0045] S2, see Figure 3 ,This embodiment uses a feature alignment module based on 2.5D spatial direction continuous smoothing to perform another branch of alignment operation on the source domain image.

[0046] Specifically, a 3D image data ( The encoder uses continuous slices (valued from 1 to 3) to construct a 2.5D image block as the input of the encoder, obtains the corresponding feature map and cuts it sequentially, and then feeds it into the Transformer-based horizontal-vertical feature attention extraction network module (called orthogonal attention module) to improve the learning performance of the features; then a new 2.5D image block is output by the decoder with the same design as the encoder, and a discriminator is designed to improve the domain adaptation performance of the extracted 2.5D image block using an anti-learning strategy.

[0047] according to The value of , for each slice in the determined direction, three iterative learnings are performed. By weighing the accuracy and the amount of calculation, this embodiment During the iteration process, the analysis Zhang Zhi Slices (in Figure 3 In For example), a total of Each iteration uses a weighted moving average strategy (EMA) with a decay rate (α) of 0.9 to update the network weights. Afterwards, in the inference phase, the middle slices of each new 2.5D image block are superimposed, that is, merged in the spatial direction to form a complete 3D image of the class target domain, and the class target domain image under the feature domain is obtained, that is, the class target domain image 2.

[0048] Among them, the adversarial loss and reconstruction loss are used to train the generative network so that the source domain adapts to the target domain, which is expressed as follows:

[0049] ;

[0050] ;

[0051] In the formula, , are adversarial loss and reconstruction loss respectively, , The source domain image block and the target domain image block are respectively, Generator for feature alignment module with continuous and smooth feature in 2.5D space direction, The discriminator of the feature alignment module is a continuous and smooth feature alignment module in 2.5D space direction. , They are and The expected sample distribution of .

[0052] S3. Large models can usually accommodate more knowledge, but they have a large number of parameters and consume more computing resources. In addition, when the amount of training data is small, large models are more likely to overfit. Therefore, this implementation adopts the method of distilling the knowledge of large models into small models to solve the above problems.

[0053] Specifically, in the segmentation stage, a multi-teacher model strategy is adopted to input the generated first-class target domain image 1 and class target domain image 2 into the first teacher model and the second teacher model respectively, and supervised learning is performed according to the labels of the source domain. After training, the fusion guides the learning of the student model, such as Figure 4 shown.

[0054] Here, Dice loss and cross entropy loss are used to train the teacher model, which is specifically expressed as follows:

[0055] ;

[0056] ;

[0057] In the formula, , are the losses of the first teacher model and the second teacher model, respectively. , are the weights of Dice loss and cross entropy loss, respectively. , are the predicted values ​​of the first teacher model (transform domain branch) and the second teacher model (feature domain branch), respectively. is the source domain label.

[0058] In traditional multi-teacher distillation (MKD), the average result of the last layer of all teacher models is used as information to guide student training. In this embodiment, the multi-teacher model segmentation loss function based on entropy is used to train the student model. By using entropy calculation to give weights to the results given by the teacher, it is possible to approximately represent how confident the teacher model is in predicting the output class, thereby obtaining a more natural guidance effect.

[0059] Among them, the segmentation loss function of the multi-teacher model based on entropy is as follows:

[0060] ;

[0061] ;

[0062] ;

[0063] In the formula, Split loss function for multi-teacher model; is the loss function calculated based on the three-dimensional KL divergence; For the The weights of the teacher models; , , are the height, width and depth of the three-dimensional image respectively; is the temperature coefficient; is the predicted value of the student model; represents Shannon entropy; express Divergence; is the softmax function; is a hyperparameter, which is set to 10 here.

[0064] Considering that the amount of computation required for the 3D multi-teacher segmentation neural network is large and that it may be difficult to achieve high-quality learning when there are fewer samples, this embodiment proposes to use feature extraction of multi-scale singular values ​​to solve the above problems.

[0065] Specifically, when performing knowledge distillation on the two teacher models that have been trained on the translated images, singular value decomposition (SVD) is applied to each layer of convolution weight parameters, and SVD is directly used for the fully connected layer. This embodiment uses a 5D convolution weight layer , it needs to be expanded to two dimensions. In order to preserve the characteristic relationship of the space, the expanded two-dimensional weight is ,in Then, the weight parameters trained on the translated image are decomposed using singular value decomposition, and the left / right singular vectors of the trained weight layer ( and ) is frozen, the weight parameter ( ) is transformed into the training of singular value parameters ( ), which is specifically expressed as follows:

[0066] ;

[0067] In the formula, , They are the pre-trained weight layer and the updated weight layer after mapping, respectively.

[0068] Therefore, it is possible to find a new set of singular values ​​in each layer of the trained teacher model, which not only makes the generated output match the distribution of the target domain, but also reduces the amount of parameters required for learning.

[0069] Therefore, the present invention adopts the above-mentioned unsupervised domain adaptation medical image segmentation method based on the teacher fusion model, which can effectively obtain the three-dimensional feature information of the image, while reducing the amount of calculation and improving the accuracy of image segmentation.

[0070] Finally, it should be noted that the above embodiments are only used to illustrate the technical solution of the present invention rather than to limit it. Although the present invention has been described in detail with reference to the preferred embodiments, those skilled in the art should understand that they can still modify or replace the technical solution of the present invention with equivalents, and these modifications or equivalent replacements cannot cause the modified technical solution to deviate from the spirit and scope of the technical solution of the present invention.

Claims

1. Unsupervised domain adaptation medical image segmentation method based on teacher fusion model, characterized in that: The following steps are involved: S1, frequency domain alignment module based on 3D wavelet transform, aligns the source domain image and the target domain image by exchanging the low frequency spectrum to obtain the target domain image 1; S2, based on the 2.5D spatial direction continuous smooth feature alignment module, by constructing 2.5D image blocks to align the feature domain, we get the target domain image 2: S21, obtain continuous slices of the source domain image, extract features through the encoder, cut the extracted feature map into blocks sequentially, and perform feature enhancement through the orthogonal attention module; S22, decoding the enhanced features through a decoder, and iteratively updating the same target slice using a weighted moving average strategy to obtain a 2.5D image block after feature alignment, merging the target slices in the obtained 2.5D image block in a spatial direction to obtain a second target domain image; S3. Input the target domain image 1 and the target domain image 2 into the first teacher model and the second teacher model respectively, and perform supervised training according to the source domain label, fuse the trained outputs, and guide the student model training.

2. The unsupervised domain adaptation medical image segmentation method based on the teacher fusion model according to claim 1, characterized in that: The feature alignment module based on continuous smoothness in 2.5D spatial direction uses adversarial loss and reconstruction loss to train the network, as follows: In the formula, , are adversarial loss and reconstruction loss respectively, , are the source domain image block and the target domain image block respectively, Generator for feature alignment module with continuous and smooth feature in 2.5D space direction, The discriminator of the feature alignment module is a continuous and smooth feature alignment module in 2.5D space direction. , They are and The expected sample distribution of .

3. The unsupervised domain adaptation medical image segmentation method based on the teacher fusion model according to claim 2 is characterized in that: In step S3, both the first teacher model and the second teacher model are trained using Dice loss and cross entropy loss, as follows: In the formula, , are the losses of the first teacher model and the second teacher model, respectively. , are the weights of cross entropy loss and Dice loss respectively, , are the predicted values ​​of the first teacher model and the second teacher model, respectively. is the source domain label.

4. The unsupervised domain adaptation medical image segmentation method based on the teacher fusion model according to claim 3 is characterized in that: The student model is trained using the entropy-based multi-teacher model segmentation loss function as follows: In the formula, is the segmentation loss function for the multi-teacher model, is the loss function calculated based on the three-dimensional KL divergence, For the The weights of the teacher models, , , are the height, width and depth of the 3D image, respectively. is the temperature coefficient, is the predicted value of the student model, represents the Shannon entropy, express Divergence, is the softmax function, is a hyperparameter.

5. The unsupervised domain adaptation medical image segmentation method based on the teacher fusion model according to claim 4, characterized in that: In step S3, the student model training includes performing knowledge distillation based on the trained first teacher model and the second teacher model, mapping the convolution layer weight parameters to singular value parameters using singular value decomposition, and training by updating the singular value parameters, as follows: In the formula, , They are the pre-trained weight layer and the weight layer updated after mapping, , are the left and right singular vectors, respectively, are the singular value parameters mapped after pre-training of the first teacher model and the second teacher model, No. The singular value parameters updated at the iteration.

Images