An unsupervised domain adaptive image segmentation method based on uncertainty estimation

Abstract

The present invention relates to the field of image segmentation technology, and specifically to an unsupervised domain adaptive image segmentation method based on uncertainty estimation, comprising the following steps: constructing a generative segmentation model based on a U‑Net model and a conditional variational autoencoder; constructing a teacher network and a student network based on the structure of the generative segmentation model, and training the teacher network with a source domain model; obtaining pseudo labels of target domain data and training the student network; predicting and calculating entropy values ​​for the teacher network and the student network, and incorporating them into the optimization function of the student network training; and performing image segmentation through the trained student network. The present invention introduces a generative model of a probabilistic segmentation network, which can randomly sample to obtain multiple segmentation results, and obtain the uncertainty of the model through the differences between them. An autoencoder branch is also added to the model, so that the model extracts rich features and regularizes the model at the same time.

Description

Technical Field

[0001] The present invention relates to the technical field of image segmentation, and in particular to an unsupervised domain adaptive image segmentation method based on uncertainty estimation. Background Art

[0002] Image segmentation using deep learning algorithms has been applied in many fields, but its accuracy is highly dependent on training data. In practical applications, a large number of labeled training data sets are usually required for training. In particular, semantic segmentation tasks require manual pixel-by-pixel annotation of labels, which is a time-consuming and laborious task. In order to solve the segmentation of unlabeled data, unsupervised domain adaptation is an effective method. This method usually first trains a model using a labeled source domain, then uses the model to predict the target domain data to obtain pseudo labels, and finally trains the target domain data together. However, due to the domain migration between the training domain and the test domain, the generated pseudo labels often contain noise. Directly using noisy labels for training the target domain will affect the performance of the model.

[0003] In the prior art, manually set thresholds are used to filter out pseudo-labels with low confidence. However, this threshold is closely related to the similarity between the source domain and the target domain, and is difficult to estimate in advance. If the threshold is set too high, some pseudo-labels that should be accepted may be excluded, resulting in the model being unable to fully utilize these data for learning. Moreover, for different segmentation categories, it is difficult to determine the same threshold as the same threshold. Pixels in small or edge locations that are difficult to segment usually have lower confidence. Therefore, the setting of the threshold must take into account both the confidence and the position and size of the pixel. Moreover, a fixed threshold is difficult to match specific needs, and each pixel in the segmentation map needs to be processed differently.

[0004] In order to achieve unsupervised domain adaptive segmentation, some existing technologies use adversarial networks to minimize the distribution differences between domains at the pixel level, feature level, semantic level, etc. to achieve domain alignment. However, their essence is to align the common features of the source domain and the target domain, and some features unique to the target domain will be ignored.

[0005] In order to better analyze the confidence of model predictions, it is crucial to measure the uncertainty of the model. Some existing technologies use dropout to quantify the uncertainty of the model. Since some neurons are randomly discarded during each forward propagation, the model output obtained by each forward propagation may be different. Although the uncertainty can be effectively calculated, this may cause the model to produce inconsistent outputs. Introducing the Bayesian method into the model that produces a certain output can predict the uncertainty of the model, but the Bayesian network obtains uncertainty by evaluating the distribution of model parameters. The Bayesian method is not very scalable for high-dimensional features, so the application effect on some high-resolution images is poor.

[0006] In summary, existing semantic segmentation requires a large amount of manual annotation per pixel, which is an expensive and time-consuming task. When there is a domain shift between the training dataset and the test dataset, the pseudo-labels predicted by the model of the training set for the test set usually contain noise. Using the noisy pseudo-labels to train the test set will cause the model to learn incorrect representations. Low-confidence labels are filtered out by manually setting a specific threshold, but this threshold is difficult to determine. Using an adversarial network to align the same features of the two domains will ignore the features unique to the test domain. Summary of the invention

[0007] The purpose of the present invention is to provide an unsupervised domain adaptive image segmentation method based on uncertainty estimation, which solves the problems that the existing manual annotation of semantic segmentation is expensive and time-consuming, the use of pseudo-labels containing noise to train the test set will cause the model to learn incorrect representations, it is difficult to manually set a specific threshold when filtering low-confidence labels, and the use of adversarial networks to align the same features of two domains will ignore the unique features of the test domain.

[0008] In order to achieve the above object, the technical solution of the present invention is as follows: an unsupervised domain adaptive image segmentation method based on uncertainty estimation, comprising the following steps:

[0009] Build a generative segmentation model based on the U-Net model and conditional variational autoencoder;

[0010] The teacher network and student network are constructed based on the generative segmentation model structure, and the teacher network is trained using source domain data.

[0011] Obtain pseudo labels for target domain data and train the student network;

[0012] The entropy values ​​of the predictions of the teacher network and the student network are calculated and incorporated into the optimization function of the student network training;

[0013] Image segmentation is performed using the trained student network.

[0014] Furthermore, the weights of the teacher network are used as the initial weights of the student network, and the weights of the teacher network during training are the exponential moving average of the weights of the student network;

[0015] First, the teacher network is trained with full supervision using labeled source domain data;

[0016] Then, the target domain data is sent to the teacher network, and the prediction results are used as pseudo labels of the target domain data. The target domain data with pseudo labels is used to perform full-supervised training on the student network.

[0017] Furthermore, the target domain data is augmented and sent to the student network, and the original data is sent to the teacher network;

[0018] The entropy values ​​of the predictions of the teacher network and the student network are calculated, and the maximum value is taken as the uncertainty index, and then the uncertainty value is incorporated into the optimization objective function of the student network training;

[0019] During the student network training process, noisy pseudo-labels are corrected pixel by pixel using automatic thresholding.

[0020] Furthermore, the training of the generative segmentation model includes:

[0021] The input image X is fed into the prior encoder and the U-Net model, and the input image X and its label Y are concatenated and fed into the posterior encoder;

[0022] The prior encoder is used to transform the input data into prior (X)∈R N , with variance σ prior (X)∈R N The axial Gaussian distribution parameters are encoded into the latent space R N The prior probability distribution P in ;

[0023] The posterior encoder is used to transform the input data into post (X,Y)∈R N , with variance σ post (X,Y)∈R N The axial Gaussian distribution parameters are encoded into the posterior probability distribution Q in the latent space;

[0024] The prior encoder and the posterior encoder map their respective inputs to a low-dimensional latent space R N ;

[0025] Under the condition of the preset input image, multiple random samplings are performed from the prior probability distribution, the mean is taken and concatenated with the output of the U-Net model, and the expected number of categories is obtained through 1X1 convolution.

[0026] Furthermore, the weight of the teacher network is used as the initial weight of the student network. During training, the weight of the teacher network is the exponential moving average of the weight of the student network. In the mth iteration, the weight is expressed as follows:

[0027] θ′ m =ηθ′ m-1 +(1-η)θ m

[0028] In the formula, θ' m is the teacher network weight of the mth iteration, θ' m-1 is the teacher network weight of the m-1th iteration, θ m is the student network weight of the mth iteration, η is the control exponential moving average update speed, and its value is 0.999.

[0029] Furthermore, the labeled dataset D L As the source domain, its input is represented as x s , the label is y s ; Unlabeled private dataset D U As the target data set, its input is represented as x u ;

[0030] The teacher network is trained with full supervision using the following loss function:

[0031] L s =L(t(x s ),y s )+D KL (Q||P)+βL2

[0032] In the formula, L is the cross entropy loss, t represents the teacher model, and D KL is the divergence loss between the prior probability distribution and the posterior probability distribution. L2 means that the VAE branch in the U-Net model uses the mean square error to measure the reconstruction loss between the reconstructed image and the original image. β is set to 0.1. During the training process, the cross entropy loss and divergence loss are focused on. The addition of the reconstruction branch is to add additional guidance and regularization to the model, so its loss weight is reduced.

[0033] The target domain data x u Send it to the teacher network to get the prediction y' u , y' u The student network is trained using a loss function that is regarded as pseudo labels of the target domain data.

[0034] Furthermore, the entropy value expression is as follows:

[0035]

[0036] Where P(i) is the probability output predicted by the teacher network and the student network, and the entropy values ​​of the predicted outputs of the teacher network and the student network are calculated, which are expressed as H t and H s , take its maximum value as the final uncertainty index D e =max{H t ,H s};

[0037] The entropy value is negatively exponentially included in the cross entropy loss function of the student network training. The loss function of the student network training is expressed as follows:

[0038] L u =E[exp{-D e}L(s(x u ),y' u )]+D KL (P||Q)+βL2

[0039] Where D e The value of is from 0 to positive infinity. The negative exponent is used to compress the entropy value to a range of 0 to 1. The larger the entropy value of the model predicts, the greater the uncertainty it represents. When calculating the corresponding loss, its weight should be reduced. Therefore, it is necessary to take a negative exponent for the entropy value. s represents the student model.

[0040] The above scheme has the following beneficial effects:

[0041] 1. The present invention can use a large number of labeled public datasets to solve the semantic segmentation of a small number of sample unlabeled private datasets.

[0042] 2. The present invention introduces a generative model of a probabilistic segmentation network, which can obtain multiple segmentation results by random sampling, and obtain the uncertainty of the model through the differences between them. An autoencoder branch is also added to the model, so that the model can extract rich features and regularize the model at the same time.

[0043] 3. The present invention adopts teacher and student networks for training, obtains model uncertainty through the entropy values ​​predicted by the two, and then incorporates it into the optimization function of model training as an automatic threshold to correct pseudo-labels, which is more scalable than the manually set fixed threshold method in the existing method.

[0044] Additional aspects and advantages of the present invention will be given in part in the following description and in part will be obvious from the following description, or will be learned through practice of the present invention. BRIEF DESCRIPTION OF THE DRAWINGS

[0045] Figure 1 A schematic diagram of a neural network architecture of a generative segmentation model of an embodiment of an unsupervised domain adaptive image segmentation method based on uncertainty estimation of the present invention;

[0046] Figure 2 A schematic diagram of a model training method of an embodiment of an unsupervised domain adaptive image segmentation method based on uncertainty estimation of the present invention;

[0047] Figure 3 A labeled source domain data graph of an embodiment of the unsupervised domain adaptive image segmentation method based on uncertainty estimation of the present invention;

[0048] Figure 4 This is an unlabeled target domain data graph of an embodiment of the unsupervised domain adaptive image segmentation method based on uncertainty estimation of the present invention. DETAILED DESCRIPTION

[0049] The technical solution of the present invention will be described clearly and completely below in conjunction with the accompanying drawings. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention.

[0050] The following is further described in detail through specific implementation methods:

[0051] Implementation example Figure 1-Figure 4 As shown: An unsupervised domain adaptive image segmentation method based on uncertainty estimation comprises the following steps:

[0052] S100. Building a generative segmentation model based on U-Net model and conditional variational autoencoder.

[0053] Among them, the U-Net model can provide accurate segmentation in many fields. The conditional variational autoencoder can simulate complex distributions and support sampling. It captures the uncertainty of the model in the form of probability distribution. Under this model, multiple random samplings of the distribution of the latent space can obtain multiple different but credible segmentation results. The uncertainty of the model can be estimated based on the differences between these results. In addition, in order to extract richer feature information and alleviate overfitting, an autoencoder branch is added to the bottleneck layer of the U-Net model to reconstruct the image.

[0054] In the training process of the generative segmentation model, the input image X is sent to the prior encoder and the U-Net model, and the input image X and its label Y are concatenated and sent to the posterior encoder; the prior encoder is used to convert the input data into a prior (X)∈R N , with variance σ prior (X)∈R N The axial Gaussian distribution parameters are encoded into the latent space R NThe prior probability distribution P in the posterior encoder is used to transform the input data through the mean u post (X,Y)∈R N , with variance σ post (X,Y)∈R N The axial Gaussian distribution parameters of are encoded into the posterior probability distribution Q in the latent space.

[0055] The prior encoder and the posterior encoder map their respective inputs to a low-dimensional latent space R N ; Under the condition of the preset input image, multiple random samplings are performed from the prior probability distribution, the mean is taken and concatenated with the output of the U-Net model, and the expected number of categories is obtained through 1X1 convolution. During the training process, the KL divergence loss of the prior probability distribution and the posterior probability distribution is minimized:

[0056]

[0057] Where P(X) represents the prior probability distribution and Q(X) represents the posterior probability distribution.

[0058] The parameters of the prior probability distribution will also change. Since the input of the posterior encoder includes label information, the prior probability distribution will learn a similar distribution only under the input image. Randomly sample a sample z from the prior probability distribution P i ∈R N , which can be expressed as:

[0059] z i ~P(·|X)=N(u prior (X),diag(σ prior (X)))

[0060] Where N represents Gaussian distribution, and the diagonal matrix diag here indicates that each dimension is independent of each other.

[0061] The sample is expanded to the same shape as the output of the U-Net segmentation model, and then a 1X1 convolution is performed to obtain the final prediction; the output S can be expressed as:

[0062] S=f 1X1 (f Unet (X),z i )

[0063] In the formula, f Unet (X) represents the encoder branch of the Unet network, z i represents the sample sampled from the latent space, which can be expressed as:

[0064] z~N(u(X),diag(σ(X)))

[0065] Where z is the latent variable sampled from a Gaussian distribution given an input image. The mean and standard deviation of this Gaussian distribution are functions of X and can be adjusted based on the specific input condition X.

[0066] S200. Construct a teacher network and a student network based on the generative segmentation model structure, and train the teacher network with the source domain model.

[0067] The training of the teacher-student model adopts the Mean Teacher method. The teacher network t uses the generative segmentation model structure in S100 to construct the teacher network and the student network. The weight of the teacher network is used as the initial weight of the student network. The weight of the teacher network during training is the exponential moving average (EMA) of the student network weight. In the mth iteration, the weight is expressed as follows:

[0068] θ′ m =ηθ′ m-1 +(1-η)θ m

[0069] In the formula, θ' m is the teacher network weight of the mth iteration, θ' m-1 is the teacher network weight of the m-1th iteration, θ m is the student network weight of the mth iteration, η is the control exponential moving average update speed, and its value is 0.999.

[0070] Labeled dataset D L As the source domain dataset, its input is represented as x s , the label is y s First, we use the labeled source domain dataset (x s ,y s ) performs the training process in S100, and uses the corresponding loss function to perform full supervision training on the teacher network t, and its loss function is expressed as:

[0071] L s =L(t(x s ),y s )+D KL (Q||P)+βL2

[0072] In the formula, L is the cross entropy loss, t represents the teacher model, and D KL is the divergence loss between the prior probability distribution and the posterior probability distribution. L2 means that the VAE branch in the U-Net model uses the mean square error to measure the reconstruction loss between the reconstructed image and the original image. β is set to 0.1. During the training process, the cross entropy loss and divergence loss are focused on. The addition of the reconstruction branch is to add additional guidance and regularization to the model, so its loss weight is reduced.

[0073] S300. Obtain pseudo labels for target domain data and train the student network.

[0074] The student network s uses the model structure of S100 and the unlabeled private dataset D U As the target domain dataset, its input is represented as x u . The target domain data x u After data enhancement, it is fed into the teacher network t in step (S200) to obtain the output prediction, which is expressed as:

[0075] y' u =t(x u )

[0076] S400. Calculate the entropy values ​​of the teacher network and the student network and incorporate them into the optimization function of the student network training;

[0077] The entropy value is calculated for the predicted output of the teacher network t and the student network s. The entropy value expression is as follows:

[0078]

[0079] Where P(i) is the probability output predicted by the teacher network and the student network, and the entropy values ​​of the predicted outputs of the teacher network and the student network are calculated, which are expressed as H t and H s , take its maximum value as the final uncertainty index D e =max{H t ,H s}; Take the negative exponent of the entropy value and incorporate it into the cross entropy loss function of the student network training. The loss function of the student network training is expressed as follows:

[0080] L u =E[exp{-D e}L(s(x u ),y' u )]+αD KL (P||Q)+βL2

[0081] Where D e The value of is from 0 to positive infinity. The negative exponent is used to compress the entropy value to a range of 0 to 1. The larger the entropy value of the model predicts, the greater the uncertainty it represents. When calculating the corresponding loss, its weight should be reduced. Therefore, it is necessary to take a negative exponent for the entropy value. s represents the student model.

[0082] S500. Image segmentation using the trained student network.

[0083] The original input x of the target domain dataset uThe trained student network s in S400 is sent to the network, and then multiple random samples are taken from the prior probability distribution P. The mean is taken as the output of the variational autoencoder and the output prediction of the U-Net model for concatenation. Finally, the final output prediction is obtained through a 1X1 convolution.

[0084] Obviously, the above embodiments are merely examples for the purpose of clear explanation, and are not intended to limit the implementation methods. For those skilled in the art, other different forms of changes or modifications can be made based on the above description. It is not necessary and impossible to list all the implementation methods here. The obvious changes or modifications derived therefrom are still within the scope of protection of the invention.

Claims

1. An unsupervised domain adaptive image segmentation method based on uncertainty estimation, characterized in that: The steps include: Build a generative segmentation model based on the U-Net model and conditional variational autoencoder; Based on the generative segmentation model structure, the teacher network and the student network are constructed, and the source domain model is trained on the teacher network. Obtain pseudo labels for target domain data and train the student network; The entropy values ​​of the predictions of the teacher network and the student network are calculated and incorporated into the optimization function of the student network training; Specifically, the target domain data is augmented and sent to the student network, and the original data is sent to the teacher network; The entropy values ​​of the predictions of the teacher network and the student network are calculated, and the maximum value is taken as the uncertainty index, and then the uncertainty value is incorporated into the optimization objective function of the student network training; During the student network training process, the noisy pseudo-labels are corrected pixel by pixel through automatic thresholding; Image segmentation is performed using the trained student network.

2. The unsupervised domain adaptive image segmentation method based on uncertainty estimation according to claim 1, characterized in that: The weight of the teacher network is used as the initial weight of the student network, and the exponential moving average method is used to update the teacher network; First, the teacher network is trained with full supervision using labeled source domain data; Then, the target domain data is fed into the teacher network, and the obtained prediction results are used as pseudo labels of the target domain data. The student network is trained with full supervision using the target domain data with pseudo labels.

3. The unsupervised domain adaptive image segmentation method based on uncertainty estimation according to claim 1, characterized in that: Training of the generative segmentation model involves: The input image X is fed into the prior encoder and the U-Net model, and the input image X and its label Y are concatenated and fed into the posterior encoder; The prior encoder is used to transform the input data into prior (X)∈R N , with variance σ prior (X)∈R N The axial Gaussian distribution parameters are encoded into the latent space R N The prior probability distribution P in ; The posterior encoder is used to transform the input data into post (X, Y)∈R N , with variance σ post (X, Y)∈R N The axial Gaussian distribution parameters are encoded into the posterior probability distribution Q in the latent space; The prior encoder and the posterior encoder map their respective inputs to a low-dimensional latent space R N ; Under the condition of the preset input image, multiple random samplings are performed from the prior probability distribution, the mean is taken and concatenated with the output of the U-Net model, and the expected number of categories is obtained through 1X1 convolution.

4. The unsupervised domain adaptive image segmentation method based on uncertainty estimation according to claim 2, characterized in that: The weight of the teacher network is used as the initial weight of the student network. During training, the weight of the teacher network is the exponential moving average of the weight of the student network. In the mth iteration, the weight is expressed as follows: θ′ m =eth′ m-1 +(1-η)θ m In the formula, θ′ m is the teacher network weight of the mth iteration, θ′ m-1 is the teacher network weight of the m-1th iteration, θ m is the student network weight of the mth iteration, η is the control exponential moving average update speed, and its value is 0.

999.

5. The unsupervised domain adaptive image segmentation method based on uncertainty estimation according to claim 3, characterized in that: Labeled dataset D L As the source domain, its input is represented as x s , the label is y s ; Unlabeled private dataset D U As the target data set, its input is represented as x u ; The teacher network is trained with full supervision using the following loss function: L s =L(t(x s ),y s )+D KL (Q||P)+βL2 In the formula, L is the cross entropy loss, t represents the teacher model, and D KL is the divergence loss between the prior probability distribution and the posterior probability distribution. L2 means that the VAE branch in the U-Net model uses the mean square error to measure the reconstruction loss between the reconstructed image and the original image. β is set to 0.

1. During the training process, the cross entropy loss and divergence loss are focused on. The addition of the reconstruction branch is to add additional guidance and regularization to the model, so its loss weight is reduced. The target domain data x u Send it to the teacher network to get the prediction y′ u , y′ u The student network is trained using a loss function that is regarded as pseudo labels of the target domain data.

6. The unsupervised domain adaptive image segmentation method based on uncertainty estimation according to claim 5, characterized in that: The entropy value expression is as follows: Where P(i) is the probability output predicted by the teacher network and the student network, and the entropy values ​​of the predicted outputs of the teacher network and the student network are calculated, which are expressed as H t and H s , take its maximum value as the final uncertainty index D e =max{H t , H s }; The entropy value is negatively exponentially included in the cross entropy loss function of the student network training. The loss function of the student network training is expressed as follows: L u =E[exp{-D e }L(s(x u ),y′ u )]+D KL (P||Q)+βL2 Where D e The value of is from 0 to positive infinity. The negative exponent is used to compress the entropy value to a range of 0 to 1. The larger the entropy value of the model predicts, the greater the uncertainty it represents. The corresponding loss should be reduced in weight. Therefore, it is necessary to take a negative exponent for the entropy value. s represents the student model, and y′ u represents a pseudo label.

Images