[Paper Review] VoxelEmbed: 3D Instance Segmentation and Tracking with Voxel Embedding based Deep Learning

Mengyang Zhao, Quan Liu+7 more|arXiv (Cornell University)|Jun 22, 2021

Cell Image Analysis TechniquesBiochemistry, Genetics and Molecular Biology25 references17 citations

TL;DR

This paper proposes VoxelEmbed, a novel 3D instance segmentation and tracking method using voxel embedding to capture spatial-temporal context from 3D volumetric videos. It employs a multi-stream learning framework with 3D synchronization to enable end-to-end, memory-efficient inference on a single 12GB GPU, achieving state-of-the-art performance on four ISBI Cell Tracking Challenge datasets, including competitive results on sparsely annotated data.

ABSTRACT

Recent advances in bioimaging have provided scientists a superior high spatial-temporal resolution to observe dynamics of living cells as 3D volumetric videos. Unfortunately, the 3D biomedical video analysis is lagging, impeded by resource insensitive human curation using off-the-shelf 3D analytic tools. Herein, biologists often need to discard a considerable amount of rich 3D spatial information by compromising on 2D analysis via maximum intensity projection. Recently, pixel embedding-based cell instance segmentation and tracking provided a neat and generalizable computing paradigm for understanding cellular dynamics. In this work, we propose a novel spatial-temporal voxel-embedding (VoxelEmbed) based learning method to perform simultaneous cell instance segmenting and tracking on 3D volumetric video sequences. Our contribution is in four-fold: (1) The proposed voxel embedding generalizes the pixel embedding with 3D context information; (2) Present a simple multi-stream learning approach that allows effective spatial-temporal embedding; (3) Accomplished an end-to-end framework for one-stage 3D cell instance segmentation and tracking without heavy parameter tuning; (4) The proposed 3D quantification is memory efficient via a single GPU with 12 GB memory. We evaluate our VoxelEmbed method on four 3D datasets (with different cell types) from the ISBI Cell Tracking Challenge. The proposed VoxelEmbed method achieved consistent superior overall performance (OP) on two densely annotated datasets. The performance is also competitive on two sparsely annotated cohorts with 20.6% and 2% of data-set having segmentation annotations. The results demonstrate that the VoxelEmbed method is a generalizable and memory-efficient solution.

Motivation & Objective

• Address the challenge of large-scale 3D biomedical video analysis, which is hindered by resource-intensive human curation and the loss of 3D spatial information when projecting to 2D.
• Overcome limitations of existing 2D-based pixel embedding methods by generalizing to 3D voxel embedding to preserve volumetric context.
• Develop a memory-efficient, end-to-end framework for simultaneous 3D instance segmentation and tracking without heavy hyperparameter tuning.
• Enable robust performance on both densely and sparsely annotated 3D datasets, reflecting real-world biological imaging constraints.

Proposed method

• Propose a voxel embedding strategy that generalizes 2D pixel embedding to 3D by incorporating volumetric context across spatial and temporal dimensions.
• Design a multi-stream learning architecture using 2D convolutional networks to efficiently learn spatial-temporal embeddings across 3D volumes.
• Introduce a 3D synchronization algorithm inspired by slice propagation to reconstruct full volumetric masks from 2D predictions.
• Use cosine embedding loss to ensure consistent and discriminative embeddings for individual cells across frames.
• Implement end-to-end training with a single-stage framework that jointly optimizes segmentation and tracking via shared embedding space.
• Leverage mean-shift clustering with a bandwidth of 0.1 to group voxels into instance-level segments based on learned embeddings.

Experimental results

Research questions

• RQ1 Can voxel embedding generalize 2D pixel embedding to 3D to improve instance segmentation and tracking in volumetric time-lapse microscopy?
• RQ2 Can a multi-stream 2D network design effectively learn 3D spatial-temporal context without requiring full 3D convolutional networks?
• RQ3 Can the proposed method achieve high performance on both densely and sparsely annotated 3D datasets with minimal hyperparameter tuning?
• RQ4 Is the method memory-efficient enough to run on a single GPU with 12 GB VRAM while maintaining competitive accuracy?

Key findings

• VoxelEmbed achieved the highest overall performance (OP) of 0.876 on the Fluo-N3DH-SIM+ dataset and 0.926 on Fluo-C3DH-A549-SIM, outperforming baseline methods including RSHN and KTH-SE.
• On the Fluo-N3DH-CHO dataset with 20.6% frame annotation, VoxelEmbed achieved OP of 0.910, ranking first in tracking (TRA: 0.958) and second in segmentation (SEG: 0.862).
• On the Fluo-N3DH-CE dataset with only 2% frame annotation, VoxelEmbed achieved OP of 0.807, ranking second in overall performance and first in tracking (TRA: 0.897).
• The method demonstrated strong generalization, maintaining competitive performance across diverse cell types and annotation sparsity levels.
• VoxelEmbed achieved these results using only a single 12GB GPU, confirming its memory efficiency and practicality for real-world biomedical imaging pipelines.
• The ablation study (VoxelEmbed-D) showed that the multi-stream design contributes to performance, with OP of 0.879 on Fluo-N3DH-SIM+ and 0.919 on Fluo-C3DH-A549-SIM, indicating robustness and consistency.