keejkrej/spotiflow-onnx

Spotiflow ONNX — General

ONNX export of the general pretrained Spotiflow model for 2D fluorescent spot detection in microscopy images.

Model Details

Property	Value
Architecture	Spotiflow (U-Net backbone + stereographic flow head)
Input	float32 [B, 1, H, W] — single-channel grayscale image
Output 0	float32 [B, 1, H, W] — heatmap (pre-sigmoid logits)
Output 1	float32 [B, 3, H, W] — stereographic flow (z, y, x)
ONNX opset	16
Pretrained variant	general

Usage

Python (onnxruntime)

import onnxruntime as ort
import numpy as np

session = ort.InferenceSession("model.onnx")
image = np.random.rand(1, 1, 512, 512).astype(np.float32)
outputs = session.run(None, {"input": image})
heatmaps, flow = outputs[0], outputs[-1]  # flow has shape [B, 3, H, W]

Rust (spotiflow-rs)

use spotiflow_rs::{SpotiflowSession, PredictParams};

let mut session = SpotiflowSession::new("model.onnx")?;
let (spots, heatmaps, flows) = session.predict(&image_f32, h, w, PredictParams::default())?;

Export

Exported from the official PyTorch weights using a wrapper that converts the model's dictionary output to an explicit (heatmaps, flow) tuple for correct ONNX tracing:

python scripts/export_onnx.py --model general --output model.onnx

License

BSD 3-Clause — same as the original Spotiflow repository.

Copyright (c) 2023, Albert Dominguez Mantes, Martin Weigert.

Citation

@article{dominguez2024spotiflow,
  title={Spotiflow: accurate and efficient spot detection for fluorescence microscopy with deep stereographic flow regression},
  author={Dominguez Mantes, Albert and Herrera, Antonio and Khven, Irina and Schlaeppi, Anjalie and Aho, Eftychia and Erskine, Amber and Laubscher, Eleonora and Hendriks, Gert-Jan and Thiran, Jean-Philippe and Bhatt, Deepak K and Wegner, Joerg D and Weigert, Martin},
  journal={bioRxiv},
  year={2024},
  publisher={Cold Spring Harbor Laboratory}
}