Unsupervised Domain Adaptation for the Histopathological Cell Segmentation through Self-Ensembling
Editor: TBA
Abstract
Histopathological images are generally considered as the golden standard for clinical diagnosis and cancer grading.
Accurate segmentation of cells/nuclei from histopathological images is a critical step to obtain reliable morphological information for quantitative analysis.
However, cell/nuclei segmentation relies heavily on well-annotated datasets, which are extremely labor-intensive, time-consuming, and expensive in practical applications.
Meanwhile, one might want to ﬁne-tune pretrained models on certain target datasets.
But it is always diﬃcult to collect enough target training images for proper ﬁne-tuning.
Therefore, there is a need for methods that can transfer learned information from one domain to another without additional target annotations.
In this paper, we propose a novel framework for cell segmentation on the unlabeled images through the unsupervised domain adaptation with self-ensembling.
It is achieved by applying generative adversarial networks (GANs) for the unsupervised domain adaptation of cell segmentation crossing diﬀerent tissues.
Images in the source and target domain can be diﬀerentiated through the learned discriminator.
Meanwhile, we present a self-ensembling model to consider the source and the target domain together as a semi-supervised segmentation task to reduce the diﬀerences of outputs.
Additionally, we introduce conditional random ﬁeld (CRF) as post-processing to preserve the local consistency on the outputs.
We validate our framework with unsupervised domain adaptation on three public cell segmentation datasets captured from diﬀerent types of tissues, which achieved superior performance in comparison with state-of-the-art.
[Figure]
Reconstructed Image dis L adv L [Figure]
[Figure]
Target Predict
[Figure]
Shared Weights
Conditional Random Field
Rotate & em dice L
L m s e L r ec on s L 1. Introduction The ﬁeld of computer-aided digital pathology plays an important role in diagnosis and treatment. In pathological images, it is self-evident that the importance of cell is almost second to none as it is the basic component of creatures.
As it always takes doctors a large amount of time, labor, and fund to do segmentation manually, a critical need for more accurate, robust, and low-cost nuclei segmentation methods come into being.
Nowadays, cell segmentation in histopathology images has been extensively studied with a variety of deep learning methods.
Inspired by the advance of Fully Convolutional Networks (FCN) (Long and et al., 2015), there is a variety of deep learning methods proposed in the segmentation ﬁeld like U-Net (Ronneberger and et al., 2015) and DeepLab (Chen and et al., 2017).
For example, Chen and et al.
(2016) proposed a deep contour-aware network (DCAN) to establish better segmentation by a multi-task learning framework that learns not only probability maps but also clear contours.
However, well-annotated datasets for cell segmentation are extremely limited in clinical diagnosis.
And unsupervised domain adaptation comes to researcher’s vision as it can reduce manual annotation works and tackle generalization problems among data collected from diﬀerent clinical sites or diﬀerent modalities.
Liu and et al.
(2020) proposed a method based on Cycada (Hoﬀman and et al., 2018) which adds a task re-weighting mechanism along with a nuclei inpainting mechanism to make the framework perform better on data from diﬀerent organs.
Haq and et al.
(2020) proposed a framework based on GAN (Goodfellow and et al., 2014) along with a reconstruction network to do segmentation on unlabeled data from diﬀerent organs.
Despite the above methods have already taken cross domain problems into consideration, there are still multiple challenges in this ﬁeld.
Firstly, current methods may lack robustness and stable performance when segmenting cross-domain cells.
This limitation is critical for cell segmentation tasks, where the annotations are hard to obtain.
Secondly, existing works only consider the domain adaptation within the same tissue, which cannot be widely applicable in clinical diagnosis.
Taking the challenges above into account, this paper proposes a novel framework for cross-domain cell segmentation.
Particularly, as presented in Fig. 1, we ﬁrst apply a semisupervised method with self-ensembling.
The framework can extract features from the source domain with data augmentation to improve the segmentation performance.
Then the unsupervised domain adaptation can eﬀectively transfer cell segmentation results across diﬀerent tissues or segmentation using diﬀerent staining methods.
The contribution of this paper can be summarized in three aspects: 1) we develop a new framework that contains semi-supervised segmentation with domain adaptation on the target domain; 2) we introduce a semi-supervised framework to do data augmentation in the feature extraction part and a CRF (Boykov and et al., 2004) module in the post-processing part to improve the robustness of the cell segmentation; 3) extensive and comprehensive experiments are carried out on three datasets from diﬀerent tissues to demonstrate the eﬀectiveness of our proposed methods.
2. Methodology
Semi-supervised Segmentation: Formally, in the cell segmentation problem, the histopathological image patches from the datasets are input X of size
H × W × 3.
Then, we want to predict the segmentation output (cid:98) Y of size H ×
W × 1.
In the source domain, we also have binary masks with pixel-wise ground-truth label Y of size H × W × 1 in our framework.
The segmentation network takes images X as input and uses the segmentation predictions (cid:98) Y of the same width and height as output
, i.e., (cid:98) Y = S ( X ).
We train S to generate the predictions (cid:98) Y s through a semi-supervised segmentation method while using the Y s as the ground-truth input of the source domain.
As for the target domain, there is no label for segmentation in the unsupervised domain adaptation problem.
And the GAN method sometimes may not be able to obtain enough information from the features in the segmentation network.
Therefore, we use the entropy minimization loss to control the weight of the labeled examples, increase the conﬁdence of the segmentation output, and make the model more stable.
In practice, the dice-coeﬃcient loss and the entropy minimization loss are more eﬀective than the normally used binary cross-entropy loss.
So we use both the dice-coeﬃcient loss and the entropy minimization loss as our segmentation loss:
[Equation]
where Y (cid:48) s and (cid:98) Y (cid:48)
s are ﬂatten Y s and (cid:98) Y s respectively.
In practice, the images from the source and the target domain may diﬀer a lot in staining result, clarity, and direction.
Therefore, the data augmentation and the robustness should be taken into consideration.
Accordingly, we apply a self-ensembling method by using rotate transformation in a generalized form.
And we optimize the consistency loss with a teacher model, which shares its weights with the student model.
The key point in the teacher-student learning-based semi-supervised segmentation network is the smoothness assumption.
For example, data points close to each other in the image space are more likely to be close in the label space.
To be speciﬁc, the semi-supervised segmentation tasks can be learned by optimizing a mean square error loss:
[Equation]
where (cid:98) Y s,π is the prediction of the source image which goes through the transformationconsistent regularization named π .
In above, the target data goes through the model twice to get two predictions under diﬀerent perturbations.
The model assumes a dual role as a teacher and as a student.
As a student, it learns a fully-supervised segmentation network.
As a teacher, it generates targets to be used by itself as a student for learning through optimizing the mean square error loss.
By doing so, the self-ensembling method is applied to the segmentation model to make the performance more stable, which may work better in the diagnosis.
The overall segmentation loss function is then deﬁned as:
[Equation]
where L em , L dice and L mse are supervised term (the former two) and regularization term.
As L em and L dice have similar importance, we set the weights of them at 1.
As for L mse , we use a time-dependent warming up function λ as a weighting factor.
This weighting function is a Gaussian ramp-up curve that slowly drop down as:
[Equation]
where E denotes the training epoch, k scales the maximum value of the weighting function, and e deﬁnes the peak.
In our experiments, we empirically set k to 1.0 and e to 30.
Training S with the annotated source data teaches S to make accurate predictions.
At this stage, the segmentation network may not generate correct predictions for target images as there are discrepancies between the source and the target.
Therefore, the S needs to generate target domain predictions closer to the source domain predictions by making the distribution of target predictions (cid:98) Y t closer to (cid:98) Y s .
We deﬁne the adversarial loss as:
[Equation]
where (cid:98) Y t = S ( X t ), and H (cid:48) and W (cid:48) are the height and width of discriminator output D ( (cid:98) Y t ).
This adversarial loss helps S to fool the discriminator so that it could consider (cid:98) Y t as a source domain segmentation prediction.
This semi-supervised network can be treated as the generator of a GAN (Goodfellow and et al., 2014).
And to make the predictions closer to each other, we also need a discriminator.
Discriminator: We introduce a discriminator D into the framework.
It can take the predictions as its input and distinguish whether the input comes from the source domain or the target domain.
To train D, we use a cross-entropy loss as:
[Equation]
where z=0 when D takes target domain prediction as input, and z=1 when input comes from source domain.
Moreover, it is possible that these target predictions are not well-correlated with the target input image.
A network for reconstructing images from the predictions to a similar appearance as input can ensure the correlation between the input and the prediction.
Reconstructor: We use a Reconstruction network R in our framework and consider the segmentation network S as an encoder and the Reconstruction network R as a decoder to reconstruct original images from predictions.
It takes the predictions (cid:98) Y t as inputs and produce the reconstructed image as the output R ( (cid:98) Y t ).
We calculate the reconstruction loss as:
[Equation]
where R ( (cid:98) Y t ) is the output of the Reconstructor for input (cid:98) Y t , and H, W, C are the height, width, and number of channels of the input image X t .
Overall, we optimize the following total loss when training our framework:
[Equation]
where the λ adv and λ recons are the weights to balance above losses.
At this point, our results have learned information from both the labeled source domain and the unlabeled target domain.
Because neighbouring voxels share substantial spatial context, the segmentation results produced by the CNN are smooth.
However, local minimization training and noise may still result in some spurious outputs, like small isolated regions or holes in the predictions, which result from the lack of regional spatial information.
Accordingly, we employ a fully connected conditional random ﬁeld (Kr¨ahenb¨uhl and et al., 2011) as a post-processing step to achieve more structured predictions and constrain the spatial consistency of the results.
For an input image X and its segmentation prediction (cid:98) Y , the Gibbs energy in the CRF model is given by:
[Equation]
the unary potential is the negative log-likelihood ψ u (cid:16) (cid:98) Y i (cid:17) =
− logP ( (cid:98) Y i | X ), where P ( (cid:98) Y i | X ) is the model’s output for pixel i.
The pairwise potentials in our model have the form:
[Equation]
where k ( m ) is Gaussain kernel k ( m )
( f i , f j ) =
exp ( − 12 ( f i − f j ) T Λ ( m ) ( f i − f j )).
The vectors f i and f j are feature vectors for pixels i and j in an arbitrary feature space, w ( m ) are linear combination weights, µ is a label compatibility function.
Each kernel k ( m ) is characterized by a symmetric, positive-deﬁnite precision matrix Λ ( m ) , which deﬁnes its shape.
The contrast-sensitive two-kernel potentials are used for our image segmentation and labeling, consisting of the appearance kernel and the smoothness kernel.
They are deﬁned [Table] in terms of the color vectors I i and I j and positions p i and p j :
[Equation]
the appearance kernel means that nearby pixels which have similar color are likely to be in the same class.
The degrees of nearness and similarity are controlled by parameters Θ α and Θ β .
The smoothness kernel removes small regions that isolated with big ones.
Finally, the weights w (1) and w (2) deﬁne the relative strength of the two factors.
3. Experiment
Datasets: There are three datasets used in our experiments.
Irshad and et al.
(2014) released the KIRC dataset which consists of 463 images of 400 × 400 pixel size from Kidney
Renal Clear cell carcinoma (KIRC).
Naylor and et al.
(2018) released the TNBC dataset that consists of 50 images of 512 × 512 pixel size from Triple Negative Breast Cancer Cell (TNBC).
Hou and et al.
(2020) released the TCIA dataset that consists of 1356 images of 256 × 256 pixel size from 14 diﬀerent cancer types.
We use 97 images from Stomach adenocarcinoma (STAD) among the dataset to maintain the images from one dataset come from one tissue.
All the datasets are extracted at 40 × magniﬁcation and come from whole slide images (WSI).
Settings: We use U-Net (Ronneberger and et al., 2015) as our backbone.
We use 80% of the images for training, 10% for validation and 10% for evaluation.
During the training, we employ Adam optimizer (Kingma and et al., 2014) to optimize the losses with learning rates of 0.0001, 0.001, and 0.001 used in the segmentation network, discriminator, and reconstructor respectively.
We use 0.001 and 0.01 as λ adv and λ recons respectively and train in total 500 epochs.
All experiments are carried out by using Pytorch on a Linux system with 2 RTX 2080Ti and take about 14 GB memory of the graphic cards for 12 hours.
Results:
In our experiment, we use TNBC (Naylor and et al., 2018) and TCIA (Hou and et al., 2020) datasets as the source domain respectively and the other two datasets as the target domain respectively.
Besides, our proposed method is compared with 3 recently proposed related methods.
The ﬁrst one is DA-ADV, a UDA method based on the GAN method which also uses a discriminator like ours and proposed by Dong and et al.
The second one is CBST, another UDA method proposed by Zou and et al.
This is a popular classbalanced self-training framework by generating pseudo labels, which is a diﬀerent method to transfer the domain from ours.
The third one is CellSegUDA, an [Figure] unsupervised adversarial domain adaptation method proposed by Haq and et al.
This method is based on the GAN method and achieved excellent performance on cell segmentation tasks.
We also perform ablation experiments to validate the eﬀectiveness of each component in our framework.
To evaluate the segmentation accuracy of the nuclei instances, we use both the Dice and
IoU metrics as our validation metrics.
As shown in Table 1, compared with the former three baseline methods, our framework’s performance has a signiﬁcant improvement under both the IOU and the Dice metrics.
Since the self-ensembling method can enhance the robustness and the GAN network can beneﬁt from it, our framework works better in comparison with the other three methods.
And in ablation experiments, with individual components removed, the performance of the framework falls, which proves that every component could contribute to the performance of the framework.
And Figure 2 shows the visualization of segmentation results of our framework and 3 compared methods, which indicates that there are signiﬁcant improvements in our framework.
In the results of the target domain, our framework works better than other methods and get a more accurate segmentation result, which beneﬁts from the self-ensembling method and the CRF post-processing step.
4. Conclusions
In this paper, we propose a novel unsupervised domain adaptation method for the cell segmentation across diﬀerent tissues.
The method is based on the GAN framework and suitable for the segmentation of whole-slide images from diﬀerent tissues.
We improve the segmentation performance by introducing a teacher-student learning-based semi-supervised segmentation network, which could help with the data augmentation and to overcome the problem of low robustness.
Meanwhile, we use the conditional random ﬁeld method as our post-processing step to achieve more structured predictions and constrain the spatial consistency of the results.
Based on this work we may apply self-supervised methods in this ﬁeld and study how to make the domain adaptation more accurate in the future