Unlabeled Data Guided Semi-supervised
Histopathology Image Segmentation
^††thanks: 978-1-7281-6215-7/20/$31.00 ©2020 IEEE.

MUNIT [18] is a state-of-the-art approach for cross-domain transformation, which explicitly disentangles image representation into content and style and provides a certain degree of explainability of the extracted latent image representation. It assumes that the source domain and target domain have distinguishable style spaces. But, in our problem, this is not the case (i.e., our source and target domains significantly overlap), and MUNIT cannot yield diversified images by transferring style from one image to another since both images are from the same dataset. Image-based methods can transfer styles between two arbitrary images; but, their representations of style and content are both in high-dimensional latent space (e.g., feature maps [17] or matrices [27, 16]), and thus it is hard to explore the image characteristics distributions. We will show how to utilize domain-based and image-based methods to (1) generate style-diversified images for model training and (2) make the image characteristics concisely represented so that their distributions can be effectively explored for segmentation.

We extract features with an Image Generation Module. To capture the local contents, we consider image characteristics distributions on image patches uniformly cropped. Given a set of $M$ image patches, $\mathcal{X}=\{x_{a}\}_{a=1}^{M}$, from our dataset, each patch $x_{a}$ can be encoded into a style vector $s_{a}=E^{s}(x_{a})$ by a style encoder $E^{s}$ and a content vector $c_{a}=E^{c}(x_{a})$ by a content encoder $E^{c}$. All the $s_{a}$’s and $c_{a}$’s form a style space $S=\{s_{a}\}_{a=1}^{M}$ and a content space $C=\{c_{a}\}_{a=1}^{M}$, respectively. A new image patch can be synthesized by a generator $G(c_{a},s_{b})$ combining any content vector $c_{a}$ and style vector $s_{b}$. The total loss consists of our style matching loss, GAN loss, and reconstruction loss [18], defined as:

where the $w_{i}$’s are hyper-parameters for controlling the weight of each term. We design the style matching loss for achieving two goals: (1) preventing the training collapse (which would cause all information being encoded into content and the generator ignoring the style); (2) making the point distribution $\{s_{a}\}_{a=1}^{M}$ in the low-dimensional latent space $S$ reflect the image style distribution. The core of the style matching loss is an interpolation relation, which associates the point-wise distance in the embedded space with the similarity between image patches.

To attain the first goal, when transferring a patch $x_{a}$ to the style $s_{b}=E^{s}(x_{b})$ of a patch $x_{b}$, we encourage the generated patch $x_{g1}=G(E^{c}(x_{a}),s_{b})$ to have the target style $s_{b}$, by minimizing a style similarity metric [17] between $x_{g1}$ and $x_{b}$:

where $J[x]$ is a Gram matrix calculated on the vectorized VGG features of layer $\ell$, $\alpha_{\ell}$ is the weight of layer ${\ell}$, and $N_{\ell}$ is the number of filters in layer ${\ell}$.

For the second goal, we encourage that two patches with a higher similarity attain a smaller distance in $S$. We train the model by enforcing the interpolation property: For a linear interpolation in the latent space, $s_{g2}=(1-\lambda)E^{s}(x_{a})+\lambda E^{s}(x_{b})$, where $\lambda\in[0,1]$, the style of the generated patch $x_{g2}=G(E^{c}(x_{a}),s_{g2})$ smoothly transfers to $s_{b}$ as the distance between $s_{g2}$ and $s_{b}$ gets closer. This property can be learned by optimizing the following style matching loss:

where $\lambda$ is uniformly selected from $[0,1]$ in training. By the interpolation property, those patches with similar styles tend to have closer style vectors, yet dissimilar patches tend to have style vectors far apart from one another. Hence, the distribution of style vectors in the style space is encouraged to reflect the patch style distribution in the given dataset. Only Eq. (3) is practically used in our training, as Eq. (2) is a special case of Eq. (3) with $\lambda=1$.

After obtaining the extracted content set $C=\{c_{a}\}_{a=1}^{M}$ and style set $S=\{s_{a}\}_{a=1}^{M}$ from all the patches, we conduct agglomerative clustering to attain content clusters $\{C_{i}\}_{i=1}^{m}$ of $C$ and style clusters $\{S_{j}\}_{j=1}^{n}$ of $S$, where $m$ and $n$ are the numbers of such clusters, respectively. A patch space $H$ is defined by the Cartesian product of $\{C_{i}\}_{i=1}^{m}$ and $\{S_{j}\}_{j=1}^{n}$: $H=\{C_{i}\}_{i=1}^{m}\times\{S_{j}\}_{j=1}^{n}=\{H_{ij}=\{(c,s)\ |\ c\in C_{i},s\in S_{j}\}\}$. Each $H_{ij}$ corresponds to possible patches whose content vectors belong to content cluster $C_{i}$ and style vectors belong to style cluster $S_{j}$. We explore the statistics of the patch space $H$ based on the information of each $H_{ij}$, whose numbers of labeled patches and unlabeled patches are denoted by $N(H_{ij}^{label})$ and $N(H_{ij}^{unlabel})$, respectively.

The same reconstruction loss for images and the latent space as in MUNIT [18] is applied to the generated images. The content and style should be consistent after decoding and encoding. The latent reconstruction loss is computed as:

where $L^{x}_{recon}$, $L^{c}_{recon}$, and $L^{s}_{recon}$ are the image, content, and style reconstruction losses, respectively.

We seek to attain appearance realism of the generated images by training the discriminator using samples generated with interpolated styles. The generated image distribution is imposed to be the same as the original image distribution. The realism loss function is:

where $x_{a}$ follows the original image distribution ($x_{a}\sim p(x)$), $s=(1-\lambda)*s_{a}+\lambda*s_{b}$ with $s_{a},s_{b}\sim p(s)$, and $c_{a}$ is encoded from $x_{a}$ with $c_{a}=E^{c}(x_{a})$.

Using extracted $c_{i}\in C$ and $s_{j}\in S$, we can generate a set of patches that may potentially help segmentation models in a basic random generation setting (by uniform sampling). However, by doing so, (1) biomedically invalid patches may be generated without considering content similarity between the source images and target images, and (2) underrepresented and rare hard cases are overwhelmed by other “common” image characteristics. In this section, we first discuss a basic random generation strategy with content matching, and then further improve the effectiveness by proposing our distribution matching policy and hard case covering policy.

Still, $|S_{gen}|\gg|S_{label}|$, and we need to strike a balance between the effectiveness of the augmented set and avoiding excessively perturbing the original labeled (training) data. $S_{gen}$ is used with a probability $R_{a}$ (a probability $1-R_{a}$ for $S_{label}$).

Thus, we propose to sample patches based on the statistics of the unlabeled data. Especially, we focus on the underrepresented clusters with image characteristics that appear frequently in the unlabeled set but rarely in the labeled set. Each $H_{ij}$ is selected with a probability $N(H^{unlabel}_{ij})/{\Sigma_{k,l}N(H^{unlabel}_{kl})}$. Then we uniformly select samples from $S_{gen}\cap H_{ij}$ (or $S_{label}\cap H_{ij}$) with probability $R_{a}$ (or $1-R_{a}$). In this way, the underrepresented clusters benefit the most as they get higher chances to be included in the augmented training data.

To address this issue, we propose to identify these hard cases automatically. Inspired by bootstraping-based methods [28, 1], we quantify the uncertainty of a segmentation model by its prediction variance when the input is transferred to different target styles. For $n$ style clusters $\{S_{j}\}^{n}_{j=1}$, we select one representative target style $Rep(S_{l})$ from each $S_{l}$ with the minimum sum of distances to all the other styles $s_{k}\in S_{l}$. For each $H_{ij}$, its uncertainty $U_{ij}$ for a segmentation model $Seg$ is calculated as:

Selecting uncertain clusters more frequently will help the segmentation model reduce potential prediction errors. Thus, we upweight the probability of sampling uncertain clusters. Each $H_{ij}$ is selected with a probability $U_{ij}/{\Sigma_{k,l}U_{kl}}$. This policy has two advantages comparing to [1]: (1) our method does not require training separate versions of the segmentation model; (2) our uncertainty value can better assess the segmentation quality, as the experiments show that our coefficient $R^{2}$ is 15% higher.

Overall, these two policies are complementary to each other and can be combined with Mixed policy: sample patches with an average probability of distribution matching and hard case covering (i.e., the probability is $N(H^{unlabel}_{ij})/(2\Sigma_{k,l}N(H^{unlabel}_{kl}))+U_{ij}/(2\Sigma_{k,l}U_{kl})$), which comprehensively covers the unlabeled data especially for the hard cases.

We use two main histopathology image datasets in our experiments: GlaS [26] of glands and MoNuSeg [25] of nuclei. We uniformly crop patches of size $384\times 384$, with a fixed step size 64 along the width and height, from all the images. The GlaS dataset has 85 images (1697 patches) for training and 80 images (1559 patches) for test. We use 8 style clusters and 5 content clusters to study the data distributions. The MoNuSeg dataset has 30 high-resolution images, with 16 images (1296 patches) of 4 organs for training and 14 images (1134 patches) of 7 organs for test (note that 3 test organs are not seen in training). We use 7 style clusters and 3 content clusters. In both the datasets, we treat either the training set or a subset of the training set as the labeled set and the test set or a subset of the training set (by ignoring the labels of this subset) as the unlabeled set in our experiments.

We use four segmentation models: DCN [24], MildNet [2], FullNet [29], and CIA-Net [30]. Here, DCN and MildNet are common segmentation models. FullNet and CIA-Net attain state-of-the-art performance on the GlaS and MoNuSeg datasets, respectively. The default value of the probability $R_{a}$ is set as 0.15. The weights $w_{1}$, $w_{2}$, and $w_{3}$ in the loss function of Eq. (1) are set as 0.002, 1, and 10, respectively. Basic augmentation operations (e.g., flipping, rotation) are applied in all the experiments (denoted as ${\rm Aug_{basic}}$).

Our experiments consist of three parts. In qualitative results, we show that our image generation module can effectively generate realistic images and cluster patches based on the style and content similarity. In quantitative results, we evaluate the effectiveness of our new method on improving segmentation performance under both the inductive learning and transductive learning settings of SSL scenarios. In ablation study, we evaluate the model sensitivity to different generation algorithms and policies.

Transductive learning means in the training process, the test images are shown as unlabeled images to the model (i.e., $S_{unlabel}$ is the original test set). In biomedical imaging, transductive learning has found wide applications. For example, in high-throughput experiments, a large amount of images (with potentially different styles) needs to be accurately segmented for further analysis. Transductive learning can be very useful in such scenarios. Tables II shows the segmentation results. We evaluate the segmentation models (DCN [24], MildNet [2], and CIA-Net [30]) with ${\rm Aug_{basic}}$ on our proposed random generation and mixed policy. First, our method can effectively boost the performance of all the models, yielding better results than the state-of-the-art performance. Second, our method works well especially for the rare hard cases (e.g., unseen organs in MoNuSeg).

The contribution of each component in our generation policy is evaluated, including content matching (CM), distribution matching (DM), and hard case covering (HC). The performance is affected when only one policy is used or CM is not applied. The three components of our policy can help improve segmentation performance in a complementary manner.