Human-machine Interactive Tissue Prototype Learning for Label-efficient Histopathology Image Segmentation

When pathologists read WSIs, they can recognize different tissues based on comparing local cells’ visual appearance and surrounding micro-environments. Similarly, we need an encoder that can project local visual histopathology patterns into a discriminate space to identify different tissues. Inspired by recent studies on contrastive learning, we adopt SimCLR [7] to pre-train the encoder $f_{\omega}(\cdot)$ without labels in a self-supervised fashion.

The detail of this step is shown in Fig. 2 Step 1. Given a set of WSIs $\mathcal{I}=\{i_{1},i_{2},...,i_{m}\}$, we crop them into non-overlap patches $\mathcal{P}=\{p_{1},p_{2},...,p_{n}\}(n\gg m)$ with small resolution (empirically set as 128$\times$128 in this study). During the training stage, given N patches in a mini-batch, we get 2N patches by applying different augmentations on each patch. Two augmented patches from the same patch are regarded as positive pairs, others are treated as negative pairs. We employ ResNet18 [11] backbone, where the final linear classification layer and global average pooling layer are removed, as the encoder $f_{\omega}(\cdot)$. Each patch $p_{i}\in\mathcal{R}^{128\times 128\times 3}$ is encoded by $f_{\omega}(\cdot)$ to get $h_{i}=f_{\omega}(p_{i}),h_{i}\in\mathcal{R}^{4\times 4\times 512}$, and then feed into global average pooling layer $avg(\cdot)$ to generate an embedding space with a dimension of $\mathcal{R}^{512}$. Same as [7], a projection head $g(\cdot)$ is used to map $h_{i}$ to $z_{i}=g(avg(h_{i})),z_{i}\in\mathcal{R}^{128}$ where the contrastive loss is applied. To achieve our goal, the contrastive loss is utilized to pull the positive pairs close and push the negative pairs away in the embedding space. Given a sample‘s embedding $s$ and $t^{+},t^{-}$ as its positive and negative samples’ embeddings, the formula of contrastive loss is defined as: $L_{s,t^{+},t^{-}}=-\log\frac{\exp(\frac{sim(s,t^{+})}{\tau})}{\exp(\frac{sim(s,t^{+})}{\tau})+\sum_{k^{-}}\exp(\frac{sim(s,t^{-})}{\tau})}$, where $sim(\cdot)$ denotes the similarity function and $\tau$ denotes the temperature parameter. In practice, we utilize cosine similarity $cosine(x,y)=\frac{x\cdot y}{|x||y|}$ and set $\tau$ to 0.5. After training, $f_{\omega}(\cdot)$ is used to project histopathology patches into a discriminative embedding space of $\mathcal{R}^{512}$ for prototype identification.

It is reminded that except for SimCLR [7], other self-supervised pre-training strategies [10, 9] can be adopted to obtain the encoder. Here, we use SimCLR to demonstrate the effectiveness of our framework and leave how different self-supervised strategies can affect the results as a future work.

Having obtained a discriminative histopathology image encoding latent space, we then use clustering to mine potential tissue prototypes (i.e. cluster centroids), which can capture inter- and intra-class heterogeneity of different tissues and represent the entire encoding space in a dictionary. To map these prototypes with corresponding labels, only in this step do we need to integrate experts’ examination. In practice, we can ask an expert to examine sampled representative patches to efficiently determine the labels of corresponding clusters without laboriously evaluating all the patches as shown in Fig. 2 Step 2. In our experiments, we simulate the visual inspection process with ground-truth labels for easy reproducibility and evaluation of the segmentation performance.

To begin with, we illustrate how to perform the clustering with the pre-trained $f_{\omega}(\cdot)$. To reduce computational overhead, we randomly sample a subset $\mathcal{P}_{sub}$ from the entire histopathology patch set $\mathcal{P}$. Then, $f_{\omega}(\cdot)$ and $avg(\cdot)$ are utilized to project patches $\mathcal{P}_{sub}$ into the discriminative embedding space to obtain the embedding set $\mathcal{E}_{sub}$. Next, we adopt unsupervised clustering on $\mathcal{E}_{sub}$ to find tissue prototypes without accessing any manual label. For computation efficiency, K-Means++ [2] is employed to generate $k$ clusters from $\mathcal{E}_{sub}$. Naturally, the first problem comes: how to determine a proper cluster number $k$?

Intuitively, the selection of $k$ plays a trade-off between reducing annotation workload and purifying clusters, i.e., a larger $k$ results in smaller tissue clusters with higher inter-class similarity, but costs more labor as the pathologist needs to inspect. In order to select a good cluster number without accessing any manual labels, we adopt the elbow method [13] to determine proper $k$ where model performance increment is no longer worth additional cost. We use an annotation-free distance-based measurement: inter-class embedding squared distance $D$ to measure the model performance. $D$ is given by: $D=\sum_{j}^{k}{\sum_{*}{||centroid_{j},p_{(j,*)}||_{2}}}$, where $centroid_{j}$ denotes the centroid of cluster $j$ and $p_{(j,*)}$ denotes the embeddings of patches in cluster $j$. Generally, $D$ decreases when $k$ grows up. To search a proper $k$, we can plot the curve of relative reduction $R_{v}=D_{v-1}-D_{v},v\geq 2$, where $v$ is potential values of $k$ (e.g., $v\in[2,3,...,50]$), and choose the elbow point where the reduction of $D$ slows down as the final cluster number.

It is noted that identifying proper cluster number is an open problem, and we also evaluate other clustering methods. Our experiments show that KMeans++ is better with current setting in our following ablation study.

Based on the obtained clusters, then there comes the second problem: how can human experts efficiently determine the tissue types for every cluster? We propose to let human experts only evaluate $t$ sampled representative patches to give the prototype-level labels instead of inspecting all the patches. The $t$ representative patches are sampled by the central sampling strategy. That is, given a target cluster, we sample $t$ patches which are closest to the cluster centroid. More sampling strategies will be discussed in the experiments. To simulate pathologists’ the visual examination process, here we count the proportion of each tissue according to the pixel-level labels. If there exists one tissue type with a proportion of pixels exceeding 80%, we consider this cluster to represent this tissue and add the prototype (the cluster centroid embedding) into the prototype dictionary. Otherwise, we assume the cluster may be a mixture of different tissue types (i.e. the border area between organs), and drop it.

To help readers better understand, we in Fig. 3 illustrate the dictionary building process with the BCSS dataset [1]. We visualize the randomly sampled 20000 patches used for dictionary building with T-SNE in Fig. 3 (a) which is colored according to the manual pixel-level labels. Particularly, these patches are divided into 6 types, of which 5 types (tumor, stroma, inflammatory infiltration, necrosis, and others) are defined by the dataset providers, and the mixture type defined by us. Our tissue dictionary is built on the clustering results drawn in Fig. 3 (b) by K-Means++ with $k=30$. Intuitively, we can observe that the clustering results maintain a high degree of consistency with the manual labels. By sampling 10 representative patches from each cluster center and calculating the tissue proportion, we can find that 13/5/3/1/3 clusters are identified as tumors/stroma/inflammatory infiltration/necrosis/others, correspondingly. Note that Cluster 1 is formed by blank background patches, therefore spreading in a centrosymmetric distribution. There are 5 cluster with multiple tissue mixture patterns, which are excluded from the dictionary.

After above process, we can build the prototype dictionary $\mathcal{C}$ with $l$ prototypes after the above process. The prototype dictionary can be noted as $\mathcal{C}=\{c_{1}:y_{1},c_{2}:y_{2},...,c_{l}:y_{l}\}$, where $\{c_{1},c_{2},...,c_{l}\}$ are prototype embeddings, and $\{y_{1},y_{2},...y_{l}\}$ are prototype-level labels given by the pathologist.

Having obtained the pre-trained $f_{\omega}(\cdot)$ and prototype dictionary $\mathcal{C}$, we can generate coarse segmentation masks for a given WSI $i\in\mathcal{R}^{H\times W\times 3}$ by the process shown in Fig. 2 Step 3. We first feed WSI $i$ into the encoder $f_{\omega}(\cdot)$ to have a $v\in\mathcal{R}^{\frac{H}{32}\times\frac{W}{32}\times 512}$ semantic vector map. Then, we can query $\mathcal{C}$ to determine the tissue type for each location. Intuitively, for each location, we can directly retrieve the nearest prototype to the current semantic vector from $\mathcal{C}$ and label each location with the corresponding tissue type. We termed such a query setting as Direct-Query (DQ). However, DQ treats each location independently ignoring intrinsic similarity between the embeddings. Thus, we further improve DQ with an enhanced Cluster-then-Query (CQ) setting to suppress outlier query results. Particularly, in order to explore the intrinsic similarity between the embeddings in the given feature map, we divide the $\frac{H}{32}\times\frac{W}{32}$ locations into different groups by performing K-Means++ on the $\frac{H}{32}\times\frac{W}{32}$ 512-d vector set. After that, the clustering centroids of different groups are used to query $\mathcal{C}$ for labeling each location. As for the clustering number of the target WSI, we find that it is better to decide the clustering number according to tissue types in the given WSI. According to the DQ based query results, we can roughly know that there may exist $d$ tissue types for the target WSI. Considering the possibility of spatial dispersion of the same tissue in images, We divide the current feature map into $\gamma\times d$ clusters. $\gamma\geq 1$ is the scale factor and set to 5 empirically.

Now, we have a $\frac{H}{32}\times\frac{W}{32}\times 1$ tissue segmentation map, which can be upsampled to restore its resolution for the given WSI. The coarse map has zigzag edges but it can offer dense supervision for segmentation network training. To refine the final prediction, we further use these coarse masks as the pseudo tissue masks to train a segmentation network from scratch. By such, we finally build a bridge from the prototype learning to semantic segmentation.

Datasets: Our experiments are conducted on two public datasets: CAMELYON16: we use the train set which contains 270 WSIs focusing on sentinel lymph nodes with pixel-wise cancerous/normal region annotation [3]. BCSS: contains 151 Region-of-Interest cropped Images from breast cancer WSIs in TCGA with tissue-level annotations [1]. Five tissue types are labeled: tumor, stroma, inflammatory infiltration, necrosis, and ‘others’ (including background, blood, etc). Similar to previous works [17, 21], we generate the 2048$\times$2048 / 1024$\times$1024 image-level data from CAMELYON16/BCSS with the specimen-level pixel size of $1335nm$/$250nm$ to get a proper view of biological tissues as the image-level data. For BCSS, we follow the official train-test split with 2151/976 images for training/testing. For CAMELYON16, we randomly divide patients into the train-test set by 2:1 following the hold-out setting and have 1348/848 images for training/testing.

Note that we achieve better segmentation results with a lower annotation burden than other methods. To compare the labeling cost, we invited 3 pathologists to annotate randomly sampled 20 BCSS patches with 1024x1024 pixels for pixel/image-level labels and use the average time to measure the labeling cost for BCSS training set reported in Table. 2. Obviously, our method is much more labor-saving. It is also reminded that we used 1024 $\times$ 1024 cropped BCSS images but the original WSIs can be with a resolution of $>$10000 $\times$ 10000, i.e. it will be more time-consuming to give pixel/image-level labels for the original WSIs. But our method will not be affected by larger resolutions.