DoNet: Deep De-overlapping Network for Cytology Instance Segmentation

Problem formulation: Following the setting of cell instance segmentation, we are provided with a dataset with $K$ images and the corresponding annotations $\mathcal{D}=\left\{(\mathcal{X}_{k},\mathcal{Y}_{k})\right\}_{k=1}^{K}$, each of which contains annotations of bounding boxes $\mathcal{B}_{k}=\left\{b_{k,i}\right\}^{N_{k}}_{i=1}$, object categories $\mathcal{C}_{k}=\left\{c_{k,i}\right\}^{N_{k}}_{i=1}$, and instance masks $\mathcal{E}_{k}=\left\{e_{k,i}\right\}^{N_{k}}_{i=1}$, where $N_{k}$ denotes the number of instances in the $k$-th image. Here, we focus on segmenting the nucleus and cytoplasm for cell instance $c_{k,i}\in\left\{\text{nuclei},\text{cytoplasm}\right\}$. For each cell cluster, based on the localization relationship among the cells, we decompose the instance mask into the intersection region $\mathcal{O}_{k}=\left\{o_{k,i}\right\}^{N_{k}}_{i=1}$ and the complement region $\mathcal{M}_{k}=\left\{m_{k,i}\right\}^{N_{k}}_{i=1}$ via the logical operation (see in Figure 1), which is used to model their relationship in the DoNet.

Workflow: The flowchart of the proposed DoNet is provided in Figure 2. It takes Mask R-CNN [11] as the base model, followed by a Dual-path Region segmentation Module (DRM) which takes the instance features as input to predict intersection regions $\hat{o}_{k,i}$ and complement regions $\hat{m}_{k,i}$ of cytoplasm via the intersection and complement layers (Section 3.2). After decomposing cell clusters according to the position relationship, a Consistency-guided Recombination Module (CRM) takes instance features from DRM and RoIAlign layer to generate the recombined masks $\hat{e}^{r}_{k,i}$ for consistency regularization (Section 3.2). Then, the recombined cytoplasm prediction is utilized as prior knowledge for intra-cellular object prediction throughout the Mask-guided Region Proposal (MRP) (Section 3.3). We will present the details of our framework in the following sections.

Coarse Mask Segmentation: Previous research works on Mask R-CNN has shown competitive performance in instance segmentation [11, 13, 9], adopted into the proposed DoNet for coarse mask segmentation. Mask R-CNN consists of two stages. The first stage utilizes Feature Pyramid Network (FPN) for feature extraction and a Region Proposal Network (RPN) for candidate object bounding boxes generation. The second stage generates Region-of-Interest (RoI) features $f^{roi}_{k,i}$ via the RoIAlign layer and predicts object classes $\hat{c}_{k,i}$ and bounding boxes $\hat{b}_{k,i}$ from the detection head, and semantic masks $\hat{e}^{c}_{k,i}$ from the Instance Mask Head ($H_{i}$), respectively. Due to limited perception capability in overlapping regions, $\hat{e}^{c}_{k,i}$ may contain ambiguous boundaries. Thus, we denote $\hat{e}^{c}_{k,i}$ as the coarse mask, which provides information for sub-region decomposition in DRM and suppresses the interference from background mimics. Following standard losses in [11], a multi-task loss $\mathcal{L}_{coarse}$ for coarse mask segmentation is formed as,

$$\mathcal{L}_{coarse}=\mathcal{L}_{reg}+\mathcal{L}_{cls}+\mathcal{L}_{cmask},$$

(1)

where $\mathcal{L}_{reg}$ denotes the smooth-L1 loss for bounding box regression, $\mathcal{L}_{cls}$ denotes cross-entropy (CE) loss for classification, and $\mathcal{L}_{cmask}$ denotes pixel-wise CE loss for segmentation.

Dual-path Region Segmentation Module (DRM): As shown in Figure 2, DRM consists of an Intersection Mask Head $H_{o}$ and a Complement Mask Head $H_{m}$ with the same architecture. Let $f_{k,i}^{c}$ denotes the rich semantic feature before the coarse mask prediction in the Instance Mask Head. Both $H_{o}$ and $H_{m}$ take the concatenation of $f^{roi}_{k,i}$ and $f_{k,i}^{c}$ as input to predict the intersection and complement regions $\hat{o}_{k,i},\hat{m}_{k,i}$ in clusters. See the supplementary material for details of the Concatenation Unit. Specifically, each head consists of 4 convolutional layers to generate features in $14\times 14\times 256$, followed by one deconvolutional layer to get the semantic mask with a resolution of $28\times 28\times 1$. Here, we add the pixel-wise CE loss to both heads as the explicit constraint for decomposition,

$\displaystyle\mathcal{L}_{dec}=\frac{1}{K}\sum_{k=1}^{K}\frac{1}{N_{k}}\sum_{i=1}^{N_{k}}\left(\mathcal{L}_{ce}(\hat{o}_{k,i},o_{k,i})+\mathcal{L}_{ce}(\hat{m}_{k,i},m_{k,i})\right).$

(2)

Semantic Consistency-guided Recombination Module (CRM): To further enhance the perception capability for overlapping instances, CRM is designed to encourage DoNet to perceive integral instances. Specifically, let $f_{k,i}^{o}$ and $f_{k,i}^{m}$ denote the features before the last layer in the intersection mask head and the complement mask head. The rich semantic feature for the overlapping and non-overlapping region is considered as the residual information for the complex areas which is fused with $f^{roi}_{k,i}$ and then fed into the recombined mask head. See the supplementary material for details of the Fusion Unit. Noted that we reuse the Instance Mask Head ($H_{i}$) here to predict the integral instance. This combination facilitates CRM to leverage contextual information of overlapping instances, leading to the improvement of perception capability. The refined mask $\hat{e}_{k,i}^{r}$ from CRM is optimized by segmentation loss $\mathcal{L}_{rmask}$,

$$\mathcal{L}_{rmask}=\frac{1}{K}\sum_{k=1}^{K}\frac{1}{N_{k}}\sum_{i=1}^{N_{k}}{\mathcal{L}_{ce}(\hat{e}_{k,i}^{r},e_{k,i})}.$$

(3)

Extended dataset with synthetic clusters: The semi-transparency characteristic of cytoplasm alleviates the effect of high overlap, allowing cytologists to delineate the contour of cellular instances based on expert knowledge. However, extremely challenging labeling and unavoidable label noise limit the availability of large-scale annotated datasets. To tackle this problem, we propose an instance-level data augmentor for overlapping cell data augmentation, which can generate large-scale synthetic cell clusters with controllable overlapping ratios and transparency based on the annotated cellular instances. This synthetic dataset further facilitates the generalization ability, by providing more diverse data to implicitly learn the concept of instance and its components. See the supplementary material for details of the synthetic pipeline. Unless otherwise specified, we do not use this synthetic dataset in the following comparison.

Overall Loss Function: The proposed cell instance segmentation framework can be trained in a supervised manner. The overall objective $\mathcal{L}$ is as follows,

$$\mathcal{L}=\mathcal{L}_{coarse}+\lambda_{dec}\mathcal{L}_{dec}+\lambda_{rmask}\mathcal{L}_{rmask}+\lambda_{cons}\mathcal{L}_{cons},$$

(6)

where $\mathcal{L}_{coarse}$ denotes losses for RoI extraction and coarse mask prediction, $\mathcal{L}_{dec}$ denotes the decomposition loss for intersection and complement regions segmentation, $\mathcal{L}_{rmask}$ is the segmentation loss for refined masks, $\mathcal{L}_{cons}$ supervises the semantic consistency between integral instance and sub-regions. $\lambda_{dec}$, $\lambda_{rmask}$, and $\lambda_{cons}$ are trade-off parameters controlling the importance of each component.

ISBI2014[24]: This is a widely-used dataset from Overlapping Cervical Cytology Image Segmentation Challenge¹¹1 https://cs.adelaide.edu.au/~carneiro/isbi14_challenge/index.html, which consists of 8 extended depth-of-focus (EDF) real cervical cytology images and corresponding synthetic images. It contains high-quality pixel-level annotations for both nuclei and cytoplasm with a resolution of $512\times 512$. We follow the setting in this challenge [24] to use 45, 90 and 810 images for training, validation and testing to evaluate our proposed DoNet in a supervised setting.

CPS[39]: This liquid-based cytology dataset contains 137 labeled images with a resolution of $1000\times 1000$. In total, it contains 4439 cytoplasm and 4789 nuclei annotations. We conduct 3-fold cross-validation on this dataset.

Evaluation metrics: To measure the overall performance of the proposed DoNet, we utilize four commonly-used evaluation metrics in instance segmentation: aggregated Jaccard index (AJI), average Dice coefficient (Dice), F1-score (F1)[20], mean of Average Precision (mAP) [11]. In order to compare our result on the ISBI2014 with previous studies [24], we further adopt evaluation metrics including Dice, object-based false negative rate ($\text{FN}_{o}$), and pixel-based true positive rate ($\text{TP}_{p}$).

Implementation details: We utilize the Mask R-CNN[11] in Detectron2 [36] as the baseline model. We use the ResNet-50-based FPN network in all experiments. During training, we adopt SGD with 0.9 momentum as the optimizer. We set the initial learning rate to 0.001 and add the linear warm-up in the first 1k iterations. We train the network for 60k iterations, decreasing the learning rate by a factor of 0.1 after 50k and 55k iterations.

Effects of Network Components: We perform ablation studies to investigate the effects of the different components of our proposed pipeline for DoNet. The comparison results can be seen in Table 3. By adding DRM for explicit decomposing integral instances into intersection and complement sub-regions, we observe a $3.25\%$ increase in the average mAP of cytoplasm and nuclei on ISBI2014. However, directly adding DRM without the fusion of structural and morphological information, may mislead the model. This issue is observed in the more complex CPS dataset. To alleviate this problem, DoNet takes advantage of the decompose-and-recombined strategy by adding CRM after DRM.This strategy brings a total of $7.34\%$ and $1.74\%$ mAP improvements on two datasets by strengthening the model’s perception of overlapping regions while preserving morphological information.

Design Choice for DRM: We provide detailed comparisons on the ISBI2014 dataset to demonstrate the effectiveness of different components in DRM: 1) $H_{i}$: instance mask head for coarse segmentation only; 2) $H_{o}$: intersection mask head for overlapping region segmentation; 3) $H_{m}$: complement mask head for non-overlapping region segmentation;

Design Choice for CRM: The goal of semantic consistency regularization is to enhance the model’s overlapping reasoning capability by learning the concept of recombined instances from sub-regions. We provide comparisons to demonstrate the design choice (Table 4): 1) CU + FU: integration of rich semantic features as inputs via Concatenation Unit and Fusion Unit. 2) $\mathcal{L}_{cons}$: consistency regularization between the recombined prediction $\hat{e}^{r}_{k,i}$ and the fusion of sub-regions $\hat{o}_{k,i}$, $\hat{m}_{k,i}$. As seen in Table 4, by aggregating the rich semantic feature for overlapping and non-overlapping region via Fusion Unit, we observe an increase of $1.52\%$ average mAP on the ISBI2014 dataset. Adding consistency regularization further improves $0.91\%$ mAP for cytoplasm.