3D Cell Oversegmentation Correction via Geo-Wasserstein Divergence

TL;DR

The paper tackles the persistent problem of oversegmentation in 3D cell segmentation by reframing it as a supervised binary classification task that leverages geometry-aware (Geo-Wasserstein divergence) and topology-aware features to distinguish true gaps from cell splits. A lightweight pretrained classifier, combined with cross-layer interpolation, identifies and corrects oversegmented regions, with an extension to tilted cutting surfaces for non-axis-aligned cases. Extensive experiments on plant and animal datasets demonstrate improved segmentation quality and transferability, and ablations confirm the added value of Geo-Wasserstein features and statistical descriptors. The approach offers a practical, dataset-agnostic pipeline that can be adopted by end-users to enhance diverse 3D segmentation tasks and potentially other temporally/planes-based segmentation problems.

Abstract

3D cell segmentation methods are often hindered by \emph{oversegmentation}, where a single cell is incorrectly split into multiple fragments. This degrades the final segmentation quality and is notoriously difficult to resolve, as oversegmentation errors often resemble natural gaps between adjacent cells. Our work makes two key contributions. First, for 3D cell segmentation, we are the first work to formulate oversegmentation as a concrete problem and propose a geometric framework to identify and correct these errors. Our approach builds a pre-trained classifier using both 2D geometric and 3D topological features extracted from flawed 3D segmentation results. Second, we introduce a novel metric, Geo-Wasserstein divergence, to quantify changes in 2D geometries. This captures the evolving trends of cell mask shape in a geometry-aware manner. We validate our method through extensive experiments on in-domain plant datasets, including both synthesized and real oversegmented cases, as well as on out-of-domain animal datasets to demonstrate transfer learning performance. An ablation study further highlights the contribution of the Geo-Wasserstein divergence. A clear pipeline is provided for end-users to build pre-trained models to any labeled dataset.

Figures (15)

• Figure 1: Examples of oversegmentation errors from 2D-based model (CellStitch) and 3D-based model (PlantSeg).
• Figure 2: Main pipeline for the oversegmentation correction. The framework extracts layer-to-layer EMDs (Algorithm \ref{['a2']}) and 3D shape information (Algorithm \ref{['a3']}), combining these features for binary classification to distinguish oversegmentations from natural gaps.
• Figure 3: An example of 2D mis-segmentation from the raw image stack that leads to the subsequent oversegmentation in 3D results.
• Figure 4: Comparison of 3D topological shape between over segmented cases and natural gap cases.
• Figure 5: Main pipeline for constructing rotated 3D segmentation layers for the tilted cases from 3D-based method (e.g., PlantSeg).