keejkrej/cellpose-cpsam-onnx

license: bsd-3-clause license_name: bsd-3-clause-code-cc-by-nc-data library_name: onnx tags:

• cellpose
• cell-segmentation
• microscopy
• onnx
• biology pipeline_tag: image-segmentation

Cellpose cpsam — ONNX

ONNX export of the Cellpose cpsam (Cellpose-SAM) model for cell segmentation in microscopy images.

Model Details

• Architecture: ViT-L based Transformer (SAM backbone + Cellpose readout)
• Parameters: 304M
• Input: pixel_values — float32 tensor (1, 3, 256, 256) (one tile)
• Output: output — float32 tensor (1, 3, 256, 256) — channels are (dY, dX, cellprob)
• Opset: 14

The model weights are stored as external data (model.onnx.data) due to the 2GB protobuf limit. Both files must be in the same directory for inference.

Usage

Python (onnxruntime)

import onnxruntime as ort
import numpy as np

sess = ort.InferenceSession("model.onnx", providers=["CUDAExecutionProvider"])
tile = np.zeros((1, 3, 256, 256), dtype=np.float32)  # your preprocessed tile
output = sess.run(None, {"pixel_values": tile})[0]    # (1, 3, 256, 256)
dY, dX, cellprob = output[0, 0], output[0, 1], output[0, 2]

Rust (cellpose-rs)

use cellpose_rs::{CellposeSession, SegmentParams, preprocess};

let mut session = CellposeSession::new(Path::new("model.onnx"), false)?;
let chw = preprocess::build_chw_image(phase, fluo, h, w);
let masks = session.segment(&chw, h, w, SegmentParams::default())?;

See cellpose-rs for the full Rust inference library.

Post-processing

The raw output requires Cellpose flow dynamics post-processing:

1. Threshold cellprob > 0 to get foreground
2. Scale flows: dP = dP * foreground / 5.0
3. 200-step Euler integration with bilinear interpolation
4. Histogram-based seed detection (max-pool kernel=5, threshold=10)
5. Connected-component labelling of seeds
6. Assign each foreground pixel to its nearest seed
7. Remove masks >40% of image area, fill holes, remove masks Pachitariu, M., Rariden, M., & Stringer, C. (2025).

Cellpose-SAM: superhuman generalization for cellular segmentation. bioRxiv. https://doi.org/10.1101/2025.04.28.651001

@article{pachitariu2025cellposesam,
  title={Cellpose-SAM: superhuman generalization for cellular segmentation},
  author={Pachitariu, Marius and Rariden, Maia and Stringer, Carsen},
  journal={bioRxiv},
  year={2025},
  doi={10.1101/2025.04.28.651001}
}