cellmap_flow.models.configs.bio

BioModelConfig: a BioImage Model Zoo model, run through bioimageio.core.

model is any source bioimageio.core loads: a zoo id or nickname (“conscientious-dromedary”), a DOI or URL, or the path of a model’s rdf.yaml (bioimageio.yaml) or packaged .zip. Its description (the RDF, read as format 0.5; a 0.4 one is converted) gives the rest:

  • What a chunk is. A 3D model (space axes z, y, x) is given one tile a chunk. A 2D model (y, x) is given slices_per_chunk z slices a chunk, each segmented on its own: in one call when its batch axis takes any size, a call a slice when the RDF fixes the batch at 1 or has no batch axis.

  • The tile’s size. An axis of fixed size keeps it; a parameterized one (min + n * step) gets the smallest size it takes of at least input_size: 256 a side for a 2D model and 128 for a 3D one by default.

  • The context: the halo of the model’s output, cut off on each side so that chunks meet where the model saw both sides of the seam. context sets it; the zoo’s EM U-Nets give none.

  • The voxel size: the input’s space axes’ scale and unit, in nm, unless voxel_size is given. No zoo model gives a unit, so the zoo’s EM models read at the voxel size they were trained at, from cellmap-flow’s own table (bioimage_voxel_sizes.yaml); any other needs voxel_size.

  • The output: every output tensor’s channels, one tensor after another, as the model’s own postprocessing leaves them (a sigmoid, if it has one; nothing clipped or rescaled), served as float32 unless the RDF says the output is integers (labels), which keep an integer type.

  • The weights: whichever bioimageio.core prefers of those the model has, or weight_format.

The model’s preprocessing (normalization) and postprocessing run in bioimageio.core’s prediction pipeline, which is built with the config, once, and reused for every chunk.

What it cannot run: a model with more than one required input, an input with more than one channel (RGB), or an output that is not a map over the input’s space, such as micro-SAM’s, whose masks have an object axis of data-dependent size and need prompts.

It can be finetuned (trainable_model) when it has torch weights (a state dict, which takes LoRA too, or TorchScript): the network is rebuilt from them, with the input’s normalization and the output’s sigmoid around it, and once trained serves through the inferencer’s forward instead of the pipeline.

bioimageio 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 bioimageio environment unless the entry names another env.

Attributes

logger

WEIGHT_FORMATS

TRAINABLE_WEIGHT_FORMATS

DEFAULT_INPUT_SIZE

DEFAULT_SLICES_PER_CHUNK

NM_PER_UNIT

Classes

BioModelConfig

A BioImage Model Zoo model, run through bioimageio.core.

Functions

served_dtype(→ numpy.dtype)

What chunks of outputs (output tensor descriptions) are served as.

Module Contents

cellmap_flow.models.configs.bio.logger
cellmap_flow.models.configs.bio.WEIGHT_FORMATS = ('pytorch_state_dict', 'torchscript', 'onnx', 'tensorflow_saved_model_bundle', 'keras_hdf5', 'keras_v3')
cellmap_flow.models.configs.bio.TRAINABLE_WEIGHT_FORMATS
cellmap_flow.models.configs.bio.DEFAULT_INPUT_SIZE
cellmap_flow.models.configs.bio.DEFAULT_SLICES_PER_CHUNK = 8
cellmap_flow.models.configs.bio.NM_PER_UNIT
cellmap_flow.models.configs.bio.served_dtype(outputs) → numpy.dtype

What chunks of outputs (output tensor descriptions) are served as.

float32 unless every output is declared as integers: the model’s own postprocessing (a sigmoid, labels from a watershed) has run by then, so the RDF’s dtype is what the values are. bool is served as uint8 and int64 as uint64, which neuroglancer shows; integer outputs are ids or classes, which are not negative.

class cellmap_flow.models.configs.bio.BioModelConfig(model: str, voxel_size=None, input_size=None, context=None, slices_per_chunk: int = None, weight_format: str = None, name=None, scale=None)

A BioImage Model Zoo model, run through bioimageio.core.

Parameters:

  • model – a zoo id or nickname (“conscientious-dromedary”), a DOI or URL, or the path of a model’s rdf.yaml (bioimageio.yaml) or packaged .zip.

  • voxel_size – nm per input voxel (one number, or z, y, x): the level the model reads. By default the RDF’s, from its input’s space axes’ scale and unit; for a model whose RDF gives no unit (all of the zoo’s), the one it was trained at when cellmap-flow’s table knows it (bioimage_catalog.trained_at: the zoo’s EM models); else it is needed. A 2D model’s z is the spacing of its slices, its y by default.

  • input_size – voxels a side of the tile the model is given, one number or one per space axis of the model (z, y, x; y, x for a 2D model). An axis the RDF fixes keeps its size; a parameterized one gets the smallest size it takes of at least this. By default 256 for a 2D model and 128 for a 3D one.

  • context – voxels cut off each side of the model’s output (one number, or one per space axis of the model), so chunks meet where the model saw both sides of the seam; the input is read that much larger. By default the RDF’s halo, 0 where it gives none.

  • slices_per_chunk – z slices in a chunk of a 2D model; 8 by default. A 3D model ignores it.

  • weight_format – the weights to run (“torchscript”, “pytorch_state_dict”, “onnx”, …); by default bioimageio.core’s pick of those the model has.

cli_name = 'bioimage'
default_env = 'bioimageio'
finetunable = True
model = ''
voxel_size = None
input_size = None
context = None
slices_per_chunk = None
weight_format = None
name = None
scale = None
process_chunk(idi: cellmap_flow.image_data_interface.ImageDataInterface, output_roi)

output_roi predicted: (channels, z, y, x), in output_dtype.

trainable_model()

The model as one torch module, for the finetune trainer: (B, 1, Z, Y, X) of the read shape in, (B, channels, Z’, Y’, X’) of the write shape out, as process_chunk serves it.

Built from the model’s description, the way bioimageio.core runs it: the input’s preprocessing (its normalization), the network from its torch weights (a 2D one on each z slice, the slices in one call when its batch axis takes any size), each output’s postprocessing (a sigmoid) and the same crop, the outputs’ channels one after another. Statistics are taken per patch, over the axes the RDF names, so a patch is normalized as a chunk is; an RDF that names none asks bioimageio.core for statistics of all it has seen, which a patch cannot reproduce, and gets the patch’s own.

Refuses, with a ValueError saying why, a model whose weights are not torch (ONNX, TensorFlow, Keras: no torch module to train), or whose output goes through something no gradient passes (labels, a threshold, Stardist or Cellpose instances). On the device the server uses, in float32 and in eval mode.

serve_trained(config, module)

Serve module (the trained trainable_model()) through the inferencer’s own forward.

process_chunk runs the pipeline’s weights, as loaded from the zoo, so it is switched off: the module does all process_chunk did (the normalization, each slice of a 2D model, the sigmoid, the crop), given (1, 1, Z, Y, X) of the read shape, and gives float32.

finetune_modes()

How the dashboard can finetune this model: (“lora”, “full”), (“full”,) or ().

From its weight formats only, without bioimageio or torch: the catalog’s cached entry, else a local rdf.yaml or .zip read as YAML, else (where bioimageio.core is) its description. Torch weights decide, as trainable_model picks them: a state dict takes LoRA and a full finetune, TorchScript a full one only, and a model without either, or whose weights cannot be told, none.

to_dict()

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

What the model’s description decides unless told otherwise (its voxel size, tile and context) is written only when it was given, so that an exported entry keeps following the model.