keejkrej/cellpose-cpsam-onnx

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 <15px

License

The Cellpose code is © 2020 Howard Hughes Medical Institute, licensed under the BSD 3-Clause License.

⚠️ The Cellpose-SAM model is trained on data licensed under CC-BY-NC. See the upstream notice.

Citation

If you use this model, please cite the Cellpose-SAM paper:

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}
}