cellmap_flow.models.models_config

The model config classes: ModelConfig and the built-in model types.

This is their public import path. The docs import them from here, and plugins subclass ModelConfig, ScriptModelConfig and the rest by it. The classes live in cellmap_flow.models.configs, one module per type. Config, what a type’s _get_config returns, is here too, for plugins that define a type.

Classes

Config	What a model type's _get_config returns: the model, its geometry
ModelConfig
ScriptModelConfig
DaCapoModelConfig
FlyModelConfig	A network fly_organelles trained, from a training checkpoint, a .ts or a model.pt.
BioModelConfig	A BioImage Model Zoo model, run through bioimageio.core.
CellMapModelConfig	Configuration class for a CellmapModel.
FinetuneModelConfig	Configuration class for a finetuned model.
HuggingFaceModelConfig	Configuration class for a Hugging Face model.
CellposeModelConfig	Cellpose 4 run on each z slice of a chunk.

Module Contents

class cellmap_flow.models.models_config.Config(**kwargs)

What a model type’s _get_config returns: the model, its geometry and anything else its script or checkpoint defines, as attributes.

Built from keyword arguments (a script’s globals, for a script model), which to_dict() and get() read back. axes_names defaults to x, y, z.

axes_names

kwargs

to_dict()

Returns the configuration as a dictionary.

get(key: str, default: Any = None) → Any

Gets the value of a configuration key.

class cellmap_flow.models.models_config.ModelConfig

env = None

default_env = None

display_channel = None

finetunable = False

finetune_modes()

What this model can be finetuned with: (“lora”, “full”), (“full”,) or (). For the dashboard, so it must not build the network: a type whose networks differ (a compiled one takes no LoRA adapter) reads it from what describes the model. The trainer checks the module it gets the same way (finetune.trainable.finetune_modes).

trainable_model()

The torch module the finetune trainer trains, or None for config.model as it is (finetune.model_loading.load_trainable_model).

The trainer feeds it a float tensor (B, 1, Z, Y, X), normalized as the dashboard’s input normalization does, at the input voxel size and read shape, and wants (B, C, Z’, Y’, X’) at the write shape back. A type whose config.model is not such a module (a bioimage.io prediction pipeline, a Cellpose model object) returns one that is, sharing its parameters, so that what is trained is what is served.

serve_trained(config, module)

Make config serve module, the trained trainable_model().

config is this model’s own Config (the trainer’s live server) or a new one with its geometry (a finetuned model’s, FinetuneModelConfig). By default the module is the model, run by the inferencer’s own forward; a type that predicts through its own process_chunk sets that up here instead.

validate_model_shapes = True

check_shapes_on_warmup = False

property effective_env

The environment this model’s server and finetuning job run in, or None for this one: env, else default_env (envs.effective).

property config

check_output_shape(output_shape)

Raise ValueError if a forward’s output_shape (batch first) contradicts the declared write_shape, block_shape or output_channels.

property geometry: cellmap_flow.models.geometry.ModelGeometry

The config’s geometry as one ModelGeometry, read once per built config.

property chunk_output_axes: tuple[str, ...]

The axes of a processed chunk; by default geometry.DEFAULT_OUTPUT_AXES, (“c”, “z”, “y”, “x”).

Models can override by setting config.chunk_output_axes. The server and blockwise read it only for whether a chunk has a channel axis: one that has must have it first, and its spatial axes are the raw data’s, in the raw data’s order.

property output_dtype

Returns the output dtype of the model. Defaults to np.float32.

to_dict()

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

{"type": <registered name>, **the constructor arguments given}, leaving out None as the built-in types leave out an unset name or scale. The built-in types write their own; this one serves a plugin type that does not, so exporting its config and launching its server work without one.

property launch_entry: dict

The model entry a launched server rebuilds this exact config from.

to_dict(), less what only the pipeline builder shows (a Hugging Face repo’s downloaded metadata), under the name the registry has this class under (registry.cli_name_of), so that the server rebuilds this class and not a parent it inherited cli_name from. With the entry’s own env, if it has one, which serving.launch takes off again; the server’s environment is effective_env, which it reads from the config.

property command: str

cellmap_flow_server arguments that rebuild this exact config.

Generated from to_dict() so no constructor argument is left behind (hand-written commands dropped Fly’s input/output size, which then silently fell back to 178/56, and Bio’s required voxel size), with every token shell-quoted: bsub runs it through bash -c and local launches through shlex.split. The class is named as the server CLI registers it (registry.cli_name_of), so it rebuilds this class and not a parent it inherited cli_name from.

class cellmap_flow.models.models_config.ScriptModelConfig(script_path, name=None, scale=None)

cli_name = 'script'

finetunable = True

script_path

name = None

scale = None

to_dict()

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

class cellmap_flow.models.models_config.DaCapoModelConfig(run_name: str, iteration: int, name=None, scale=None)

cli_name = 'dacapo'

finetunable = True

run_name

iteration

name = None

scale = None

to_dict()

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

class cellmap_flow.models.models_config.FlyModelConfig(checkpoint_path: str, channels: list[str] = None, input_voxel_size: tuple = None, output_voxel_size: tuple = None, name: str = None, input_size=None, output_size=None, scale=None, sigmoid: bool = True)

A network fly_organelles trained, from a training checkpoint, a .ts or a model.pt.

Parameters:

• checkpoint_path – the checkpoint file (see the module docstring).

• channels – the name of each output channel. Default: read beside the checkpoint.

• input_voxel_size – nm per voxel (one number or one per axis). Default: read beside the checkpoint; one stands for the other.

• output_voxel_size – nm per voxel (one number or one per axis). Default: read beside the checkpoint; one stands for the other.

• input_size – voxels a side of a tile in and out. Default: read beside the checkpoint, else 178 in and the output computed from the network. An output size needs its input size.

• output_size – voxels a side of a tile in and out. Default: read beside the checkpoint, else 178 in and the output computed from the network. An output size needs its input size.

• sigmoid – pass the output through a sigmoid, added unless the network ends in one (default True); False serves the network’s output.

cli_name = 'fly'

finetunable = True

checkpoint_path = ''

name = None

scale = None

sigmoid

channels

input_voxel_size

output_voxel_size

input_size = None

output_size = None

property default_env

fly for a training checkpoint or an eager model.pt, else none (TorchScript runs anywhere).

A training checkpoint is loaded into fly_organelles’ StandardUnet, and unpickling an eager model.pt imports fly_organelles’ classes: pixi.toml’s fly environment has them, cellmap-flow’s own does not. A .ts file carries its own code and runs wherever torch does.

load_eval_model()

The network to serve, in eval mode: the checkpoint’s, with a sigmoid if sigmoid.

property model

to_dict()

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

With what was read beside the checkpoint, so the entry says what is served; the sizes when given or read, not when computed. In the constructor’s order, which exported YAMLs keep.

class cellmap_flow.models.models_config.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.

class cellmap_flow.models.models_config.CellMapModelConfig(folder_path, name=None, scale=None)

Configuration class for a CellmapModel.

cli_name = 'cellmap'

finetunable = True

cellmap_model

name = None

scale = None

to_dict()

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

class cellmap_flow.models.models_config.FinetuneModelConfig(lora_adapter_path: str = None, base_model: dict = None, name: str = None, scale=None, weights_path: str = None, input_voxel_size: tuple = None, output_voxel_size: tuple = None)

Configuration class for a finetuned model.

Wraps any base ModelConfig with the finetune’s weights, which are exactly one of: a LoRA adapter (lora_adapter_path), applied with PEFT, or a full finetune’s state dict (weights_path, from a run with --lora-r 0), loaded strictly. The base model is built by its own ModelConfig, as the trainable module the trainer trained, and the weights go on top of it.

cli_name = 'finetune'

finetunable = True

lora_adapter_path = None

weights_path = None

base_model_dict = None

name = None

scale = None

input_voxel_size = None

output_voxel_size = None

property env

the weights go on the base model, so they need the base’s packages.

Type:

Its own env if it was given one, else its base model’s

property default_env

The base model type’s default_env, as the weights need the base model’s packages. A type that decides per model decides for the base.

property base_model_config

Lazily build the base ModelConfig from the stored dict.

property display_channel

a finetune serves the same channels.

Type:

Its base’s

finetune_modes()

Its base’s: continuing a finetune trains the same network.

trainable_model()

The module it serves (its base’s trainable module with the adapter or weights on), to train on from where the finetune left off; not config.model, which a type serving through its own process_chunk (cellpose) keeps as its own model object.

to_dict()

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

Surfaces key base model fields at the top level so the pipeline builder UI can display them alongside the finetune-specific fields.

class cellmap_flow.models.models_config.HuggingFaceModelConfig(repo, revision=None, name=None, scale=None)

Configuration class for a Hugging Face model.

cli_name = 'huggingface'

finetunable = True

METADATA_RETRY_SECONDS = 60

repo

revision = None

name = None

scale = None

to_dict()

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

class cellmap_flow.models.models_config.CellposeModelConfig(voxel_size, pretrained_model: str = 'cpsam_v2', output: str = 'flows', slices_per_chunk: int = 8, slice_size: int = 512, context: int = 32, batch_size: int = DEFAULT_BATCH_SIZE, diameter: float = None, flow_threshold: float = 0.4, cellprob_threshold: float = 0.0, stitch_threshold: float = 0.0, name=None, scale=None)

Cellpose 4 run on each z slice of a chunk.

Parameters:

• voxel_size – nm per voxel, input and output (one number or one per axis). Cellpose-SAM sees objects best about 30 voxels across, so pick the scale at which yours are roughly that, or give diameter.

• pretrained_model – “cpsam_v2” (default), “cpsam”, “cpdino”, “cpdino-vitb”, or the path of finetuned Cellpose weights.

• output – “flows” (the default; float32: flow_y, flow_x and the cell probability), “probability” (float32, 0 to 1, one channel) or “masks” (uint64, ids unique within a chunk).

• slices_per_chunk – z slices in a chunk.

• slice_size – voxels a side, in y and x, of each chunk’s slices.

• context – voxels read on each side in y and x beyond those, and cut off.

• batch_size – tiles per GPU pass (DEFAULT_BATCH_SIZE; None, as a blank form field sends it, is that too).

• diameter – object diameter in input voxels; None keeps the model’s own scale (about 30). Cellpose resizes each slice by 30 / diameter.

• flow_threshold – Cellpose’s mask thresholds (masks only).

• cellprob_threshold – Cellpose’s mask thresholds (masks only).

• stitch_threshold – Cellpose’s own slice linking (masks only): a mask takes the id of the mask in the slice before that it overlaps by at least this IoU, within a chunk. 0, the default, leaves each slice’s masks apart.

cli_name = 'cellpose'

default_env = 'cellpose4'

finetunable = True

finetune_output_type = 'flows'

property display_channel

The cell probability’s channel when the output has several (flows): the layer’s flow colours are dimmed by it, and its range is the layer’s.

voxel_size

pretrained_model = ''

output = 'flows'

slices_per_chunk = 8

slice_size = 512

context = 32

batch_size = 16

diameter = None

flow_threshold

cellprob_threshold

stitch_threshold

name = None

scale = None

property read_size: int

Voxels a side, in y and x, of each slice Cellpose is given.

process_chunk(idi, output_roi)

output_roi segmented: (channels, z, y, x), the output’s channels.

finetune_modes()

(“lora”, “full”), without building the network: every Cellpose 4 network is a ViT whose Linear and Conv layers take adapters.

training_patch_voxels()

The patch the trainer reads, (input, output) in voxels: one slice of one tile.

Not the serving geometry. Cellpose-SAM’s ViT adds a fixed 32 x 32 grid of position embeddings to its 8-pixel tokens, so it takes 256 x 256 images and nothing else; Cellpose serves a slice by cutting it into such tiles, and trains on random tiles of that size too. The trainer does the same: it reads tiles, and supervises the whole of each, as Cellpose does (an instance cut by its edge is left out of the flow loss, instance_flows). One slice: the network is 2D and sees each slice alone, so more slices would cost a network pass each and add nothing, and a one-slice patch centred on a painted voxel always lands on paint.

trainable_model()

Cellpose’s network as the trainer trains it: (B, 1, Z, Y, X) -> (B, 3, Z, Y, X).

Each slice normalized as Cellpose’s eval does (1st to 99th percentile), then through the network, which gives flowY, flowX and the cell probability logit (instance_flows’ channels). The network is the one config.model segments with, so what is trained is what is served; it is put in float32 (Cellpose loads it in bfloat16), which serving then uses too. Its patch embedding is kept out of LoRA (lora_exclude_patterns): Cellpose’s forward reads that layer’s weight tensor directly, so an adapter on it would never be used.

serve_trained(config, module)

Serve module (the trained trainable_model()) through Cellpose’s eval.

Training changed the network in place, LoRA adapters and all, and Cellpose’s eval runs that network, so the chunks are segmented as before. config is this model’s own config (the trainer’s live server) or a finetuned model’s new one with this geometry, which gets what segmenting needs.

config.model stays Cellpose’s model object, as it is for the base model, so the inferencer neither moves nor forwards it. Were it the module, the warmup and the shape check would forward it, and a 3-channel tile network fails the 1-channel serving geometry: as a hard error whenever the read slice happens to be one tile. The module is kept as config.trained_module.

to_dict()

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