Look in Different Views: Multi-Scheme Regression Guided Cell Instance Segmentation

There are currently two main categories for instance segmentation: single-stage and multi-stage methods. The process of multi-stage methods is to first find out the bounding boxes by means of object detection, then pixel-wise segmentation is been performed within the proposed regions. Mask R-CNN (He et al. 2017) is the most famous method under this approach, which merges a mask branch parallel to the bounding boxes branch for Faster R-CNN (Ren et al. 2017). A large number of Mask R-CNN based instance segmentation works have been proposed in recent years. For example, Cascade Mask R-CNN (Cai and Vasconcelos 2021) uses cascade regressors and detectors to enable instance segmentation head to better fit the data; Hybrid Task Cascade (HTC) (Chen et al. 2019a) crosses the bounding box regression branch and mask prediction branch, and use information from semantic labels to guide instance segmentation. Sample Consistency Network (SCNet) (Vu, Kang, and Yoo 2021) improves the network structure of Cascade Mask R-CNN to ensure that the IoU distribution in training time and inference time is close. In recent years, many excellent works relate to single-stage instance segmentation have been published, some of these methods are faster and more accurate than Mask R-CNN. For example, SOLO (Wang et al. 2020a) transforms the instance segmentation task into a classification task. The main method of SOLO is to divide the image into S*S regions and use two branches to predict category and mask of each instance. SOLOv2 (Wang et al. 2020b) further enhances the performance of SOLO.

For medical image instance segmentation tasks, especially cell instance segmentation tasks, the difficulties encountered in segmentation are mainly the dense distribution of instances, unclear instance boundaries and the variability in the appearance of different instances(Yi et al. 2018b, a). These cause difficulties for the regression of cell locations. For this task, many specialized instance segmentation networks are proposed. For example, DCAN (Li et al. 2020b) is an efficient net which uses two different branches to output the results of object detection and contours segmentation. STARDIST (Schmidt et al. 2018) uses the prior that cells are all convex polygon shapes for segmentation, which requires that the cells in the dataset be convex. CosineEmbedding (Payer et al. 2018) recognizes cells by clustering, while it tends to produce a large number of false positives due to the different clustering results of each cell.

The attention mechanism is a module inspired by human visual mechanism that helps model to focus on the more critical parts of the input image. Attention mechanism was first proposed in a language translation task (Cho et al. 2014), now it has been used for all aspects of tasks including classification (Yang et al. 2016), object detection (Schlemper et al. 2018), and image segmentation (Oktay et al. 2018) and can be divided into two categories: soft attention and hard attention (Xu et al. 2015). Guided Attention Inference Network (GAIN) (Li et al. 2020a) explains what the learner focus on, and can feeds back with direct guidance towards specific tasks. Mask-guided Contrastive Attention Model (Song et al. 2018) helps the feature extractor focus on body region instead of background region, so that the network will be more adaptable to the background. In this paper, we introduce a gaussian cell center annotation to supervise the gaussian guidance attention branch. We take inspiration of some networks based on gaussian guidance. MCNN (Zhang et al. 2016) is proposed to regress gaussian maps for different head sizes with multi-column convolutions. CSRNet (Li, Zhang, and Chen 2018) uses a series of dilated convolutional blocks to have a large field of perception and be able to extract deep features. Ma et al. present a bayesian loss (Ma et al. 2019) which directly use the point annotations as supervision, use expectations instead of pseudo labels to calculate loss. We take advantage of such methods and propose a gaussian guidance attention branch for instance segmentation.

Single cell recognized as multiple cells is a serious weakness in cell instance segmentation. When cells are elongated, they are easily to be incorrectly recognized into two separated smaller cells. Aiming at this weakness, we propose a gaussian guidance attention branch. During training phase, this attention branch will learn to generate gaussian prediction masks for cell images. With optimization, our backbone will pay more attention on the center region of cells. On this branch, we first use the instance segmentation label to construct a gaussian mask for each input image. In this paper, we use the object center positions to construct our gaussian masks. Specifically, around the center position, we generate each cell’s mask processed by a gaussian kernel. We treat the upper limit of the gaussian radius as a hyperparameter in this paper. When the length of the bounding box’s short edge is shorter than this gaussian radius, this short edge’s length will be used as the new gaussian radius to prevent the gaussian circle from overflowing the bounding box range.

In general, low-level features contain more detailed information, which is urgently needed for instance segmentation of densely packed cells. To take full advantage of the multi-level features, we present a bi-directional hourglass module. In this module, a bottom-up fusion is first applied to the pyramid features to make each layer of features obtain detailed features. Then the highest level feature map is processed by multiple up-samplers. After each up-sampling, a channel-wise concatenation is fused with the features at same stage and previous stage (requires an additional down-sampling). This module enables the output features to preserve as much detailed information as possible and necessary high level feature map’s information. Some dilation convolutional blocks are then applied to the output feature. Finally, we adopt a MSE (Mean-Square Error) loss to supervise the learning using the gaussian mask.

		$\displaystyle F_{i+1}^{mid}=f_{i}^{mid}(Concat(F_{i+1}^{in},Pool(F_{i}^{mid})))$		(1)
		$\displaystyle F_{i}^{out}=f_{i}^{out}(F_{i}^{mid}+\beta_{i}Pool(F_{i-1}^{mid})+\gamma_{i}Up(F_{i+1}^{out}))$

where $f$ means convolutional blocks while $\beta_{i}$ and $\gamma_{i}$ are learnable parameters that control the proportion learn from each input feature map. $Up$ indicates upsamplers and $i=1,2,3,\cdots$.

To solve the weakness (b) showing in Fig. 1, we propose GGAB to improve the network’s perception of the cell’s central region. This module can help the feature map to generate sufficient activation in central region, thus prevent the cells from being truncated.

Paralleled to GGAB, we propose a point-regression branch. Similar to GGAB, raw pyramid features is processed by a bi-directional hourglass module. Then, passed through some convolutional blocks, two types of masks are generated, which are center point masks and height-width predictions. These should be supervised by two types of ground truths constructed by instance segmentation labels. The center point mask has a non-zero value only at the cell-centered pixels. The height-width label is a 2-channel tensor that contains the height and width data of each boundary boxes of instance segmentation labels. Only the center point masks are used to guide the instance segmentation branch in the dual-scheme guidance module, while the height-width predictions are only used to assist the PRB to regress the correct cell center’s position.

By applying this module, we aim to solve the weakness (a) showing in Fig. 1. When two cells are densely packed, gaussian masks may have difficulty distinguishing between these two cells, because predicted gaussian masks may also be densely packed. Under this situation, instance segmentation branch may recognize these two cells as one with confusing guidance. On solving this weakness, we propose point regression branch to regress the location of cells in a fine-grained manner.

Although PRB is effective, using PRB alone without GGAB is far less effective than using two branches at the same time. Because PRB cannot generate enough activation in the central region, using PRB alone cannot efficiently resolve weakness (a) showing in Fig. 1.

Gaussian masks obtained in GGAB ($M^{G}\in\mathbb{R}^{H\times W}$) and the point masks obtained in the PRB ($M^{P}\in\mathbb{R}^{H\times W}$) contain information of each cell’s center position. Therefore, the raw pyramid features can enhance cell location ability guided by these two masks. Through dual-scheme guidance, MSRNet can adapt to the cells of various appearances. Specifically, The prediction masks need to be processed by a series of max-pooling layers, then be concatenated into the feature map of the same size in pyramid features. At the same time, output features of the lower level layers will be processed by a downsampler and be concatenated to upper level input features. Convolutional blocks are then applied to the features after concatenation is completed, and the output feature maps have the same number of channels as the raw pyramid features. Finally, raw pyramid features need to be added to these feature maps, which are used for enable the fused feature map to retain the features in the original pyramid features. The framework of our dual-scheme guidance module is shown in Fig. 4. Through this guidance, we obtain the input pyramid features of instance segmentation branch (Ren et al. 2015).

$$F_{i}^{cat}=\left\{\begin{aligned} &Concat(F_{i}^{in},M_{i}^{G},M_{i}^{P},F_{i-1}^{out}),&&i=1,2,3,4\\ &Concat(F_{i}^{in},M^{G},M^{P}),&&n=0\end{aligned}\right.$$

(2)

$$\hfill F_{i}^{out}=f(F_{i}^{cat})+F_{i}^{in},n=0,1,2,3,4$$

(3)

where $F_{i}^{out}$ represents the feature map of the i^th level of features in the output pyramid features, similarly, the $F_{i}^{in}$ represents the i^th level of input pyramid features. The $M_{i}^{G}$ and $M_{i}^{P}$ represents $M^{G}$ and $M^{P}$ feature maps be processed by i downsamplers.

Such processing allows the feature maps at each level to retain the detailed information while accept the useful information from dual-scheme prediction masks.

By using a multi-task loss, we train our MSRNet in an end-to-end manner.

$$\mathcal{L}_{G\_Pred}=\frac{1}{N\times H\times W}\sum_{n=1}^{N}\sum_{\begin{subarray}{c}w=1\\ h=1\end{subarray}}^{H\times W}\left|\left|\widehat{G}(w,h)-G(w,h)\right|\right|^{2}$$

(4)

The gaussian mask prediction loss $\mathcal{L}_{G\_Pred}$ measures the distance between our predict gaussian masks and gaussian ground truths. $N$ represents the batch size. $G_{(w,h)}$ denotes the value of ground truths at the point $(w,h)$, while the $\widehat{G}_{(w,h)}$ denotes the value of predictions at the point $(w,h)$.

$$\mathcal{L}_{P\_Loc}=-\frac{1}{N}\sum_{\begin{subarray}{c}w=1\\ h=1\end{subarray}}^{H\times W}(1-P_{(w,h)})^{\delta_{1}}(\widehat{P}_{(w,h)})^{\delta_{2}}log(1-\widehat{P}_{(w,h)})$$

(5)

where $P_{(w,h)}$ indicates the value of ground truths at the point $(w,h)$, while the $\widehat{P}_{(w,h)}$ represents the value of predictions at the point $(w,h)$. $N$ represents the batch size, and $\delta_{1}$ and $\delta_{2}$ are hyper-parameters.

$$\mathcal{L}_{HW\_Reg}=\frac{1}{n}\sum_{i=0}^{n}(\left|\widehat{c}_{n\_w}-\lfloor\frac{c_{n\_w}}{S}\rfloor\right|+\left|\widehat{c}_{n\_h}-\lfloor\frac{c_{n\_h}}{S}\rfloor\right|)$$

(6)

where n represents the total number of cells in the input image. The $\widehat{c}_{n\_w}$ and $\widehat{c}_{n\_h}$ are width and length of each cell in our prediction. Only when the predicted cell center location is exactly the same as the actual cell location, the cell length and width will be calculated in the Loss function. The ${c_{n\_w}}$ and ${c_{n\_h}}$ represents the ground truths of height and width of cells. S represents the scale_factor. Combining these two parts, the $\mathcal{L}_{P\_Reg}$ is as follows:

$$\mathcal{L}_{P\_Reg}=\beta_{1}\mathcal{L}_{P\_Loc}+\beta_{2}\mathcal{L}_{HW\_Reg}$$

(7)

$\mathcal{L}_{P\_Loc}$ represents the loss of cell centroid regression, while $\mathcal{L}_{HW\_Reg}$ is the loss used in height-width regression. In summary, The multi-task loss is as follows:

$$\mathcal{L}=\lambda_{1}\mathcal{L}_{I}+\lambda_{2}\mathcal{L}_{G\_Pred}+\lambda_{3}\mathcal{L}_{P\_Reg}$$

(8)

where $\mathcal{L}_{I}$ is similar to the Loss function used in the Cascade Mask R-CNN (Cai and Vasconcelos 2021), $\mathcal{L}_{G\_Pred}$ is the loss used in predicting the gaussian masks in the GGAB, $\mathcal{L}_{P\_Reg}$ is the loss used in PRB. The $\lambda_{1}$, $\lambda_{2}$ and $\lambda_{3}$ are all hyper-parameters.

The performance of our method are compared with other SOTA instance segmentation methods, such as Mask R-CNN (He et al. 2017), DCAN (Li et al. 2020b), CosineEmbedding (Payer et al. 2018) and SCNet (Vu, Kang, and Yoo 2021). Specially, due to the CA2.5 and SCIS dataset is very new, few result is available from other papers. So we use a series of SOTA methods on these datasets to get the baseline result. The results of the SOTA instance segmentation methods are shown in Tab. 1 and Tab. 2. Obviously, our proposed method shows the best AP50 and AP75 scores compared with other SOTA methods. The visual comparisons of our proposed method with our baseline SCNet are shown in Fig. 5, clearly, especially under dense conditions, our method performs better performance. The first line of this figure shows a simple case that our approach does not break the validity of the baseline method. Other rows contain some hard cases. For example, a single cell be recognized into two different cells (the cell in the upper left corner of the second row, and the cells on the left and right side of the fifth row, etc.). Besides, cells of the third and forth row show the densely packed cells be recognized into one cell. The last line shows an image that is very difficult for instance segmentation, for which our method has achieve some improvements compared to baseline.

Fig. 7 shows the difference between MSRNet and baseline in terms of feature maps, from which we can directly explain the improvements of instance segmentation results. The first two columns show why our method can correctly recognize densely packed cells: in the original feature map, the feature regions of two cells are almost connected and indistinguishable, while our method’s feature maps shows significant activation in each cell’s central region, and the two activation do not overlap each other. MSRNet can use this feature to distinguish densely packed cells. The third column demonstrates why our method improves the recognition of elongated cells. MSRNet pays more attention to the central part of the cell, making it less likely be recognized into two cells. The fourth column demonstrates the improvement of the detection capability of our proposed MSRNet. MSRNet effectively reduces the probability of missing small cells by using multiple regression. In particular, fine-grained point-regression branch is useful for the detection and segmentation of small cells. Fig. 8 shows the gaussian prediction masks and point prediction masks. As can be seen from the figure, these two branches are able to correctly predict the location of each cell centers’ position.