ONNX export#

Export model prediction maps to ONNX. From a source checkout, run uv sync --locked --extra onnx to install the export packages. Install an ONNX runtime separately to run the file.

The file contains prediction maps. Convert them to classes and instance masks separately, using the same preprocessing as the checkpoint. Image height and width are fixed; batch size can vary.

Family

Input

Outputs, in order

CellPose

image: floating (B, 3, H, W)

flow_map: (B, 2, H, W); type_map: raw logits (B, classes, H, W)

Omnipose

image: floating (B, 3, H, W)

flow_map: (B, 2, H, W); type_map: raw logits (B, classes, H, W)

StarDist

image: floating (B, 3, H, W)

binary_map: foreground values (B, 1, H, W); ray_map: (B, n_rays, H, W); type_map: raw logits (B, classes, H, W)

HoverNet

image: floating (B, 3, H, W)

hv_map: horizontal/vertical distances (B, 2, H, W); type_map: raw logits (B, classes, H, W)

InstanSeg

image: floating (B, 3, H, W)

instanseg_map: reconstruction values (B, 4, H, W); type_map: raw logits (B, classes, H, W)

CellVit

image: floating (B, 3, H, W)

hv_map: horizontal/vertical distances (B, 2, H, W); type_map: raw logits (B, classes, H, W)

B is batch size, H height and W width. out_size can change the output dimensions. The exporter uses the network’s device and dtype, even inside autocast, and restores its training mode afterward. It creates output folders and may overwrite the destination file.

from pathlib import Path

import torch

from cellseg_models_pytorch.models.cellpose import CellPose

model = CellPose.from_pretrained(
    Path("checkpoints/cellpose.safetensors"), device=torch.device("cpu")
)
output_path = model.export_onnx(Path("exports/cellpose.onnx"))

CellPose and Omnipose#

Omnipose uses export_cellpose_onnx too. Build its network with the same architecture and class count as the training run, then load your weights:

from cellseg_models_pytorch.models.cellpose import export_cellpose_onnx
from cellseg_models_pytorch.models.cellpose.cellpose_unet import omnipose_nuclei

network = omnipose_nuclei(3, enc_name="resnet18", enc_pretrain=False)
network.load_state_dict(
    torch.load("checkpoints/omnipose.pt", map_location="cpu", weights_only=True)
)
omnipose_path = export_cellpose_onnx(network, Path("exports/omnipose.onnx"))
export_cellpose_onnx(model, output_path, input_shape=(1, 3, 256, 256), opset_version=18, dynamic_batch=True)#

Export CellPose or Omnipose prediction maps to ONNX.

The exported graph contains only the dense neural-network prediction stage. Flow integration and instance reconstruction remain in the existing Python post-processing pipeline.

Parameters:

  • model (CellPose | Module) – A CellPose instance, its underlying network, or an Omnipose network created with omnipose_nuclei. Load trained weights first.

  • output_path (str | Path) – Destination .onnx file. Parent folders are created; an existing file may be overwritten.

  • input_shape (Tuple[int, int, int, int]) – Example (batch, channels, height, width) input shape. Spatial dimensions are fixed in the exported graph. With dynamic batches, the tracing batch is at least two to avoid specializing batch one. Defaults to (1, 3, 256, 256).

  • opset_version (int) – ONNX opset version. Defaults to 18.

  • dynamic_batch (bool) – Mark the input batch dimension dynamic. Output batch dimensions inherit the same symbolic dimension through the exported graph. Defaults to True.

Returns:

Path to the exported ONNX model.

Raises:

  • ValueError – input_shape does not contain four positive integer dimensions.

  • TypeError – model is not the expected wrapper or a PyTorch module.

  • ImportError – onnx or onnxscript is unavailable.

  • RuntimeError – The PyTorch version is too old or export fails.

Return type:

Path

Note

Requires onnx and onnxscript, in addition to the package dependencies. The exporter checks for PyTorch 2.7 or later; the package’s own minimum PyTorch version still applies. Training/evaluation state of every network module is restored on success or failure. Example inputs use the network’s device and floating dtype; exporting does not move the model.

The input is named image; outputs are flow_map and type_map. Outputs have batch, channel, height and width axes. Cell-type outputs are raw logits. Preprocessing and Python reconstruction are not exported.

CellPose.export_onnx(output_path, input_shape=(1, 3, 256, 256), opset_version=18, dynamic_batch=True)#

Save this network’s dense prediction stage as an ONNX file.

Parameters:

  • output_path (str | Path) – Destination .onnx path. Parent folders are created and an existing file may be overwritten.

  • input_shape (Tuple[int, int, int, int]) – Example (batch, channels, height, width) shape. Spatial dimensions are fixed in the exported graph.

  • opset_version (int) – ONNX operator-set version supported by your deployment runtime.

  • dynamic_batch (bool) – Allow the batch dimension to vary at inference time.

Returns:

Path to the exported model. Python instance reconstruction is not included.

Return type:

Path

See also

export_cellpose_onnx for dependencies, output names and state restoration.

StarDist#

export_stardist_onnx(model, output_path, input_shape=(1, 3, 256, 256), opset_version=18, dynamic_batch=True)#

Export the StarDist neural-network forward pass to ONNX.

The exported graph contains only the dense neural-network prediction stage. StarDist polygon reconstruction and overlapping-object suppression remain in the existing Python post-processing pipeline.

Parameters:

  • model (StarDist | Module) – A high-level StarDist instance or its underlying .model module.

  • output_path (str | Path) – Destination .onnx file. Parent folders are created; an existing file may be overwritten.

  • input_shape (Tuple[int, int, int, int]) – Example (batch, channels, height, width) input shape. Spatial dimensions are fixed in the exported graph. With dynamic batches, the tracing batch is at least two to avoid specializing batch one. Defaults to (1, 3, 256, 256).

  • opset_version (int) – ONNX opset version. Defaults to 18.

  • dynamic_batch (bool) – Mark the input batch dimension dynamic. Output batch dimensions inherit the same symbolic dimension through the exported graph. Defaults to True.

Returns:

Path to the exported ONNX model.

Raises:

  • ValueError – input_shape does not contain four positive integer dimensions.

  • TypeError – model is not the expected wrapper or a PyTorch module.

  • ImportError – onnx or onnxscript is unavailable.

  • RuntimeError – The PyTorch version is too old or export fails.

Return type:

Path

Note

Requires onnx and onnxscript, in addition to the package dependencies. The exporter checks for PyTorch 2.7 or later; the package’s own minimum PyTorch version still applies. Training/evaluation state of every network module is restored on success or failure. Example inputs use the network’s device and floating dtype; exporting does not move the model.

The input is named image; outputs are binary_map, ray_map and type_map. Outputs have batch, channel, height and width axes. Cell-type outputs are raw logits. Preprocessing and Python reconstruction are not exported.

StarDist.export_onnx(output_path, input_shape=(1, 3, 256, 256), opset_version=18, dynamic_batch=True)#

Save this network’s dense prediction stage as an ONNX file.

Parameters:

  • output_path (str | Path) – Destination .onnx path. Parent folders are created and an existing file may be overwritten.

  • input_shape (Tuple[int, int, int, int]) – Example (batch, channels, height, width) shape. Spatial dimensions are fixed in the exported graph.

  • opset_version (int) – ONNX operator-set version supported by your deployment runtime.

  • dynamic_batch (bool) – Allow the batch dimension to vary at inference time.

Returns:

Path to the exported model. Python instance reconstruction is not included.

Return type:

Path

See also

export_stardist_onnx for dependencies, output names and state restoration.

HoverNet#

Load your trained checkpoint, then export its prediction maps:

from cellseg_models_pytorch.models.hovernet import HoverNet

hovernet = HoverNet.from_pretrained(
    Path("checkpoints/hovernet.pt"), device=torch.device("cpu")
)
hovernet_path = hovernet.export_onnx(Path("exports/hovernet.onnx"))
export_hovernet_onnx(model, output_path, input_shape=(1, 3, 256, 256), opset_version=18, dynamic_batch=True)#

Save HoverNet prediction maps as an ONNX file.

Parameters:

  • model (HoverNet | Module) – A trained HoverNet instance or its PyTorch network.

  • output_path (str | Path) – Destination .onnx file. Parent folders are created; an existing file may be overwritten.

  • input_shape (tuple[int, int, int, int]) – Example (batch, channels, height, width) shape. Image size stays fixed; a dynamic export traces with at least two images.

  • opset_version (int) – ONNX operator-set version, default 18.

  • dynamic_batch (bool) – Allow the number of images in a batch to vary.

Returns:

Path to the exported file. Input image is floating BCHW. Outputs hv_map (horizontal/vertical distance maps) and type_map (class logits) use the same axes. Preprocessing and instance reconstruction stay in Python.

Raises:

  • ValueError – Input dimensions are not four positive integers.

  • TypeError – The model is not a HoverNet instance or a PyTorch module.

  • ImportError – Export requires the optional onnx and onnxscript packages.

  • RuntimeError – The PyTorch version is unsupported or export fails.

Return type:

Path

Note

Export uses the network’s device and floating dtype. Every module’s training mode is restored after success or failure.

HoverNet.export_onnx(output_path, input_shape=(1, 3, 256, 256), opset_version=18, dynamic_batch=True)#

Save prediction maps for use in an ONNX runtime.

Parameters:

  • output_path (str | Path) – Destination .onnx file; an existing file may be overwritten.

  • input_shape (tuple[int, int, int, int]) – Example (batch, channels, height, width) shape. Image height and width stay fixed in the exported file.

  • opset_version (int) – ONNX operator-set version, default 18.

  • dynamic_batch (bool) – Allow the number of images in a batch to vary.

Returns:

Exported file path. Instance reconstruction remains in Python.

Return type:

Path

See also

export_hovernet_onnx for output names and optional dependencies.

InstanSeg#

The four reconstruction channels contain two coordinate values, a scale value and a seed value. They are raw network outputs. InstanSeg’s pixel classifier stays in Python; keep its trained weights with the model used for reconstruction.

from cellseg_models_pytorch.models.instanseg import InstanSeg

instanseg = InstanSeg.from_pretrained(
    Path("checkpoints/instanseg.pt"), device=torch.device("cpu")
)
instanseg_path = instanseg.export_onnx(Path("exports/instanseg.onnx"))
export_instanseg_onnx(model, output_path, input_shape=(1, 3, 256, 256), opset_version=18, dynamic_batch=True)#

Save InstanSeg prediction maps as an ONNX file.

Parameters:

  • model (InstanSeg | Module) – A trained InstanSeg instance or its PyTorch network.

  • output_path (str | Path) – Destination .onnx file. Parent folders are created; an existing file may be overwritten.

  • input_shape (tuple[int, int, int, int]) – Example (batch, channels, height, width) shape. Image size stays fixed; a dynamic export traces with at least two images.

  • opset_version (int) – ONNX operator-set version, default 18.

  • dynamic_batch (bool) – Allow the number of images in a batch to vary.

Returns:

Path to the exported file. Input image is floating BCHW. Outputs instanseg_map (four reconstruction channels) and type_map (class logits) use the same axes. Preprocessing and instance reconstruction stay in Python.

Raises:

  • ValueError – Input dimensions are not four positive integers.

  • TypeError – The model is not a InstanSeg instance or a PyTorch module.

  • ImportError – Export requires the optional onnx and onnxscript packages.

  • RuntimeError – The PyTorch version is unsupported or export fails.

Return type:

Path

Note

Export uses the network’s device and floating dtype. Every module’s training mode is restored after success or failure. The attached pixel classifier is not exported; keep its trained weights with the Python model used for InstanSeg reconstruction.

InstanSeg.export_onnx(output_path, input_shape=(1, 3, 256, 256), opset_version=18, dynamic_batch=True)#

Save prediction maps for use in an ONNX runtime.

Parameters:

  • output_path (str | Path) – Destination .onnx file; an existing file may be overwritten.

  • input_shape (tuple[int, int, int, int]) – Example (batch, channels, height, width) shape. Image height and width stay fixed in the exported file.

  • opset_version (int) – ONNX operator-set version, default 18.

  • dynamic_batch (bool) – Allow the number of images in a batch to vary.

Returns:

Exported file path. Instance reconstruction remains in Python.

Return type:

Path

See also

export_instanseg_onnx for output names and optional dependencies.

CellVit#

Use the SAM encoder and image size from training. Export is tested with samvit_base_patch16_224 at 224 × 224 pixels. Other SAM encoders and image sizes need separate runtime checks.

from cellseg_models_pytorch.models.cellvit import CellVit

cellvit = CellVit.from_pretrained(
    Path("checkpoints/cellvit.pt"), device=torch.device("cpu")
)
cellvit_path = cellvit.export_onnx(
    Path("exports/cellvit.onnx"), input_shape=(1, 3, 224, 224)
)
export_cellvit_onnx(model, output_path, input_shape=(1, 3, 224, 224), opset_version=18, dynamic_batch=True)#

Save CellVit prediction maps as an ONNX file.

Parameters:

  • model (CellVit | Module) – A trained CellVit instance or its PyTorch network.

  • output_path (str | Path) – Destination .onnx file. Parent folders are created; an existing file may be overwritten.

  • input_shape (tuple[int, int, int, int]) – Example (batch, channels, height, width) shape. Tested with samvit_base_patch16_224 at 224 × 224 pixels. Image size stays fixed; a dynamic export traces with at least two images.

  • opset_version (int) – ONNX operator-set version, default 18.

  • dynamic_batch (bool) – Allow the number of images in a batch to vary.

Returns:

Path to the exported file. Input image is floating BCHW. Outputs hv_map (horizontal/vertical distance maps) and type_map (class logits) use the same axes. Preprocessing and instance reconstruction stay in Python.

Raises:

  • ValueError – Input dimensions are not four positive integers.

  • TypeError – The model is not a CellVit instance or a PyTorch module.

  • ImportError – Export requires the optional onnx and onnxscript packages.

  • RuntimeError – The PyTorch version is unsupported or export fails.

Return type:

Path

Note

Export uses the network’s device and floating dtype. Every module’s training mode and SAM attention settings are restored afterward, even if export fails.

CellVit.export_onnx(output_path, input_shape=(1, 3, 224, 224), opset_version=18, dynamic_batch=True)#

Save prediction maps for use in an ONNX runtime.

Parameters:

  • output_path (str | Path) – Destination .onnx file; an existing file may be overwritten.

  • input_shape (tuple[int, int, int, int]) – Example (batch, channels, height, width) shape. Image height and width stay fixed in the exported file.

  • opset_version (int) – ONNX operator-set version, default 18.

  • dynamic_batch (bool) – Allow the number of images in a batch to vary.

Returns:

Exported file path. Instance reconstruction remains in Python.

Return type:

Path

See also

export_cellvit_onnx for output names and optional dependencies.