Skip to content
BranchesTags

Repository files navigation

cellpose-rs

Rust translation of Cellpose — a generalist deep learning model for cell and nucleus segmentation.

  • 2026-06-28: vendored candle with 3d support
  • 2026-06-27: each file passes 2 audits. real data testing still needed though
  • 2026-06-25: Proper audit is ongoing

This is an LLM-mediated faithful (hopefully) translation, not the original code!

Most users should probably first see if the existing original code works for them, unless they have reason otherwise. The original source may have newer features and it has had more love in terms of fixing bugs. In fact, we aim to replicate bugs if they are present, for the sake of reproducibility! (but then we might have added a few more in the process)

There are however cases when you might prefer this Rust version. We generally agree with this page but more specifically:

  • We have had many issues with ensuring that our software works using existing containers (Docker, PodMan, Singularity). One size does not fit all and it eats our resources trying to keep up with every way of delivering software
  • Common package managers do not work well. It was great when we had a few Linux distributions with stable procedures, but now there are just too many ecosystems (Homebrew, Conda). Conda has an NP-complete resolver which does not scale. Homebrew is only so-stable. And our dependencies in Python still break. These can no longer be considered professional serious options. Meanwhile, Cargo enables multiple versions of packages to be available, even within the same program(!)
  • The future is the web. We deploy software in the web browser, and until now that has meant Javascript. This is a language where even the == operator is broken. Typescript is one step up, but a game changer is the ability to compile Rust code into webassembly, enabling performance and sharing of code with the backend. Translating code to Rust enables new ways of deployment and running code in the browser has especial benefits for science - researchers do not have deep pockets to run servers, so pushing compute to the user enables deployment that otherwise would be impossible
  • Old CLI-based utilities are bad for the environment(!). A large amount of compute resources are spent creating and communicating via small files, which we can bypass by using code as libraries. Even better, we can avoid frequent reloading of databases by hoisting this stage, with up to 100x speedups in some cases. Less compute means faster compute and less electricity wasted
  • LLM-mediated translations may actually be safer to use than the original code. This article shows that running the same code on different operating systems can give somewhat different answers. This is a gap that Rust+Cargo can reduce. Typesafe interfaces also reduce coding mistakes and error handling, as opposed to typical command-line scripting

But:

  • This approach should still be considered experimental. The LLM technology is immature and has sharp corners. But there are opportunities to reap, and the genie is not going back to the bottle. This translation is as much aimed to learn how to improve the technology and get feedback on the results.
  • Translations are not endorsed by the original authors unless otherwise noted. Do not send bug reports to the original developers. Use our Github issues page instead.
  • Do not trust the benchmarks on this page. They are used to help evaluate the translation. If you want improved performance, you generally have to use this code as a library, and use the additional tricks it offers. We generally accept performance losses in order to reduce our dependency issues
  • Check the original Github pages for information about the package. This README is kept sparse on purpose. It is not meant to be the primary source of information.
  • Translation parity is tracked file-by-file in TOAUDIT.md. A file is only marked done after two clean audits in a row.

Features

  • Cell and nucleus segmentation in 2D and 3D images
  • ViT-SAM (Cellpose-SAM) model architecture via candle
  • TIFF and PNG image I/O (including multi-page TIFF stacks)
  • Image normalization (percentile-based, tile-based, lowhigh)
  • Tiled inference for arbitrarily large images
  • Flow field computation and Euler-integration mask reconstruction
  • Mask post-processing (hole filling, small mask removal, 3D stitching)
  • CLI with the same interface as the Python version
  • Usable as a Rust library (cellpose crate)

Compatibility note: save_mpl writes the Cellpose four-panel summary image using the same panel data path as Python, but byte-identical matplotlib titles, axes, spacing, and layout are out of scope for this translation.

Not included: GUI, Python bindings, denoise/deblur restoration models, BioImage.IO export helpers, and distributed segmentation helpers.

Installation

Requires a current stable Rust toolchain. This repository does not currently declare a minimum supported Rust version.

# Clone the repository
git clone https://github.com/your-org/cellpose-rs.git
cd cellpose-rs

# CPU-only build
cargo build --release

# CUDA GPU build (requires CUDA toolkit 12+)
cargo build --release --features cuda

# CUDA 12.8 runtime used for local Quadro RTX 5000 validation
export CUDA_ROOT=/usr/local/cuda-12.8
export CUDA_COMPUTE_CAP=75
export PATH="$CUDA_ROOT/bin:$PATH"
export LD_LIBRARY_PATH="$CUDA_ROOT/lib64:${LD_LIBRARY_PATH:-}"

# Apple Metal GPU build (macOS only)
cargo build --release --features metal

Model weights

The default CP-SAM model is cached at ~/.cellpose/models/cpsam, matching Python's cpsam model name. The cache directory can be overridden with CELLPOSE_LOCAL_MODELS_PATH.

To prepare the default model cache:

cellpose --download_model

The CLI and CellposeModel::new_default(...) also download the default model on demand if it is missing.

Usage

As a library

Basic segmentation

use std::path::Path;
use cellpose::CellposeModel;
use cellpose::models::EvalParams;
use cellpose::core::Device;
use cellpose::io;

// Load the pretrained model
let model_path = Path::new("/home/user/.cellpose/models/cpsam");
let model = CellposeModel::new(model_path, Device::Cpu, None).unwrap();

// Read an image
let img = io::imread(Path::new("cells.png")).unwrap();

// Run segmentation with default parameters
let params = EvalParams::default();
let output = model.eval(&img, &params).unwrap();

// output.masks: labeled segmentation (0 = background, 1+ = cell IDs)
// output.flows: flow field (2, H, W)
// output.cellprob: cell probability map (H, W)
let num_cells = output.masks.iter().copied().max().unwrap_or(0);
println!("Found {} cells", num_cells);

Custom parameters

use cellpose::models::EvalParams;

let params = EvalParams {
    diameter: Some(45.0),          // expected cell diameter in pixels
    flow_threshold: 0.4,           // flow error threshold (higher = more permissive)
    cellprob_threshold: 0.0,       // cell probability threshold
    min_size: 15,                  // remove masks smaller than this (pixels)
    ..Default::default()
};

let output = model.eval(&img, &params).unwrap();

Flows only (skip mask computation)

let params = EvalParams {
    compute_masks: false,
    ..Default::default()
};

let output = model.eval(&img, &params).unwrap();
// output.flows and output.cellprob are computed; output.masks is all zeros

Saving results

use cellpose::io::{self, SaveMasksOptions};
use std::path::Path;

// Save masks as PNG
let masks_f32 = output.masks.mapv(|v| v as f32);
let options = SaveMasksOptions {
    png: true,
    tif: false,
    suffix: "_cp_masks",
    ..Default::default()
};
io::save_masks(
    Some(&img),
    &masks_f32,
    &[],
    Path::new("cells.png"),  // base filename
    &options,
).unwrap();
// Creates "cells_cp_masks.png"

Image I/O and transforms

use cellpose::io;
use cellpose::transforms::{self, NormalizeParams};
use std::path::Path;

// Read an image (supports PNG, JPEG, TIFF including multi-page stacks)
let img = io::imread(Path::new("image.tif")).unwrap();
println!("Shape: {:?}", img.shape()); // e.g. [512, 512, 3]

// Convert to standard format (H, W, 3)
let img = transforms::convert_image(&img, None, None, false).unwrap();

// Normalize (percentile-based, default 1st-99th)
let params = NormalizeParams::default();
let normalized = transforms::normalize_img(&img, &params).unwrap();

// Discover all image files in a directory
let files = io::get_image_files(
    Path::new("/path/to/images"),
    "_masks",    // exclude files ending with this
    None,        // no filename filter
    false,       // don't search subdirectories
).unwrap();

// Save an image
io::imsave(Path::new("output.tif"), &img).unwrap();

Batch processing

let files = io::get_image_files(Path::new("data/"), "_masks", None, false).unwrap();
let params = EvalParams::default();

for path in &files {
    let img = io::imread(path).unwrap();
    let output = model.eval(&img, &params).unwrap();
    let n = output.masks.iter().copied().max().unwrap_or(0);
    println!("{}: {} cells", path.display(), n);

    let masks_f32 = output.masks.mapv(|v| v as f32);
    let options = SaveMasksOptions {
        png: true,
        tif: false,
        suffix: "_cp_masks",
        ..Default::default()
    };
    io::save_masks(Some(&img), &masks_f32, &[], path, &options).unwrap();
}

3D segmentation

let stack = io::imread(Path::new("zstack.tif")).unwrap(); // [Z, H, W] or [Z, H, W, C]

let params = EvalParams {
    diameter: Some(30.0),
    stitch_threshold: 0.5,  // IoU threshold for stitching masks across z-slices
    ..Default::default()
};

let output = model.eval_3d(&stack, &params).unwrap();
// output.masks shape: [Z, H, W]
// output.flows shape: [3, Z, H, W] — (dZ, dY, dX)

Command line

# Segment all images in a directory
cellpose --dir /path/to/images --save_png

# Segment a single image with a specific diameter
cellpose --image_path image.tif --save_tif --diameter 30

# 3D segmentation
cellpose --dir /path/to/stacks --do_3d --save_tif

# Adjust thresholds
cellpose --dir /path/to/images --flow_threshold 0.4 --cellprob_threshold 0.0 --min_size 15

# Use GPU (requires CUDA build and matching CUDA runtime libraries)
cellpose --dir /path/to/images --use_gpu --backend auto --save_png

--backend auto and --backend candle currently select the Candle backend. --backend burn selects the feature-gated Burn prototype when the crate is built with an executable Burn backend feature. The same selector is available from Rust as cellpose::InferenceBackend.

The Burn prototype is available for backend development:

cargo test --features burn-ndarray burn_backend
cargo test -p cellpose-cli --features burn backend_accepts_burn_for_future_feature_gating

This verifies Burn weight staging, tensor conversion, and the executable ndarray prototype. burn_backend::BurnCpsamModel implements the shared InferenceNetwork contract, and Burn model construction uses the same backend-neutral CP-SAM checkpoint index as the Candle loader.

Additional compile-checked Burn backend feature sets are available:

cargo check --features burn-cuda --all-targets
cargo check --features burn-wgpu --all-targets
cargo check --features burn-ndarray --all-targets

On the Quadro RTX 5000 / CUDA 12.8 test machine, the real-checkpoint Burn CUDA prototype runs but is much slower than Candle CUDA, so Candle remains the recommended GPU backend for normal inference.

Project structure

Cargo.toml           # Workspace root + library crate (cellpose)
src/
  lib.rs             # Public API re-exports
  error.rs           # Error types
  transforms.rs      # Image normalization, tiling, conversion
  io.rs              # Image I/O (TIFF, PNG), Cellpose result files, ROI output
  core.rs            # Device management, tiled network execution
  dynamics.rs        # Flow diffusion, Euler integration, mask computation
  network.rs         # ViT-SAM neural network (candle)
  models.rs          # CellposeModel API
  utils.rs           # Mask post-processing, hole filling, stitching
  metrics.rs         # Evaluation metrics
  train.rs           # Training pipeline
  plot.rs            # Summary figures, outlines, mask overlays
cellpose-cli/        # CLI binary
  Cargo.toml
  src/main.rs
tests/
  integration_demo_images.rs   # I/O + transforms on real images
  integration_inference.rs     # End-to-end inference tests
orig-cellpose/       # Original Python cellpose (reference)

Benchmarks

Current inference benchmark for demodata/demo_images/img00.png, CP-SAM, single image, --diameter 30, Quadro RTX 5000:

Implementation Cells Inference Peak RSS
Python Cellpose / PyTorch CUDA 222 1.45s warm 2,057,472 KB
Rust / patched Candle CUDA 222 2.58s timed / 2.63s warm throughput 1,471,540 KB

On img00.png-img04.png, Rust/Candle CUDA preserves the Python cell counts exactly and currently measures 9.36s throughput-only vs 5.09s for Python/PyTorch CUDA (1.84x). The remaining broad-batch gap is mainly SAM encoder tile throughput, with flow-error QC diffusion the largest remaining postprocessing cost.

Inference time includes tiling, normalization, network forward pass, flow dynamics, and mask computation. CUDA builds must load the same CUDA runtime libraries they were built against; on the validation machine this required LD_LIBRARY_PATH=/usr/local/cuda-12.8/lib64:$LD_LIBRARY_PATH. If CUDA is requested directly through the Rust API and cannot initialize, model creation now returns the CUDA initialization error instead of silently falling back to CPU.

This repository carries a local [patch.crates-io] copy of candle-transformers 0.8.4 with Candle SAM's absolute positional embedding add changed from same-shape + to broadcast_add, enabling batched encoder execution for CP-SAM tiles.

On the validated Quadro RTX 5000 (sm75), Rust uses F32 execution even when Python-compatible bfloat16 mode is requested because BF16 is not supported on that GPU generation.

For benchmark throughput runs without internal profiling synchronization:

make bench-inference-rust BENCH_RUST_EXTRA_ARGS=--no-timing

For the current five-image Cellpose parity set:

make bench-inference-rust BENCH_IMAGE_LIST=demodata/demo5_images.txt BENCH_RUNS=1
make bench-inference-python BENCH_IMAGE_LIST=demodata/demo5_images.txt BENCH_RUNS=1

To compare saved Rust and Python benchmark JSON files against the current speed/RSS/parity gates:

make compare-inference-benchmarks

For mask-overlap parity, export masks from each implementation and compare with explicit IoU gates:

make bench-inference-demo5-mask-parity

The target above expands to:

make bench-inference-rust BENCH_IMAGE_LIST=demodata/demo5_images.txt BENCH_RUNS=1 \
  BENCH_MASK_OUT_DIR=benchmark-results/masks/rust-demo5 \
  BENCH_RUST_JSON=benchmark-results/rust-candle-demo5-mask-parity.json
make bench-inference-python BENCH_IMAGE_LIST=demodata/demo5_images.txt BENCH_RUNS=1 \
  BENCH_MASK_OUT_DIR=benchmark-results/masks/python-demo5 \
  BENCH_PYTHON_JSON=benchmark-results/python-pytorch-demo5-mask-parity.json
make compare-inference-benchmarks \
  BENCH_RUST_JSON=benchmark-results/rust-candle-demo5-mask-parity.json \
  BENCH_PYTHON_JSON=benchmark-results/python-pytorch-demo5-mask-parity.json \
  BENCH_MIN_FOREGROUND_IOU=1.0 BENCH_MIN_MEAN_BEST_IOU=1.0

For local CUDA validation, including actual CUDA model placement, the sm75 F32 dtype policy, and the img00.png 222-cell baseline:

make test-cuda-inference

Roadmap

  • Phase 1: Image I/O, transforms, normalization, tiling
  • Phase 2: Neural network inference (ViT-SAM in candle), flow dynamics, mask computation
  • Phase 3: Full CLI (PNG/TIFF/flows/outlines/TXT output, progress bar, GPU support)
  • Phase 4: Training pipeline
  • Continue file-by-file translation audits in TOAUDIT.md

Citation

If you use this software, please cite the original Cellpose papers:

Pachitariu, M., Rariden, M., & Stringer, C. (2025). Cellpose-SAM: superhuman generalization for cellular segmentation. bioRxiv.

Stringer, C., Wang, T., Michaelos, M., & Pachitariu, M. (2021). Cellpose: a generalist algorithm for cellular segmentation. Nature Methods, 18(1), 100-106.

License

BSD-3-Clause

About

No description, website, or topics provided.

Resources

Readme
Activity
Custom properties

Stars

0 stars

Watchers

1 watching

Forks

0 forks
Report repository

Releases

Packages

Contributors

Languages