Learning continuous spatial transcriptomics vector fields for accurate, scalable, deep cell segmentation

This repository contains the code and supplementary materials for the paper:

"Learning continuous spatial transcriptomics vector fields for accurate, scalable, deep cell segmentation"

We introduce a novel framework for cell segmentation in spatial transcriptomics that models cellular structure as a continuous, locally conditioned vector field. This formulation enables accurate reconstruction of cells across both 2D and 3D datasets, while scaling efficiently to large molecule-resolved datasets.

Table of contents

• Background
• Method Overview
• Key Features
• Datasets
• Evaluation
• Results
• Installation
• Usage
• Pretrained Models
• Reproducibility
• Contributing

Background

Spatial transcriptomics enables measurement of gene expression with spatial resolution, but accurate cell segmentation remains a major challenge, especially in:

• High-density regions
• Irregular cell geometries
• 3D volumetric datasets

Most existing methods rely on clustering or global probabilistic models, which struggle to capture local spatial structure and continuous organization.

Method Overview

Our approach formulates segmentation as a continuous spatial field inference problem rather than discrete clustering.

Core ideas:

• Tiling + halo regions
  Large datasets are processed as overlapping spatial tiles with halo regions to preserve neighborhood context.

• Local neighborhood reasoning
  Each transcript is encoded using spatial position and gene identity, and interacts with nearby transcripts via:

  • k-nearest neighbors or radius search
  • distance-aware feature aggregation

• Continuous vector field prediction
  The model predicts:

  • Cell occupancy
  • Cell center likelihood
  • Offset vectors pointing to cell centers

• Implicit grouping via vector fields
  Transcripts belonging to the same cell follow consistent directional flows toward a shared center.

• Global reconstruction
  Local predictions are merged via overlap-based stitching into a coherent segmentation.

Key Features

• ✅ Continuous vector field formulation (no clustering required)
• ✅ Works on both 2D and 3D spatial transcriptomics
• ✅ Handles irregular geometries and variable densities
• ✅ Scalable via tiling and parallelization
• ✅ Reduces:
  • Molecular contamination
  • Spatial leakage
• ✅ Improves downstream biological analyses

Datasets

We evaluate across diverse datasets spanning:

• Multiple tissues (brain, prostate, lung, liver, colon, stomach and kidney)
• Multiple technologies (MERFISH, STARmap, Xenium, CosMx and Stereo-seq)
• Both 2D and 3D acquisitions

(See paper for full details)

Evaluation

We evaluate segmentation quality using multi-level metrics for 2D and 3D (ovrlpy package), among them:

Molecular consistency

• Spurious co-expression of mutually exclusive genes

Spatial consistency

• Boundary leakage
• Spatial contamination

Cellular coherence

• Gene expression entropy
• Transcript/gene counts per cell

(See paper for full details)

Results

🔬 Segmentation Performance

• Ranked #1
• Strong improvements in 3D and 2D datasets
• More accurate:
  • Cell counts
  • Transcript assignment
  • Cellular completeness

🧬 Reduced Contamination

• Lower spurious gene co-expression and 3D spatial contamination
• Reduced cross-cell transcript mixing

🧠 Better Biological Signal

• Improved clustering stability
• Better marker gene separability
• Stronger spatial organization

🧫 Downstream Insights

Accurate segmentation enables discovery of biological phenomena:

• Tumor immune exclusion patterns
• Spatial organization of stromal compartments
• Improved ligand–receptor interaction detection

⚡ Efficiency

• Faster inference than competing methods
• Fewer failed segmentation regions
• Scales to large datasets via parallelization

Pretrained Models

We provide pretrained models trained across multiple tissues.

Capabilities:

• ✅ Zero-shot generalization across datasets
• ✅ Strong performance without retraining
• ✅ Fine-tuning further improves results

Pretrained models:

• Maintain low contamination
• Preserve spatial coherence
• Generalize across tissues (e.g., → liver)

Reproducibility

• Consistent performance across biological replicates
• Low variance across metrics
• Stable across parameter settings

The method operates in a broad high-performance regime, reducing sensitivity to tuning.

Usage

All codes to run Rubicon and the evaluation framework are provided in this Github repository.

Pretrain model is also provided in this Github repository.

Contributing

Contributions to improve this repository are welcome! If you have suggestions or improvements, please open an issue or submit a pull request.