cellmap_flow.models.configs.onnx

OnnxModelConfig: an ONNX export as the browser mode runs it, served with ONNX Runtime.

The model is one of:

path: a local .onnx file. A model finetuned in the browser (the file its Download finetuned model saves) carries its manifest in the file’s metadata_props, under cellmap_flow.manifest (MANIFEST_KEY), so the file alone says its voxel sizes, block sizes and channels. Any other file says none of that: the entry gives input_voxel_size, and the block sizes and the channel count are read from the graph’s input and output shapes when they are fixed, else from input_size, output_size and channels. A relative path in a YAML is taken from the YAML file’s folder (config_relative_paths), so an exported YAML can name the downloaded file beside it.
folder: one model of a Hugging Face repo of web exports, repo (cellmap/onnx-web-models by default) at revision: a folder with a manifest.json and one .onnx per variant, a precision and a block size. file names the variant’s file (what the browser’s Export Config writes: the one it ran); without it, variant (fp32 by default) and input_size (by default the largest block) choose one. Both files are downloaded with huggingface_hub, as the huggingface type’s are.

A manifest’s fields are web/src/models/manifest.ts’s, which parse_manifest reads the same way: the voxel sizes in nm, the channels, output (what the output is: the exports’ “raw logits”), input_norm (cellmap-flow steps), io (the input and output tensors’ names) and the variants. Anything the entry gives (voxel sizes, sizes, channels, output) wins over the manifest.

Each forward is what the browser feeds (web/src/infer/stages.ts): one block of float32 [1, 1, z, y, x] at the input size, normalized by the chain, in and [1, C, z', y', x'] out ([1, z', y', x'] is one channel); an fp16 export whose graph takes float16 is given float16, and its output is served as float32. A 2D student (input and output one slice thick, such as setup_53_2d) is given one slice a call, slices_per_chunk slices a chunk, as the bioimage type runs a 2D model. The session is wrapped in a torch module (_onnx_module), so that the inference runner serves it as it serves any network: the warmup forward checks the declared shapes against the graph’s. What its output is, the runner reads off the file’s graph (output_graph, read_graph): an export’s last Conv says logits, so the dashboard suggests a SigmoidPostprocessor (or the affinity chain for affinities), as for any model whose output is logits. The price is one more copy each way: the runner puts a block on the device, and the module hands it to ONNX Runtime from the host (148³ float32 is 13 MB, milliseconds a block next to the forward’s tenths of a second).

The manifest’s input_norm is what the browser feeds the model through; here, as for every type, the session’s chain (a YAML’s json_data) does the normalizing. declared_input_norm gives it to the dashboard’s model advice, which offers it with an Apply button, and the browser’s Export Config writes the chain the page ran with.

Not finetunable: the finetune trainer trains a torch module, with LoRA’s adapters or in full, and an ONNX graph is not one; its weights are tensors in a protobuf, and no backward pass runs through ONNX Runtime’s inference sessions. The browser’s own finetune trains an export’s last layer, and its result is served here.

onnxruntime is imported only when the model is built: the CLIs and the dashboard’s model form import every type. The server runs in pixi.toml’s onnx environment (onnxruntime-gpu on Linux, onnxruntime on the Mac) unless the entry names another env. The manifest is read without it (read_metadata, read_manifest), so cellmap_flow add and the dashboard’s Add a model can say what a file is.

Attributes

`logger`
`MANIFEST_KEY`
`VARIANTS`
`DEFAULT_VARIANT`
`DEFAULT_SLICES_PER_CHUNK`
`MANIFEST_RETRY_SECONDS`
`SCALAR_TENSOR_BYTES`

Classes

OnnxModelConfig

An ONNX export, the browser mode's kind of model, served with ONNX Runtime.

Functions

`read_metadata`(path) → Dict[str, str]	The `metadata_props` of the ONNX model at `path` (ModelProto field 14), by key.
`read_graph`(path)	The graph of the ONNX model at `path`, as `portable.activation`
`read_manifest`(path) → Optional[dict]	The cellmap-flow manifest the ONNX file at `path` carries (MANIFEST_KEY), parsed; None without one.
`parse_manifest`(raw[, where]) → dict	`raw`, a manifest's JSON, checked and widened as the browser's `parseManifest` does.
`choose_variant`(manifest[, file, variant, input_size]) → dict	The manifest's variant to serve: the one in `file`, else of precision

Module Contents

cellmap_flow.models.configs.onnx.logger

cellmap_flow.models.configs.onnx.MANIFEST_KEY = 'cellmap_flow.manifest'

cellmap_flow.models.configs.onnx.VARIANTS = ('fp32', 'fp16', 'int8w', 'int8w_fp16')

cellmap_flow.models.configs.onnx.DEFAULT_VARIANT = 'fp32'

cellmap_flow.models.configs.onnx.DEFAULT_SLICES_PER_CHUNK = 8

cellmap_flow.models.configs.onnx.MANIFEST_RETRY_SECONDS = 60

cellmap_flow.models.configs.onnx.read_metadata(path) → Dict[str, str]

The metadata_props of the ONNX model at path (ModelProto field 14), by key.

Read without onnx or onnxruntime, so the dashboard’s environment can read it: each of the ModelProto’s fields is stepped over (_fields) but the metadata’s entries; the graph (field 7), nearly all of the file, is never read.

Raises:: ValueError – the file is not a protobuf message with a graph, so not an ONNX model (a failed download’s HTML page, say).

cellmap_flow.models.configs.onnx.SCALAR_TENSOR_BYTES = 256

cellmap_flow.models.configs.onnx.read_graph(path)

The graph of the ONNX model at path, as portable.activation reads it: (nodes, outputs), each node {"op" (its op_type), "inputs", "outputs"}, a Clip’s with min and max where they are constants, and the graph’s output names.

Read without onnx, as read_metadata reads the metadata (the onnx environment has onnxruntime, not onnx): the ModelProto’s graph (field 7) is a GraphProto, whose nodes (field 1) and outputs (field 12, a ValueInfoProto named by its field 1) are read, and whose initializers (field 5), the weights, are stepped over unless one is a single number (SCALAR_TENSOR_BYTES), as a Clip’s bounds are.

Raises:: ValueError – the file is not an ONNX model.

cellmap_flow.models.configs.onnx.read_manifest(path) → dict | None

The cellmap-flow manifest the ONNX file at path carries (MANIFEST_KEY), parsed; None without one.

Raises:: ValueError – the file is not an ONNX model, or its manifest is not one.

cellmap_flow.models.configs.onnx.parse_manifest(raw, where='manifest.json') → dict

raw, a manifest’s JSON, checked and widened as the browser’s parseManifest does.

The exports write a variant’s block as one number (a cube) and io as a sentence; here they are three numbers and {"input", "output"}. A variant of a precision not in VARIANTS is left out.

Raises:: ValueError – a field is missing or of the wrong kind; it says which.

cellmap_flow.models.configs.onnx.choose_variant(manifest, file=None, variant=None, input_size=None) → dict

The manifest’s variant to serve: the one in file, else of precision variant (fp32 by default, else the first there is) with an input of input_size voxels, by default the largest block.

Raises:: ValueError – no variant is the one asked for; the message lists them.

class cellmap_flow.models.configs.onnx.OnnxModelConfig(path: str = None, folder: str = None, repo: str = None, revision: str = None, file: str = None, variant: str = None, input_voxel_size: tuple = None, output_voxel_size: tuple = None, input_size: tuple = None, output_size: tuple = None, channels: list[str] = None, output: str = None, slices_per_chunk: int = None, name: str = None, scale=None)

An ONNX export, the browser mode’s kind of model, served with ONNX Runtime.

Parameters:

path – a local .onnx file. Relative in a YAML: from the YAML’s folder.
folder – a model’s folder in repo (setup_54), instead of path.
repo – the Hugging Face repo of folder; cellmap/onnx-web-models.
revision – the repo’s branch, tag or commit; its main branch.
file – the .onnx of folder to serve: one of its manifest’s variants (what the browser’s Export Config writes).
variant – the precision to serve when file is not given: fp32 (the default), fp16, int8w or int8w_fp16.
input_voxel_size – nm per input voxel (one number, or z, y, x). The manifest’s; needed for a file without one.
output_voxel_size – nm per output voxel: by default the input’s times the manifest’s ratio of the two, else the input’s.
input_size – voxels of the block the model reads (one number or z, y, x): the manifest’s variant’s, else the graph’s input. For a folder, it chooses the variant (the largest by default).
output_size – voxels it writes: the variant’s, else the graph’s output.
channels – the output channels’ names; the manifest’s, else channel_0, channel_1, … for as many as the graph outputs.
output – what the output is, as the manifests say it (“raw logits”).
slices_per_chunk – z slices a chunk of a 2D student (a model whose input and output are one slice thick), each one call; 8 by default. A 3D model ignores it.

cli_name = 'onnx'

default_env = 'onnx'

finetunable = False

config_relative_paths = ('path',)

path = None

folder = None

repo = None

revision = None

file = None

variant = None

input_voxel_size = None

output_voxel_size = None

input_size = None

output_size = None

channels = None

output = None

slices_per_chunk = None

name = None

scale = None

property hub_repo: str

declared_input_norm(data_range=None) → list | None

The input normalization the model’s manifest declares, as cellmap-flow steps ([{"name": ..., **params}]), or None. The manifest gives each step’s values, so data_range is not used.

For the dashboard’s model advice (suggest_input_norm’s declared), which asks on every poll (_declared_manifest).

declared_geometry() → dict | None: ModelConfig.declared_geometry from the entry, the manifest and its variant, as the server works it out (_geometry) but without the graph: None when they leave a size to the graph, or the manifest cannot be read now (_declared_manifest).

output_graph(): The served file’s graph (read_graph), read back from the output it is fetched by; the manifest’s output words are not consulted, as the graph says it for certain.

to_dict()

This config as a model entry, which registry.build_model rebuilds.

What the manifest or the graph decides unless told otherwise is written only when it was given, so that an exported entry keeps following the model.