Coverage for src/bioimageio/core/digest_spec.py: 83%
247 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-08 15:59 +0000
1from __future__ import annotations
3import collections.abc
4import importlib.util
5import sys
6from itertools import chain
7from pathlib import Path
8from tempfile import TemporaryDirectory
9from typing import (
10 Any,
11 Callable,
12 Dict,
13 Iterable,
14 List,
15 Literal,
16 Mapping,
17 NamedTuple,
18 Optional,
19 Sequence,
20 Tuple,
21 Union,
22)
23from zipfile import ZipFile, is_zipfile
25import numpy as np
26import xarray as xr
27from loguru import logger
28from numpy.typing import NDArray
29from typing_extensions import TypeAlias, Unpack, assert_never
31from bioimageio.spec._internal.io import HashKwargs, PermissiveFileSource
32from bioimageio.spec.common import FileDescr, FileSource
33from bioimageio.spec.model import AnyModelDescr, v0_4, v0_5
34from bioimageio.spec.model.v0_4 import CallableFromDepencency, CallableFromFile
35from bioimageio.spec.model.v0_5 import (
36 ArchitectureFromFileDescr,
37 ArchitectureFromLibraryDescr,
38 ParameterizedSize_N,
39)
40from bioimageio.spec.utils import load_array
42from .axis import AxisId, AxisInfo, AxisLike, PerAxis
43from .block_meta import split_multiple_shapes_into_blocks
44from .common import Halo, MemberId, PerMember, SampleId, TotalNumberOfBlocks
45from .io import load_tensor
46from .sample import (
47 LinearSampleAxisTransform,
48 Sample,
49 SampleBlockMeta,
50 SampleBlockWithOrigin,
51 sample_block_meta_generator,
52)
53from .stat_measures import Stat
54from .tensor import Tensor
56TensorSource: TypeAlias = Union[
57 Tensor, xr.DataArray, NDArray[Any], PermissiveFileSource
58]
61def import_callable(
62 node: Union[
63 ArchitectureFromFileDescr,
64 ArchitectureFromLibraryDescr,
65 CallableFromDepencency,
66 CallableFromFile,
67 v0_5.CustomProcessingDescr,
68 ],
69 /,
70 **kwargs: Unpack[HashKwargs],
71) -> Callable[..., Any]:
72 """import a callable (e.g. a torch.nn.Module) from a spec node describing it"""
73 if isinstance(node, CallableFromDepencency):
74 module = importlib.import_module(node.module_name)
75 c = getattr(module, str(node.callable_name))
76 elif isinstance(node, ArchitectureFromLibraryDescr):
77 module = importlib.import_module(node.import_from)
78 c = getattr(module, str(node.callable))
79 elif isinstance(node, CallableFromFile):
80 c = _import_from_file_impl(node.source_file, str(node.callable_name), **kwargs)
81 elif isinstance(node, (ArchitectureFromFileDescr, v0_5.CustomProcessingDescr)):
82 c = _import_from_file_impl(node.source, str(node.callable), sha256=node.sha256)
83 else:
84 assert_never(node)
86 if not callable(c):
87 raise ValueError(f"{node} (imported: {c}) is not callable")
89 return c
92tmp_dirs_in_use: List[TemporaryDirectory[str]] = []
93"""keep global reference to temporary directories created during import to delay cleanup"""
96def _import_from_file_impl(
97 source: FileSource, callable_name: str, **kwargs: Unpack[HashKwargs]
98):
99 src_descr = FileDescr(source=source, **kwargs)
100 # ensure sha is valid even if perform_io_checks=False
101 # or the source has changed since last sha computation
102 src_descr.validate_sha256(force_recompute=True)
103 assert src_descr.sha256 is not None
104 source_sha = src_descr.sha256
106 reader = src_descr.get_reader()
107 # make sure we have unique module name
108 module_name = f"{reader.original_file_name.split('.')[0]}_{source_sha}"
110 # make sure we have a unique and valid module name
111 if not module_name.isidentifier():
112 module_name = f"custom_module_{source_sha}"
113 assert module_name.isidentifier(), module_name
115 source_bytes = reader.read()
117 module = sys.modules.get(module_name)
118 if module is None:
119 try:
120 td_kwargs: Dict[str, Any] = (
121 dict(ignore_cleanup_errors=True) if sys.version_info >= (3, 10) else {}
122 )
123 if sys.version_info >= (3, 12):
124 td_kwargs["delete"] = False
126 tmp_dir = TemporaryDirectory(**td_kwargs)
127 # keep global ref to tmp_dir to delay cleanup until program exit
128 # TODO: remove for py >= 3.12, when delete=False works
129 tmp_dirs_in_use.append(tmp_dir)
131 module_path = Path(tmp_dir.name) / module_name
132 if reader.original_file_name.endswith(".zip") or is_zipfile(reader):
133 module_path.mkdir()
134 ZipFile(reader).extractall(path=module_path)
135 else:
136 module_path = module_path.with_suffix(".py")
137 _ = module_path.write_bytes(source_bytes)
139 importlib_spec = importlib.util.spec_from_file_location(
140 module_name, str(module_path)
141 )
143 if importlib_spec is None:
144 raise ImportError(f"Failed to import {source}")
146 module = importlib.util.module_from_spec(importlib_spec)
148 sys.modules[module_name] = module # cache this module
150 assert importlib_spec.loader is not None
151 importlib_spec.loader.exec_module(module)
153 except Exception as e:
154 if module_name in sys.modules:
155 del sys.modules[module_name]
157 raise ImportError(f"Failed to import {source}") from e
159 try:
160 callable_attr = getattr(module, callable_name)
161 except AttributeError as e:
162 raise AttributeError(
163 f"Imported custom module from {source} has no `{callable_name}` attribute."
164 ) from e
165 except Exception as e:
166 raise AttributeError(
167 f"Failed to access `{callable_name}` attribute from custom module imported from {source} ."
168 ) from e
170 else:
171 return callable_attr
174def get_axes_infos(
175 io_descr: Union[
176 v0_4.InputTensorDescr,
177 v0_4.OutputTensorDescr,
178 v0_5.InputTensorDescr,
179 v0_5.OutputTensorDescr,
180 ],
181) -> List[AxisInfo]:
182 """get a unified, simplified axis representation from spec axes"""
183 return [AxisInfo.create(a) for a in io_descr.axes]
186def get_member_id(
187 tensor_description: Union[
188 v0_4.InputTensorDescr,
189 v0_4.OutputTensorDescr,
190 v0_5.InputTensorDescr,
191 v0_5.OutputTensorDescr,
192 ],
193) -> MemberId:
194 """get the normalized tensor ID, usable as a sample member ID"""
196 if isinstance(tensor_description, (v0_4.InputTensorDescr, v0_4.OutputTensorDescr)):
197 return MemberId(tensor_description.name)
198 elif isinstance(
199 tensor_description, (v0_5.InputTensorDescr, v0_5.OutputTensorDescr)
200 ):
201 return tensor_description.id
202 else:
203 assert_never(tensor_description)
206def get_member_ids(
207 tensor_descriptions: Iterable[
208 Union[
209 v0_4.InputTensorDescr,
210 v0_4.OutputTensorDescr,
211 v0_5.InputTensorDescr,
212 v0_5.OutputTensorDescr,
213 ]
214 ],
215) -> List[MemberId]:
216 """get normalized tensor IDs to be used as sample member IDs"""
217 return [get_member_id(descr) for descr in tensor_descriptions]
220def get_test_input_sample(model: AnyModelDescr) -> Sample:
221 if isinstance(model, v0_4.ModelDescr):
222 info = {
223 MemberId(d.name): (d, t) for d, t in zip(model.inputs, model.test_inputs)
224 }
225 else:
226 info = {d.id: d for d in model.inputs}
228 return _get_test_sample(info)
231get_test_inputs = get_test_input_sample
232"""DEPRECATED: use `get_test_input_sample` instead"""
235def get_test_output_sample(model: AnyModelDescr) -> Sample:
236 """returns a model's test output sample"""
237 if isinstance(model, v0_4.ModelDescr):
238 info = {
239 MemberId(d.name): (d, t) for d, t in zip(model.outputs, model.test_outputs)
240 }
241 else:
242 info = {d.id: d for d in model.outputs}
244 return _get_test_sample(info)
247get_test_outputs = get_test_output_sample
248"""DEPRECATED: use `get_test_output_sample` instead"""
251def _get_test_sample(
252 info: Union[
253 Mapping[MemberId, Union[v0_5.InputTensorDescr, v0_5.OutputTensorDescr]],
254 Mapping[
255 MemberId,
256 Tuple[
257 v0_4.InputTensorDescr,
258 FileSource,
259 ],
260 ],
261 Mapping[
262 MemberId,
263 Tuple[
264 v0_4.OutputTensorDescr,
265 FileSource,
266 ],
267 ],
268 ],
269) -> Sample:
270 arrays: Dict[MemberId, NDArray[Any]] = {}
271 for m, src in info.items():
272 if isinstance(src, tuple):
273 arrays[m] = load_array(src[1])
274 elif isinstance(src, (v0_5.InputTensorDescr, v0_5.OutputTensorDescr)):
275 if src.test_tensor is None:
276 raise ValueError(
277 f"Model input '{m}' has no test tensor defined, cannot create test sample."
278 )
279 arrays[m] = load_array(src.test_tensor)
280 else:
281 assert_never(src)
283 axes = {
284 m: get_axes_infos(t[0] if isinstance(t, tuple) else t) for m, t in info.items()
285 }
286 return Sample(
287 members={m: Tensor.from_numpy(arrays[m], dims=axes[m]) for m in info},
288 stat={},
289 id="test-sample",
290 )
293class IO_SampleBlockMeta(NamedTuple):
294 input: SampleBlockMeta
295 output: SampleBlockMeta
298def get_input_halo(
299 model: v0_5.ModelDescr, output_halo: Optional[PerMember[PerAxis[Halo]]] = None
300):
301 """returns which halo input tensors need to be divided into blocks with, such that
302 `output_halo` can be cropped from their outputs without introducing gaps."""
303 input_halo: Dict[MemberId, Dict[AxisId, Halo]] = {}
304 outputs = {t.id: t for t in model.outputs}
305 all_tensors = {**{t.id: t for t in model.inputs}, **outputs}
306 if output_halo is None:
307 output_halo = {
308 t.id: {
309 a.id: Halo(a.halo, a.halo)
310 for a in t.axes
311 if isinstance(a, v0_5.WithHalo)
312 }
313 for t in model.outputs
314 }
316 for t, th in output_halo.items():
317 axes = {a.id: a for a in outputs[t].axes}
319 for a, ah in th.items():
320 s = axes[a].size
321 if not isinstance(s, v0_5.SizeReference):
322 raise ValueError(
323 f"Unable to map output halo for {t}.{a} to an input axis"
324 )
326 axis = axes[a]
327 ref_axis = {a.id: a for a in all_tensors[s.tensor_id].axes}[s.axis_id]
329 input_halo_left = ah.left * axis.scale / ref_axis.scale
330 input_halo_right = ah.right * axis.scale / ref_axis.scale
331 assert input_halo_left == int(input_halo_left), f"{input_halo_left} not int"
332 assert input_halo_right == int(input_halo_right), (
333 f"{input_halo_right} not int"
334 )
336 input_halo.setdefault(s.tensor_id, {})[a] = Halo(
337 int(input_halo_left), int(input_halo_right)
338 )
340 return input_halo
343def get_block_transform(
344 model: v0_5.ModelDescr,
345) -> PerMember[PerAxis[Union[LinearSampleAxisTransform, int]]]:
346 """returns how a model's output tensor shapes relates to its input shapes"""
347 ret: Dict[MemberId, Dict[AxisId, Union[LinearSampleAxisTransform, int]]] = {}
348 batch_axis_trf = None
349 for ipt in model.inputs:
350 for a in ipt.axes:
351 if a.type == "batch":
352 batch_axis_trf = LinearSampleAxisTransform(
353 axis=a.id, scale=1, offset=0, member=ipt.id
354 )
355 break
356 if batch_axis_trf is not None:
357 break
358 axis_scales = {
359 t.id: {a.id: a.scale for a in t.axes}
360 for t in chain(model.inputs, model.outputs)
361 }
362 for out in model.outputs:
363 new_axes: Dict[AxisId, Union[LinearSampleAxisTransform, int]] = {}
364 for a in out.axes:
365 if a.size is None:
366 assert a.type == "batch"
367 if batch_axis_trf is None:
368 raise ValueError(
369 "no batch axis found in any input tensor, but output tensor"
370 + f" '{out.id}' has one."
371 )
372 s = batch_axis_trf
373 elif isinstance(a.size, int):
374 s = a.size
375 elif isinstance(a.size, v0_5.DataDependentSize):
376 s = -1
377 elif isinstance(a.size, v0_5.SizeReference):
378 s = LinearSampleAxisTransform(
379 axis=a.size.axis_id,
380 scale=axis_scales[a.size.tensor_id][a.size.axis_id] / a.scale,
381 offset=a.size.offset,
382 member=a.size.tensor_id,
383 )
384 else:
385 assert_never(a.size)
387 new_axes[a.id] = s
389 # account for postprocessing that changes the nubmer of output channels by
390 # overwriting described output shape by the intermediate output shape
391 c = AxisId("channel")
392 if c not in new_axes:
393 continue
394 for post in out.postprocessing:
395 if post.id == "cellpose_flow_dynamics":
396 new_axes[c] = 3
397 break
398 elif post.id == "stardist_postprocessing":
399 new_axes[c] = post.kwargs.n_rays + 1
400 break
402 ret[out.id] = new_axes
404 return ret
407def get_io_sample_block_metas(
408 model: v0_5.ModelDescr,
409 input_sample_shape: PerMember[PerAxis[int]],
410 ns: Mapping[Tuple[MemberId, AxisId], ParameterizedSize_N],
411 batch_size: int = 1,
412) -> Tuple[TotalNumberOfBlocks, Iterable[IO_SampleBlockMeta]]:
413 """returns an iterable yielding meta data for corresponding input and output samples"""
414 if not isinstance(model, v0_5.ModelDescr):
415 raise TypeError(f"get_block_meta() not implemented for {type(model)}")
417 block_axis_sizes = model.get_axis_sizes(ns=ns, batch_size=batch_size)
418 input_block_shape = {
419 t: {aa: s for (tt, aa), s in block_axis_sizes.inputs.items() if tt == t}
420 for t in {tt for tt, _ in block_axis_sizes.inputs}
421 }
422 output_halo = {
423 t.id: {
424 a.id: Halo(a.halo, a.halo) for a in t.axes if isinstance(a, v0_5.WithHalo)
425 }
426 for t in model.outputs
427 }
428 input_halo = get_input_halo(model, output_halo)
430 n_input_blocks, input_blocks = split_multiple_shapes_into_blocks(
431 input_sample_shape, input_block_shape, halo=input_halo
432 )
433 block_transform = get_block_transform(model)
434 return n_input_blocks, (
435 IO_SampleBlockMeta(ipt, ipt.get_transformed(block_transform))
436 for ipt in sample_block_meta_generator(
437 input_blocks, sample_shape=input_sample_shape, sample_id=None
438 )
439 )
442def get_tensor(
443 src: TensorSource,
444 descr: Union[
445 v0_4.InputTensorDescr,
446 v0_5.InputTensorDescr,
447 v0_4.OutputTensorDescr,
448 v0_5.OutputTensorDescr,
449 Sequence[AxisInfo],
450 ],
451 *,
452 extra_dims: Literal["raise", "squeeze", "stack", "squeeze_or_stack"] = "squeeze",
453 missing_dims: Literal["raise", "expand", "unstack", "unstack_or_expand"] = "raise",
454):
455 """helper to cast/load various tensor sources
457 Args:
458 src: the tensor source to load/cast
459 descr: the tensor description or a sequence of axis infos
460 extra_dims:
461 How to handle extra dimensions in the input tensors.
462 If "raise", any extra dimensions will raise an error.
463 If "squeeze", any extra singleton dimensions will be squeezed, non-singleton dimensions will raise an error.
464 If "stack", any extra dimensions will be stacked to the batch dimension. Such a stacked batch dimension then has a multi-index that can be unstacked using `Tensor.unstack_batch_multi_index()`.
465 If "squeeze_or_stack", any extra singleton dimensions will be squeezed, non-singleton dimensions will be stacked to the batch dimension.
466 """
468 if isinstance(
469 descr,
470 (
471 v0_4.InputTensorDescr,
472 v0_5.InputTensorDescr,
473 v0_4.OutputTensorDescr,
474 v0_5.OutputTensorDescr,
475 ),
476 ):
477 axes = get_axes_infos(descr)
478 else:
479 axes = descr
481 if isinstance(src, Tensor):
482 output_dims = [a.id for a in axes]
483 return src.transpose(
484 output_dims, extra_dims=extra_dims, missing_dims=missing_dims
485 )
486 elif isinstance(src, xr.DataArray):
487 output_dims = [a.id for a in axes]
489 return Tensor.from_xarray(src).transpose(
490 axes=output_dims, extra_dims=extra_dims, missing_dims=missing_dims
491 )
492 elif isinstance(src, np.ndarray):
493 return Tensor.from_numpy(src, dims=axes)
494 else:
495 return load_tensor(src, axes=axes)
498def create_sample_for_model(
499 model: AnyModelDescr,
500 *,
501 stat: Optional[Stat] = None,
502 sample_id: SampleId = None,
503 inputs: Union[PerMember[TensorSource], TensorSource],
504 extra_dims: Literal["raise", "squeeze", "stack", "squeeze_or_stack"] = "stack",
505) -> Sample:
506 """Create a sample from a single set of input(s) for a specific bioimage.io model
508 Args:
509 model: a bioimage.io model description
510 stat: dictionary with sample and dataset statistics (may be updated in-place!)
511 inputs: the input(s) constituting a single sample.
512 extra_dims: How to handle extra dimensions in the input tensors.
513 If "raise", any extra dimensions will raise an error.
514 If "squeeze", any extra singleton dimensions will be squeezed, non-singleton dimensions will raise an error.
515 If "stack", any extra dimensions will be stacked to the batch dimension. Such a stacked batch dimension then has a multi-index that can be unstacked using `Tensor.unstack_batch_multi_index()`.
516 If "squeeze_or_stack", any extra singleton dimensions will be squeezed, non-singleton dimensions will be stacked to the batch dimension.
517 """
519 model_inputs = {get_member_id(d): d for d in model.inputs}
520 if isinstance(inputs, collections.abc.Mapping):
521 inputs = {MemberId(k): v for k, v in inputs.items()}
522 elif len(model_inputs) == 1:
523 inputs = {list(model_inputs)[0]: inputs}
524 else:
525 raise TypeError(
526 f"Expected `inputs` to be a mapping with keys {tuple(model_inputs)}"
527 )
529 if unknown := {k for k in inputs if k not in model_inputs}:
530 raise ValueError(f"Got unexpected inputs: {unknown}")
532 if missing := {
533 k
534 for k, v in model_inputs.items()
535 if k not in inputs and not (isinstance(v, v0_5.InputTensorDescr) and v.optional)
536 }:
537 raise ValueError(f"Missing non-optional model inputs: {missing}")
539 return Sample(
540 members={
541 m: get_tensor(inputs[m], ipt, extra_dims=extra_dims)
542 for m, ipt in model_inputs.items()
543 if m in inputs
544 },
545 stat={} if stat is None else stat,
546 id=sample_id,
547 )
550def load_sample_for_model(
551 *,
552 model: AnyModelDescr,
553 paths: PerMember[Path],
554 axes: Optional[PerMember[Sequence[AxisLike]]] = None,
555 stat: Optional[Stat] = None,
556 sample_id: Optional[SampleId] = None,
557):
558 """load a single sample from `paths` that can be processed by `model`"""
560 if axes is None:
561 axes = {}
563 # make sure members are keyed by MemberId, not string
564 paths = {MemberId(k): v for k, v in paths.items()}
565 axes = {MemberId(k): v for k, v in axes.items()}
567 model_inputs = {get_member_id(d): d for d in model.inputs}
569 if unknown := {k for k in paths if k not in model_inputs}:
570 raise ValueError(f"Got unexpected paths for {unknown}")
572 if unknown := {k for k in axes if k not in model_inputs}:
573 raise ValueError(f"Got unexpected axes hints for: {unknown}")
575 members: Dict[MemberId, Tensor] = {}
576 for m, p in paths.items():
577 if m not in axes:
578 axes[m] = get_axes_infos(model_inputs[m])
579 logger.info(
580 "loading '{}' from {} with default input axes {} ",
581 m,
582 p,
583 axes[m],
584 )
585 members[m] = load_tensor(p, axes[m])
587 return Sample(
588 members=members,
589 stat={} if stat is None else stat,
590 id=sample_id or tuple(sorted(paths.values())),
591 )
594def split_sample_into_blocks_for_model(
595 sample: Sample,
596 model: v0_5.ModelDescr,
597 blocksize_parameter: int,
598 batch_size: int = 1,
599) -> Tuple[TotalNumberOfBlocks, Iterable[SampleBlockWithOrigin]]:
600 if isinstance(model, v0_4.ModelDescr):
601 raise NotImplementedError(
602 "`predict_sample_with_blocking` not implemented for v0_4.ModelDescr"
603 + f" {model.name}."
604 + " Consider using `predict_sample_with_fixed_blocking` or update the model description to format version 0.5."
605 )
607 ns = {
608 (ipt.id, a.id): blocksize_parameter
609 for ipt in model.inputs
610 for a in ipt.axes
611 if isinstance(a.size, v0_5.ParameterizedSize)
612 }
613 halo = get_input_halo(model)
615 input_block_shape = model.get_tensor_sizes(ns, batch_size=batch_size).inputs
617 return sample.split_into_blocks(
618 block_shapes=input_block_shape,
619 halo=halo,
620 pad_mode={ipt.id: ipt.pad or "symmetric" for ipt in model.inputs},
621 )
624def transpose_sample_for_model(
625 sample: Sample,
626 model: AnyModelDescr,
627 extra_dims: Literal["raise", "squeeze", "stack", "squeeze_or_stack"] = "stack",
628 missing_dims: Literal["raise", "expand", "unstack", "unstack_or_expand"] = "raise",
629) -> Sample:
630 """Transpose sample members to the order expected as inputs by the model.
632 Unexpected dimensions are stacked to the batch dimension, missing dimensions are added as singletons.
633 """
635 axes = {get_member_id(d): [a.id for a in get_axes_infos(d)] for d in model.inputs}
636 for m in axes:
637 if m not in sample.members and not (
638 isinstance(model, v0_5.ModelDescr)
639 and [d for d in model.inputs if get_member_id(d) == m][0].optional
640 ):
641 raise ValueError(
642 f"Sample is missing non-optional member {m} required by model"
643 )
645 return sample.transpose(
646 axes=axes,
647 extra_dims=extra_dims,
648 missing_dims=missing_dims,
649 )