Recovering the Imperfect: Cell Segmentation in the Presence of Dynamically Localized Proteins

Deploying off-the-shelf segmentation networks on biomedical data has become common practice, yet if structures of interest in an image sequence are visible only temporarily, existing frame-by-frame methods fail. In this paper, we provide a solution to segmentation of imperfect data through time based on temporal propagation and uncertainty estimation. We integrate uncertainty estimation into Mask R-CNN network and propagate motion-corrected segmentation masks from frames with low uncertainty to those frames with high uncertainty to handle temporary loss of signal for segmentation. We demonstrate the value of this approach over frame-by-frame segmentation and regular temporal propagation on data from human embryonic kidney (HEK293T) cells transiently transfected with a fluorescent protein that moves in and out of the nucleus over time. The method presented here will empower microscopic experiments aimed at understanding molecular and cellular function.


Introduction
In the past decades, it has become evident that proteins are very dynamic and their localization within the cell dictates which function they perform [3]. Cell biologists carry out time-lapse fluorescence microscopy experiments to study protein localization and unravel how this affects the protein's function. Optogenetics can be used to control the localization of the protein of interest [18] to observe the effects of dynamic localization patterns. Take Figure 1 for an example: HEK293T cells expressing mCherry fused to the optogenetic tool LINuS [15] were observed over time to analyze the effect of blue light on the nuclear import of the protein. In the absence of blue light, the fusion protein localizes predominantly in the cytosol. Hence, the nucleus appears dark. However, in presence of blue light, the fusion protein enters the nucleus making it appear bright and difficult to distinguish from the cytosol. Giving light repeatedly creates oscillations of the protein in and out of the nucleus in time. Learning-based methods, such as U-Net [19] and Mask R-CNN [6], succeed at segmenting structures in data with clearly visible patterns, but fail when the visibility deteriorates. When the signal in the nucleus is similar to that in the cytosol, nuclei segmentation from any network is not reliable. While in a single image it is not possible to improve these segmentations, past and future frames in a video provide additional information for refinement. An expert can play the video back and forth to infer the segmentation of ambiguous nuclei. An automated segmentation method, too, must (1) automatically identify critical frames and (2) propagate predictions from neighbouring frames.
In this paper, we address both challenges by (1) equipping Mask R-CNN with uncertainty estimation to identify erroneous predictions and (2) incorporating optical flow to improve the identified erroneous predictions by propagating certain predictions from neighboring frames. Doing so, we introduce the most recent uncertainty estimation methods in biomedical instance segmentation and solve a real task commonly experienced in signalling studies which is not yet addressed. So far, nuclear markers have been employed in the experiments so that the available automated segmentation tools can be used [4]. However, additional markers cause unreliable quantification since different proteins bleed-through and interfere with each other. They also limit the channel space needed for other proteins of interest. The presented method makes the use of a nuclear marker dispensable.

Related Work
We are the first to address instance segmentation of structures with an oscillatory fluorescent signal in biomedical videos. Our work is related to methods that benefit from temporal features in their design. Milan et al. [14] and Payer et al. [17] used Recurrent Neural Network (RNN) to aggregate temporal features. Paul et al. [16] incorporated temporal cues for segmentation via optical flow. Similarly, Jain et al. [9] propagated features instead of segmentation with flow from key frames. A similar idea was applied to instance segmentation by Bertasius and Torresani [2]. Although these methods seem close, the task is different: (1) they heavily rely on dense annotation in time to learn interpolations explicitly while we cannot afford it due to high cost of expert annotation and difficulty in fine-grained annotation of imperfect frames and (2) they benchmark only on visible objects, while we are solely interested in objects with limited visibility.

Instance-Aware Segmentation with Uncertainty Estimation
We base our model on Mask R-CNN [6] equipped with elastic deformations of U-Net [19] to create additional biomedically plausible images on-the-fly for better generalization. We incorporate uncertainty estimation in our Mask R-CNN architecture to detect erroneous predictions; see Figure 2. We consider data uncertainty (aleatoric), model uncertainty (epistemic), and their combination.
Data uncertainty. For data uncertainty, we use the modified cross-entropy loss [10] in the mask branch of Mask R-CNN. This models the data uncertainty as the learned noise scale from the data. To learn both the class scores and their noise scale, the negative expected log-likelihood is minimized for pixel i as: where c is the correct class among all classes (c) andŝ are predicted logits corrupted by Gaussian noise with standard deviation σ. σ is also learned by the network alongside with the logits. Figure 2 shows where this loss is embedded in the network. At test time, the uncertainty is computed as the entropy of the class pseudo-probabilities for each pixel as: Model uncertainty. Model uncertainty requires sampling from the model. We experimented with the latest sampling strategies. Dropout is one of the common techniques to sample from networks [10]. Ensemble replaces the sampling

Algorithm 1: Uncertainty-Based Nuclei Propagation
Input: instance segmentation masks S, uncertainty maps U, motion per consecutive frame pairs M, hard threshold θ, relative thresholds: α,β Output: tracks for frame id f in order do by dropout by individually trained networks [11]. SGDR ensemble replaces the individual trained networks by the pre-converged snapshots of the same model [7]. Pre-converged models are obtained at the end of each SGDR (SGD with warm restarts) [12] cycle. WTA (Winner-Takes-All) [8] aims for training a single network with multiple heads where at each iteration only the head with the best prediction gets penalized. EWTA (Evolving WTA) [13] is a variant of WTA which improves the trade-off between diversity and consistency.

Uncertainty-Based Nuclei Propagation
We base the nuclei propagation on the cell tracker 4 by Ronneberger et al. [19]. This overlapping-based tracker is well suited for the data at hand, which does not show large motion over time. Despite its simplicity, it is one of the top performing cell trackers in the challenge. After computing the overlapping-based tracks, we go over all the frames of each track with ascending average nuclei uncertainty order (without nuclei, uncertainty is set to infinity) and update the identified uncertain predictions by propagating masks from more certain neighbours as depicted in Figure 1. A prediction is marked uncertain if its average uncertainty is higher than a threshold θ. Neighbours are more certain if they meet a relative threshold of α in case of a single-side propagation and β in case of a two-sided propagation as detailed in Algorithm 1. This resembles the procedure that experts follow to find non-visible nuclei by sliding over time.

Motion Estimation for Biomedical Videos
Propagation of predictions over time requires motion estimation between frames to warp the certain predictions onto the less certain ones. One simple way is to warp by the shift and scaling parameters computed between the not yet updated and neighbouring nuclei predictions. This approach assumes the shape of the nuclei does not change over time; however, slight deformations can occur. Optical flow can provide fine-grained motion. Recent optical flow methods for natural images are networks [5] trained on synthetic datasets. These methods perform well on real images, but their performance deteriorates as the gap between real and synthetic images grows. Our data is very different from existing synthetic datasets and no synthetic data exists for biomedical data. Therefore, with the prior knowledge that expected flow in our videos can be well explained by smooth deformations, and being able to generate them on the fly via deformation augmentations of U-Net, we explicitly train a network to predict these deformations as explored by Sokooti et al. [20]. We warp an image (I t+1 ) backward in time with randomly generated smooth deformations (f t→t+1 ) to obtain the previous image (I t ). Then we train a FlowNet on the on-the-fly-generated image pairs and ground-truth flows. Even though the fluorescent signal in the nuclei regularly disappears making the optical flow challenging to estimate, the motion of the cytosol helps infer the motion of the nuclei.

Experiments
Implementation details. We based our implementation on the Mask R-CNN by Abdulla [1]. We incorporated publicly available elastic deformations 5 for U-Net [19] as additional augmentation. We used ResNet50-FPN as backbone. We used softmax cross entropy for the mask head with 3 classes (background, cytosol and nuclei). We trained the networks from scratch for 12k epochs except the SGDR ensemble, which was trained for 15k with 3k cycles. We used the last pre-converged models. The WTA and EWTA models were trained for 11k and then merged with a network trained for 4k. We used 4 as ensemble size for all methods. Ensembling was performed over bounding boxes. For nuclei propagation, we used θ = 0.5, α = 0.7 and β = 0.85. For Mask R-CNN and FlowNet (variant C) training we used deformations with 3 and 10 control points, respectively and deformation magnitude of 10. Please see the original papers for more details. Our code and examplar dataset is publicly available 6 . Data and annotation. While the challenge we address is common in signalling studies, there is no public dataset for this purpose. Our data was generated after 24 hours of transient transfection of human embryonic kidney (HEK293T) cells with a construct expressing the fusion protein mCherry-LINuS with a ZEISS LSM780 confocal microscope. MaskRCNN was trained on 82 cropped images from original 1800x1800 images with pixel size 0.11x0.11 and 6 cells on average in each and tested on 2 full-sized images with 73 cells in total to validate instance segmentation and uncertainty estimation. We evaluated our propagation algorithm on an unseen video with 35 full-sized frames with randomly selected 117 cells annotated by experts. In cases where the nuclei appeared ambiguous experts interpolated the annotations from past and future frames.

Instance Segmentation and Uncertainty Estimation Evaluation
We evaluated our instance segmentation by mean average precision (mAP). In Table 1 we refer to standard mAP as mAP (sm) since it is based on softmax scores. To further evaluate the quality of the uncertainty estimation, we replaced the softmax scores by average entropy over the cell and re-computed mAP (mAP (ent)). This simulates our approach as we rely on averaged uncertainties in our nuclei propagation. This measure shows how reliable the predicted uncertainties are at ranking the prediction quality in the precision-recall curves, which are used commonly to evaluate uncertainty estimation for classification tasks. In Table 1 we see that WTA merged is the best at 0.5 IoU (Intersection over Union) threshold and at 0.75 threshold ensemble is the best. In the rest of our experiments we used the WTA merged with data uncertainty since it is computationally more efficient than ensemble.

Nuclei Propagation Evaluation
We report the mean IoU of the nuclei segmentation for all experiments in Table 2. We use interpolated for nuclei completely missed before our improvement, updated for the nuclei that were segmented but our method decided to improve, non-updated for the nuclei that were chosen not to be improved by our method and all for all the mentioned cases. The first row shows the results before any propagation algorithm. We present 3 variants of motion used in propagation (column: warped with). We also explored using the certainties as pixel-wise weights in fusing segmentations from candidate neighbors and computed the weighted average to find the final mask. To isolate the gain by our uncertainty-based error detection, we created a baseline which is identical to our mean-flow variant, but performs the propagation on all the nuclei over a track. The significant improvement obtained by all our variants, especially flow variants without fusion, shows that our method can effectively improve erroneous nuclei predictions. The baseline that propagates to all frames (all ) had a lower performance on the non-updated frames, showing that propagation independent of an uncertainty measure harms nuclei with high confidence (approx. 90% of all nuclei).

Slight Signal Loss
Since the network has been trained on a very clean training data, even small shifts in the signal distribution might disturb its performance and cause failures as in Figure 3. Notice that in the first frame, we have two connected components wrongly building together the nucleus mask. As each cell can only have a single nucleus, we consider propagating one which is the most certain. The remaining connected component can be filtered with a simple post-processing step. We see that our method improves upon the initial segmentation masks significantly.

Extreme Signal Loss
The scenario in Figure 4 shows the full potential of our method in the case of extreme signal loss. Uncertainty estimation allows to rank segmentation masks and infer more reasonable masks for the less certain frames. In the last frame the network misses the nucleus completely in a very confident way. Since we set the uncertainty in such cases to infinity to ensure that it gets updated (biologically there is certainly a nucleus in each cell), we are able to recover the missing nucleus. For a biologist taking measurements from the results of our method is more reliable compared to the non-updated results. Fig. 3: Example of a nucleus getting slightly less visible due to increased noise at time t. Columns correspond to time points t − 1, t and t + 1, from left to right respectively. In the first row we show the raw images. In the second row we overlay MaskRCNN outputs with the input image and report the average uncertainty. Pixel-wise uncertainty maps are presented in the third row while updated masks are overlaid to the original images in the last row.

No Signal Loss -Segmentation Failure
If the training data is not rich with all variations of the real task at hand, which is often the case in biomedical settings, networks can also fail at generalization. Our method is not only limited to refining failures due to signal loss, on the contrary, it is generic to all failures caused by the uncertainties inherited in the network or the data. Our predicted uncertainty estimations provide a good indicator for these scenarios and allows to recover them as we show in Figure 5.

Conclusion
We addressed automated segmentation of image sequences, which cannot be analyzed frame-by-frame due to temporary uncertainties causing errors in predictions. First, we estimate uncertainty from Mask R-CNN to identify unreliable predictions. Second, we improve the less reliable predictions by propagating Fig. 4: Example of a nucleus becoming completely indistinguishable from the surrounding cytosol in all time points. Columns correspond to time points t − 1, t, t + 1 and t + 2, from left to right respectively. In the first row we show the raw images. In the second row we overlay MaskRCNN outputs with the input image and report the average uncertainty. Pixel-wise uncertainty maps are presented in the third row while updated masks are overlaid to the original images in the last row.
the more certain ones from neighbouring frames. We evaluated our method on HEK293T cells expressing a protein that oscillates in and out of the nucleus over time making the nucleus invisible temporarily. Our method improves nuclei segmentation over several baselines while keeping the cytosol and background segmentation untouched. We believe that our method will facilitate further timeseries analysis for quantitative biology to understand the effect that the dynamic localization of the protein has on the cell without additional markers. Columns correspond to time points t − 1, t and t + 1, from left to right respectively. In the first row we show the raw images. In the second row we overlay MaskRCNN outputs with the input image and report the average uncertainty. Pixel-wise uncertainty maps are presented in the third row while updated masks are overlaid to the original images in the last row.