Coverage for src / bioimageio / spec / model / v0_5.py: 71%
1657 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-05-18 15:21 +0000
1from __future__ import annotations
3import collections.abc
4import re
5import string
6import warnings
7from copy import deepcopy
8from functools import partial
9from itertools import chain
10from math import ceil
11from pathlib import Path, PurePosixPath
12from tempfile import mkdtemp
13from textwrap import dedent
14from typing import (
15 TYPE_CHECKING,
16 Any,
17 Callable,
18 ClassVar,
19 Dict,
20 Generic,
21 List,
22 Literal,
23 Mapping,
24 NamedTuple,
25 Optional,
26 Sequence,
27 Set,
28 Tuple,
29 Type,
30 TypeVar,
31 Union,
32 cast,
33 overload,
34)
36import numpy as np
37from annotated_types import Ge, Gt, Interval, MaxLen, MinLen, Predicate
38from imageio.v3 import imread, imwrite # pyright: ignore[reportUnknownVariableType]
39from loguru import logger
40from numpy.typing import NDArray
41from pydantic import (
42 AfterValidator,
43 Discriminator,
44 Field,
45 RootModel,
46 SerializationInfo,
47 SerializerFunctionWrapHandler,
48 StrictInt,
49 Tag,
50 ValidationInfo,
51 WrapSerializer,
52 field_validator,
53 model_serializer,
54 model_validator,
55)
56from typing_extensions import Annotated, Self, assert_never, get_args
58from .._internal.common_nodes import (
59 InvalidDescr,
60 KwargsNode,
61 Node,
62 NodeWithExplicitlySetFields,
63)
64from .._internal.constants import DTYPE_LIMITS
65from .._internal.field_warning import issue_warning, warn
66from .._internal.io import BioimageioYamlContent as BioimageioYamlContent
67from .._internal.io import FileDescr as FileDescr
68from .._internal.io import (
69 FileSource,
70 WithSuffix,
71 YamlValue,
72 extract_file_name,
73 get_reader,
74 wo_special_file_name,
75)
76from .._internal.io_basics import Sha256 as Sha256
77from .._internal.io_packaging import (
78 FileDescr_,
79 FileSource_,
80 package_file_descr_serializer,
81)
82from .._internal.io_utils import load_array
83from .._internal.node_converter import Converter
84from .._internal.type_guards import is_dict, is_sequence
85from .._internal.types import (
86 FAIR,
87 AbsoluteTolerance,
88 LowerCaseIdentifier,
89 LowerCaseIdentifierAnno,
90 MismatchedElementsPerMillion,
91 RelativeTolerance,
92)
93from .._internal.types import Datetime as Datetime
94from .._internal.types import Identifier as Identifier
95from .._internal.types import NotEmpty as NotEmpty
96from .._internal.types import SiUnit as SiUnit
97from .._internal.url import HttpUrl as HttpUrl
98from .._internal.utils import try_all_raise_last
99from .._internal.validation_context import get_validation_context
100from .._internal.validator_annotations import RestrictCharacters
101from .._internal.version_type import Version as Version
102from .._internal.warning_levels import INFO
103from ..dataset.v0_2 import DatasetDescr as DatasetDescr02
104from ..dataset.v0_2 import LinkedDataset as LinkedDataset02
105from ..dataset.v0_3 import DatasetDescr as DatasetDescr
106from ..dataset.v0_3 import DatasetId as DatasetId
107from ..dataset.v0_3 import LinkedDataset as LinkedDataset
108from ..dataset.v0_3 import Uploader as Uploader
109from ..generic.v0_3 import (
110 VALID_COVER_IMAGE_EXTENSIONS as VALID_COVER_IMAGE_EXTENSIONS,
111)
112from ..generic.v0_3 import Author as Author
113from ..generic.v0_3 import BadgeDescr as BadgeDescr
114from ..generic.v0_3 import CiteEntry as CiteEntry
115from ..generic.v0_3 import DeprecatedLicenseId as DeprecatedLicenseId
116from ..generic.v0_3 import Doi as Doi
117from ..generic.v0_3 import (
118 FileSource_documentation,
119 GenericModelDescrBase,
120 LinkedResourceBase,
121 _author_conv, # pyright: ignore[reportPrivateUsage]
122 _maintainer_conv, # pyright: ignore[reportPrivateUsage]
123)
124from ..generic.v0_3 import LicenseId as LicenseId
125from ..generic.v0_3 import LinkedResource as LinkedResource
126from ..generic.v0_3 import Maintainer as Maintainer
127from ..generic.v0_3 import OrcidId as OrcidId
128from ..generic.v0_3 import RelativeFilePath as RelativeFilePath
129from ..generic.v0_3 import ResourceId as ResourceId
130from .v0_4 import Author as _Author_v0_4
131from .v0_4 import BinarizeDescr as _BinarizeDescr_v0_4
132from .v0_4 import CallableFromDepencency as CallableFromDepencency
133from .v0_4 import CallableFromDepencency as _CallableFromDepencency_v0_4
134from .v0_4 import CallableFromFile as _CallableFromFile_v0_4
135from .v0_4 import ClipDescr as _ClipDescr_v0_4
136from .v0_4 import ImplicitOutputShape as _ImplicitOutputShape_v0_4
137from .v0_4 import InputTensorDescr as _InputTensorDescr_v0_4
138from .v0_4 import KnownRunMode as KnownRunMode
139from .v0_4 import ModelDescr as _ModelDescr_v0_4
140from .v0_4 import OutputTensorDescr as _OutputTensorDescr_v0_4
141from .v0_4 import ParameterizedInputShape as _ParameterizedInputShape_v0_4
142from .v0_4 import PostprocessingDescr as _PostprocessingDescr_v0_4
143from .v0_4 import PreprocessingDescr as _PreprocessingDescr_v0_4
144from .v0_4 import RunMode as RunMode
145from .v0_4 import ScaleLinearDescr as _ScaleLinearDescr_v0_4
146from .v0_4 import ScaleMeanVarianceDescr as _ScaleMeanVarianceDescr_v0_4
147from .v0_4 import ScaleRangeDescr as _ScaleRangeDescr_v0_4
148from .v0_4 import SigmoidDescr as _SigmoidDescr_v0_4
149from .v0_4 import TensorName as _TensorName_v0_4
150from .v0_4 import ZeroMeanUnitVarianceDescr as _ZeroMeanUnitVarianceDescr_v0_4
151from .v0_4 import package_weights
153SpaceUnit = Literal[
154 "attometer",
155 "angstrom",
156 "centimeter",
157 "decimeter",
158 "exameter",
159 "femtometer",
160 "foot",
161 "gigameter",
162 "hectometer",
163 "inch",
164 "kilometer",
165 "megameter",
166 "meter",
167 "micrometer",
168 "mile",
169 "millimeter",
170 "nanometer",
171 "parsec",
172 "petameter",
173 "picometer",
174 "terameter",
175 "yard",
176 "yoctometer",
177 "yottameter",
178 "zeptometer",
179 "zettameter",
180]
181"""Space unit compatible to the [OME-Zarr axes specification 0.5](https://ngff.openmicroscopy.org/0.5/#axes-md)"""
183TimeUnit = Literal[
184 "attosecond",
185 "centisecond",
186 "day",
187 "decisecond",
188 "exasecond",
189 "femtosecond",
190 "gigasecond",
191 "hectosecond",
192 "hour",
193 "kilosecond",
194 "megasecond",
195 "microsecond",
196 "millisecond",
197 "minute",
198 "nanosecond",
199 "petasecond",
200 "picosecond",
201 "second",
202 "terasecond",
203 "yoctosecond",
204 "yottasecond",
205 "zeptosecond",
206 "zettasecond",
207]
208"""Time unit compatible to the [OME-Zarr axes specification 0.5](https://ngff.openmicroscopy.org/0.5/#axes-md)"""
210AxisType = Literal["batch", "channel", "index", "time", "space"]
212_AXIS_TYPE_MAP: Mapping[str, AxisType] = {
213 "b": "batch",
214 "t": "time",
215 "i": "index",
216 "c": "channel",
217 "x": "space",
218 "y": "space",
219 "z": "space",
220}
222_AXIS_ID_MAP = {
223 "b": "batch",
224 "t": "time",
225 "i": "index",
226 "c": "channel",
227}
229WeightsFormat = Literal[
230 "keras_hdf5",
231 "keras_v3",
232 "onnx",
233 "pytorch_state_dict",
234 "tensorflow_js",
235 "tensorflow_saved_model_bundle",
236 "torchscript",
237]
240class TensorId(LowerCaseIdentifier):
241 root_model: ClassVar[Type[RootModel[Any]]] = RootModel[
242 Annotated[LowerCaseIdentifierAnno, MaxLen(32)]
243 ]
246def _normalize_axis_id(a: str):
247 a = str(a)
248 normalized = _AXIS_ID_MAP.get(a, a)
249 if a != normalized:
250 logger.opt(depth=3).warning(
251 "Normalized axis id from '{}' to '{}'.", a, normalized
252 )
253 return normalized
256class AxisId(LowerCaseIdentifier):
257 root_model: ClassVar[Type[RootModel[Any]]] = RootModel[
258 Annotated[
259 LowerCaseIdentifierAnno,
260 MaxLen(16),
261 AfterValidator(_normalize_axis_id),
262 ]
263 ]
266def _is_batch(a: str) -> bool:
267 return str(a) == "batch"
270def _is_not_batch(a: str) -> bool:
271 return not _is_batch(a)
274NonBatchAxisId = Annotated[AxisId, Predicate(_is_not_batch)]
276PreprocessingId = Literal[
277 "binarize",
278 "clip",
279 "ensure_dtype",
280 "fixed_zero_mean_unit_variance",
281 "scale_linear",
282 "scale_range",
283 "sigmoid",
284 "softmax",
285]
286PostprocessingId = Literal[
287 "binarize",
288 "clip",
289 "custom",
290 "ensure_dtype",
291 "fixed_zero_mean_unit_variance",
292 "scale_linear",
293 "scale_mean_variance",
294 "scale_range",
295 "sigmoid",
296 "softmax",
297 "zero_mean_unit_variance",
298]
301SAME_AS_TYPE = "<same as type>"
304ParameterizedSize_N = int
305"""
306Annotates an integer to calculate a concrete axis size from a `ParameterizedSize`.
307"""
310class ParameterizedSize(Node):
311 """Describes a range of valid tensor axis sizes as `size = min + n*step`.
313 - **min** and **step** are given by the model description.
314 - All blocksize paramters n = 0,1,2,... yield a valid `size`.
315 - A greater blocksize paramter n = 0,1,2,... results in a greater **size**.
316 This allows to adjust the axis size more generically.
317 """
319 N: ClassVar[Type[int]] = ParameterizedSize_N
320 """Positive integer to parameterize this axis"""
322 min: Annotated[int, Gt(0)]
323 step: Annotated[int, Gt(0)]
325 def validate_size(self, size: int, msg_prefix: str = "") -> int:
326 if size < self.min:
327 raise ValueError(
328 f"{msg_prefix}size {size} < {self.min} (minimum axis size)"
329 )
330 if (size - self.min) % self.step != 0:
331 raise ValueError(
332 f"{msg_prefix}size {size} is not parameterized by `min + n*step` ="
333 + f" `{self.min} + n*{self.step}`"
334 )
336 return size
338 def get_size(self, n: ParameterizedSize_N) -> int:
339 return self.min + self.step * n
341 def get_n(self, s: int) -> ParameterizedSize_N:
342 """return smallest n parameterizing a size greater or equal than `s`"""
343 return ceil((s - self.min) / self.step)
346class DataDependentSize(Node):
347 min: Annotated[int, Gt(0)] = 1
348 max: Annotated[Optional[int], Gt(1)] = None
350 @model_validator(mode="after")
351 def _validate_max_gt_min(self):
352 if self.max is not None and self.min >= self.max:
353 raise ValueError(f"expected `min` < `max`, but got {self.min}, {self.max}")
355 return self
357 def validate_size(self, size: int, msg_prefix: str = "") -> int:
358 if size < self.min:
359 raise ValueError(f"{msg_prefix}size {size} < {self.min}")
361 if self.max is not None and size > self.max:
362 raise ValueError(f"{msg_prefix}size {size} > {self.max}")
364 return size
367class SizeReference(Node):
368 """A tensor axis size (extent in pixels/frames) defined in relation to a reference axis.
370 `axis.size = reference.size * reference.scale / axis.scale + offset`
372 Note:
373 1. The axis and the referenced axis need to have the same unit (or no unit).
374 2. Batch axes may not be referenced.
375 3. Fractions are rounded down.
376 4. If the reference axis is `concatenable` the referencing axis is assumed to be
377 `concatenable` as well with the same block order.
379 Example:
380 An unisotropic input image of w*h=100*49 pixels depicts a phsical space of 200*196mm².
381 Let's assume that we want to express the image height h in relation to its width w
382 instead of only accepting input images of exactly 100*49 pixels
383 (for example to express a range of valid image shapes by parametrizing w, see `ParameterizedSize`).
385 >>> w = SpaceInputAxis(id=AxisId("w"), size=100, unit="millimeter", scale=2)
386 >>> h = SpaceInputAxis(
387 ... id=AxisId("h"),
388 ... size=SizeReference(tensor_id=TensorId("input"), axis_id=AxisId("w"), offset=-1),
389 ... unit="millimeter",
390 ... scale=4,
391 ... )
392 >>> print(h.size.get_size(h, w))
393 49
395 ⇒ h = w * w.scale / h.scale + offset = 100 * 2mm / 4mm - 1 = 49
396 """
398 tensor_id: TensorId
399 """tensor id of the reference axis"""
401 axis_id: AxisId
402 """axis id of the reference axis"""
404 offset: StrictInt = 0
406 def get_size(
407 self,
408 axis: Union[
409 ChannelAxis,
410 IndexInputAxis,
411 IndexOutputAxis,
412 TimeInputAxis,
413 SpaceInputAxis,
414 TimeOutputAxis,
415 TimeOutputAxisWithHalo,
416 SpaceOutputAxis,
417 SpaceOutputAxisWithHalo,
418 ],
419 ref_axis: Union[
420 ChannelAxis,
421 IndexInputAxis,
422 IndexOutputAxis,
423 TimeInputAxis,
424 SpaceInputAxis,
425 TimeOutputAxis,
426 TimeOutputAxisWithHalo,
427 SpaceOutputAxis,
428 SpaceOutputAxisWithHalo,
429 ],
430 n: ParameterizedSize_N = 0,
431 ref_size: Optional[int] = None,
432 ):
433 """Compute the concrete size for a given axis and its reference axis.
435 Args:
436 axis: The axis this [SizeReference][] is the size of.
437 ref_axis: The reference axis to compute the size from.
438 n: If the **ref_axis** is parameterized (of type `ParameterizedSize`)
439 and no fixed **ref_size** is given,
440 **n** is used to compute the size of the parameterized **ref_axis**.
441 ref_size: Overwrite the reference size instead of deriving it from
442 **ref_axis**
443 (**ref_axis.scale** is still used; any given **n** is ignored).
444 """
445 assert axis.size == self, (
446 "Given `axis.size` is not defined by this `SizeReference`"
447 )
449 assert ref_axis.id == self.axis_id, (
450 f"Expected `ref_axis.id` to be {self.axis_id}, but got {ref_axis.id}."
451 )
453 assert axis.unit == ref_axis.unit, (
454 "`SizeReference` requires `axis` and `ref_axis` to have the same `unit`,"
455 f" but {axis.unit}!={ref_axis.unit}"
456 )
457 if ref_size is None:
458 if isinstance(ref_axis.size, (int, float)):
459 ref_size = ref_axis.size
460 elif isinstance(ref_axis.size, ParameterizedSize):
461 ref_size = ref_axis.size.get_size(n)
462 elif isinstance(ref_axis.size, DataDependentSize):
463 raise ValueError(
464 "Reference axis referenced in `SizeReference` may not be a `DataDependentSize`."
465 )
466 elif isinstance(ref_axis.size, SizeReference):
467 raise ValueError(
468 "Reference axis referenced in `SizeReference` may not be sized by a"
469 + " `SizeReference` itself."
470 )
471 else:
472 assert_never(ref_axis.size)
474 return int(ref_size * ref_axis.scale / axis.scale + self.offset)
476 @staticmethod
477 def _get_unit(
478 axis: Union[
479 ChannelAxis,
480 IndexInputAxis,
481 IndexOutputAxis,
482 TimeInputAxis,
483 SpaceInputAxis,
484 TimeOutputAxis,
485 TimeOutputAxisWithHalo,
486 SpaceOutputAxis,
487 SpaceOutputAxisWithHalo,
488 ],
489 ):
490 return axis.unit
493class AxisBase(NodeWithExplicitlySetFields):
494 id: AxisId
495 """An axis id unique across all axes of one tensor."""
497 description: Annotated[str, MaxLen(128)] = ""
498 """A short description of this axis beyond its type and id."""
501class WithHalo(Node):
502 halo: Annotated[int, Ge(1)]
503 """The halo should be cropped from the output tensor to avoid boundary effects.
504 It is to be cropped from both sides, i.e. `size_after_crop = size - 2 * halo`.
505 To document a halo that is already cropped by the model use `size.offset` instead."""
507 size: Annotated[
508 SizeReference,
509 Field(
510 examples=[
511 10,
512 SizeReference(
513 tensor_id=TensorId("t"), axis_id=AxisId("a"), offset=5
514 ).model_dump(mode="json"),
515 ]
516 ),
517 ]
518 """reference to another axis with an optional offset (see [SizeReference][])"""
521BATCH_AXIS_ID = AxisId("batch")
524class BatchAxis(AxisBase):
525 implemented_type: ClassVar[Literal["batch"]] = "batch"
526 if TYPE_CHECKING:
527 type: Literal["batch"] = "batch"
528 else:
529 type: Literal["batch"]
531 id: Annotated[AxisId, Predicate(_is_batch)] = BATCH_AXIS_ID
532 size: Optional[Literal[1]] = None
533 """The batch size may be fixed to 1,
534 otherwise (the default) it may be chosen arbitrarily depending on available memory"""
536 @property
537 def scale(self):
538 return 1.0
540 @property
541 def concatenable(self):
542 return True
544 @property
545 def unit(self):
546 return None
549class ChannelAxis(AxisBase):
550 implemented_type: ClassVar[Literal["channel"]] = "channel"
551 if TYPE_CHECKING:
552 type: Literal["channel"] = "channel"
553 else:
554 type: Literal["channel"]
556 id: NonBatchAxisId = AxisId("channel")
558 channel_names: NotEmpty[List[str]]
560 @property
561 def size(self) -> int:
562 return len(self.channel_names)
564 @property
565 def concatenable(self):
566 return False
568 @property
569 def scale(self) -> float:
570 return 1.0
572 @property
573 def unit(self):
574 return None
577class _WithInputAxisSize(Node):
578 size: Annotated[
579 Union[Annotated[int, Gt(0)], ParameterizedSize, SizeReference],
580 Field(
581 examples=[
582 10,
583 ParameterizedSize(min=32, step=16).model_dump(mode="json"),
584 SizeReference(
585 tensor_id=TensorId("t"), axis_id=AxisId("a"), offset=5
586 ).model_dump(mode="json"),
587 ]
588 ),
589 ]
590 """The size/length of this axis can be specified as
591 - fixed integer
592 - parameterized series of valid sizes ([ParameterizedSize][])
593 - reference to another axis with an optional offset ([SizeReference][])
594 """
597class IndexAxisBase(AxisBase):
598 implemented_type: ClassVar[Literal["index"]] = "index"
599 if TYPE_CHECKING:
600 type: Literal["index"] = "index"
601 else:
602 type: Literal["index"]
604 id: NonBatchAxisId = AxisId("index")
606 @property
607 def scale(self) -> float:
608 return 1.0
610 @property
611 def unit(self):
612 return None
615class IndexInputAxis(IndexAxisBase, _WithInputAxisSize):
616 concatenable: bool = False
617 """If a model has a `concatenable` input axis, it can be processed blockwise,
618 splitting a longer sample axis into blocks matching its input tensor description.
619 Output axes are concatenable if they have a [SizeReference][] to a concatenable
620 input axis.
621 """
624class IndexOutputAxis(IndexAxisBase):
625 size: Annotated[
626 Union[Annotated[int, Gt(0)], SizeReference, DataDependentSize],
627 Field(
628 examples=[
629 10,
630 SizeReference(
631 tensor_id=TensorId("t"), axis_id=AxisId("a"), offset=5
632 ).model_dump(mode="json"),
633 ]
634 ),
635 ]
636 """The size/length of this axis can be specified as
637 - fixed integer
638 - reference to another axis with an optional offset ([SizeReference][])
639 - data dependent size using [DataDependentSize][] (size is only known after model inference)
640 """
643class TimeAxisBase(AxisBase):
644 implemented_type: ClassVar[Literal["time"]] = "time"
645 if TYPE_CHECKING:
646 type: Literal["time"] = "time"
647 else:
648 type: Literal["time"]
650 id: NonBatchAxisId = AxisId("time")
651 unit: Optional[TimeUnit] = None
652 scale: Annotated[float, Gt(0)] = 1.0
655class TimeInputAxis(TimeAxisBase, _WithInputAxisSize):
656 concatenable: bool = False
657 """If a model has a `concatenable` input axis, it can be processed blockwise,
658 splitting a longer sample axis into blocks matching its input tensor description.
659 Output axes are concatenable if they have a [SizeReference][] to a concatenable
660 input axis.
661 """
664class SpaceAxisBase(AxisBase):
665 implemented_type: ClassVar[Literal["space"]] = "space"
666 if TYPE_CHECKING:
667 type: Literal["space"] = "space"
668 else:
669 type: Literal["space"]
671 id: Annotated[NonBatchAxisId, Field(examples=["x", "y", "z"])] = AxisId("x")
672 unit: Optional[SpaceUnit] = None
673 scale: Annotated[float, Gt(0)] = 1.0
676class SpaceInputAxis(SpaceAxisBase, _WithInputAxisSize):
677 concatenable: bool = False
678 """If a model has a `concatenable` input axis, it can be processed blockwise,
679 splitting a longer sample axis into blocks matching its input tensor description.
680 Output axes are concatenable if they have a [SizeReference][] to a concatenable
681 input axis.
682 """
685INPUT_AXIS_TYPES = (
686 BatchAxis,
687 ChannelAxis,
688 IndexInputAxis,
689 TimeInputAxis,
690 SpaceInputAxis,
691)
692"""intended for isinstance comparisons in py<3.10"""
694_InputAxisUnion = Union[
695 BatchAxis, ChannelAxis, IndexInputAxis, TimeInputAxis, SpaceInputAxis
696]
697InputAxis = Annotated[_InputAxisUnion, Discriminator("type")]
700class _WithOutputAxisSize(Node):
701 size: Annotated[
702 Union[Annotated[int, Gt(0)], SizeReference],
703 Field(
704 examples=[
705 10,
706 SizeReference(
707 tensor_id=TensorId("t"), axis_id=AxisId("a"), offset=5
708 ).model_dump(mode="json"),
709 ]
710 ),
711 ]
712 """The size/length of this axis can be specified as
713 - fixed integer
714 - reference to another axis with an optional offset (see [SizeReference][])
715 """
718class TimeOutputAxis(TimeAxisBase, _WithOutputAxisSize):
719 pass
722class TimeOutputAxisWithHalo(TimeAxisBase, WithHalo):
723 pass
726def _get_halo_axis_discriminator_value(v: Any) -> Literal["with_halo", "wo_halo"]:
727 if isinstance(v, dict):
728 return "with_halo" if "halo" in v else "wo_halo"
729 else:
730 return "with_halo" if hasattr(v, "halo") else "wo_halo"
733_TimeOutputAxisUnion = Annotated[
734 Union[
735 Annotated[TimeOutputAxis, Tag("wo_halo")],
736 Annotated[TimeOutputAxisWithHalo, Tag("with_halo")],
737 ],
738 Discriminator(_get_halo_axis_discriminator_value),
739]
742class SpaceOutputAxis(SpaceAxisBase, _WithOutputAxisSize):
743 pass
746class SpaceOutputAxisWithHalo(SpaceAxisBase, WithHalo):
747 pass
750_SpaceOutputAxisUnion = Annotated[
751 Union[
752 Annotated[SpaceOutputAxis, Tag("wo_halo")],
753 Annotated[SpaceOutputAxisWithHalo, Tag("with_halo")],
754 ],
755 Discriminator(_get_halo_axis_discriminator_value),
756]
759_OutputAxisUnion = Union[
760 BatchAxis, ChannelAxis, IndexOutputAxis, _TimeOutputAxisUnion, _SpaceOutputAxisUnion
761]
762OutputAxis = Annotated[_OutputAxisUnion, Discriminator("type")]
764OUTPUT_AXIS_TYPES = (
765 BatchAxis,
766 ChannelAxis,
767 IndexOutputAxis,
768 TimeOutputAxis,
769 TimeOutputAxisWithHalo,
770 SpaceOutputAxis,
771 SpaceOutputAxisWithHalo,
772)
773"""intended for isinstance comparisons in py<3.10"""
776AnyAxis = Union[InputAxis, OutputAxis]
778ANY_AXIS_TYPES = INPUT_AXIS_TYPES + OUTPUT_AXIS_TYPES
779"""intended for isinstance comparisons in py<3.10"""
781TVs = Union[
782 NotEmpty[List[int]],
783 NotEmpty[List[float]],
784 NotEmpty[List[bool]],
785 NotEmpty[List[str]],
786]
789NominalOrOrdinalDType = Literal[
790 "float32",
791 "float64",
792 "uint8",
793 "int8",
794 "uint16",
795 "int16",
796 "uint32",
797 "int32",
798 "uint64",
799 "int64",
800 "bool",
801]
804class NominalOrOrdinalDataDescr(Node):
805 values: TVs
806 """A fixed set of nominal or an ascending sequence of ordinal values.
807 In this case `data.type` is required to be an unsigend integer type, e.g. 'uint8'.
808 String `values` are interpreted as labels for tensor values 0, ..., N.
809 Note: as YAML 1.2 does not natively support a "set" datatype,
810 nominal values should be given as a sequence (aka list/array) as well.
811 """
813 type: Annotated[
814 NominalOrOrdinalDType,
815 Field(
816 examples=[
817 "float32",
818 "uint8",
819 "uint16",
820 "int64",
821 "bool",
822 ],
823 ),
824 ] = "uint8"
826 @model_validator(mode="after")
827 def _validate_values_match_type(
828 self,
829 ) -> Self:
830 incompatible: List[Any] = []
831 for v in self.values:
832 if self.type == "bool":
833 if not isinstance(v, bool):
834 incompatible.append(v)
835 elif self.type in DTYPE_LIMITS:
836 if (
837 isinstance(v, (int, float))
838 and (
839 v < DTYPE_LIMITS[self.type].min
840 or v > DTYPE_LIMITS[self.type].max
841 )
842 or (isinstance(v, str) and "uint" not in self.type)
843 or (isinstance(v, float) and "int" in self.type)
844 ):
845 incompatible.append(v)
846 else:
847 incompatible.append(v)
849 if len(incompatible) == 5:
850 incompatible.append("...")
851 break
853 if incompatible:
854 raise ValueError(
855 f"data type '{self.type}' incompatible with values {incompatible}"
856 )
858 return self
860 unit: Optional[Union[Literal["arbitrary unit"], SiUnit]] = None
862 @property
863 def range(self):
864 if isinstance(self.values[0], str):
865 return 0, len(self.values) - 1
866 else:
867 return min(self.values), max(self.values)
870IntervalOrRatioDType = Literal[
871 "float32",
872 "float64",
873 "uint8",
874 "int8",
875 "uint16",
876 "int16",
877 "uint32",
878 "int32",
879 "uint64",
880 "int64",
881]
884class IntervalOrRatioDataDescr(Node):
885 type: Annotated[ # TODO: rename to dtype
886 IntervalOrRatioDType,
887 Field(
888 examples=["float32", "float64", "uint8", "uint16"],
889 ),
890 ] = "float32"
891 range: Tuple[Optional[float], Optional[float]] = (
892 None,
893 None,
894 )
895 """Tuple `(minimum, maximum)` specifying the allowed range of the data in this tensor.
896 `None` corresponds to min/max of what can be expressed by **type**."""
897 unit: Union[Literal["arbitrary unit"], SiUnit] = "arbitrary unit"
898 scale: float = 1.0
899 """Scale for data on an interval (or ratio) scale."""
900 offset: Optional[float] = None
901 """Offset for data on a ratio scale."""
903 @model_validator(mode="before")
904 def _replace_inf(cls, data: Any):
905 if is_dict(data):
906 if "range" in data and is_sequence(data["range"]):
907 forbidden = (
908 "inf",
909 "-inf",
910 ".inf",
911 "-.inf",
912 float("inf"),
913 float("-inf"),
914 )
915 if any(v in forbidden for v in data["range"]):
916 issue_warning("replaced 'inf' value", value=data["range"])
918 data["range"] = tuple(
919 (None if v in forbidden else v) for v in data["range"]
920 )
922 return data
925TensorDataDescr = Union[NominalOrOrdinalDataDescr, IntervalOrRatioDataDescr]
928class BinarizeKwargs(KwargsNode):
929 """key word arguments for [BinarizeDescr][]"""
931 threshold: float
932 """The fixed threshold"""
935class BinarizeAlongAxisKwargs(KwargsNode):
936 """key word arguments for [BinarizeDescr][]"""
938 threshold: NotEmpty[List[float]]
939 """The fixed threshold values along `axis`"""
941 axis: Annotated[NonBatchAxisId, Field(examples=["channel"])]
942 """The `threshold` axis"""
945class BinarizeDescr(NodeWithExplicitlySetFields):
946 """Binarize the tensor with a fixed threshold.
948 Values above [BinarizeKwargs.threshold][]/[BinarizeAlongAxisKwargs.threshold][]
949 will be set to one, values below the threshold to zero.
951 Examples:
952 - in YAML
953 ```yaml
954 postprocessing:
955 - id: binarize
956 kwargs:
957 axis: 'channel'
958 threshold: [0.25, 0.5, 0.75]
959 ```
960 - in Python:
962 >>> postprocessing = [BinarizeDescr(
963 ... kwargs=BinarizeAlongAxisKwargs(
964 ... axis=AxisId('channel'),
965 ... threshold=[0.25, 0.5, 0.75],
966 ... )
967 ... )]
968 """
970 implemented_id: ClassVar[Literal["binarize"]] = "binarize"
971 if TYPE_CHECKING:
972 id: Literal["binarize"] = "binarize"
973 else:
974 id: Literal["binarize"]
975 kwargs: Union[BinarizeKwargs, BinarizeAlongAxisKwargs]
978class ClipKwargs(KwargsNode):
979 """key word arguments for [ClipDescr][]"""
981 min: Optional[float] = None
982 """Minimum value for clipping.
984 Exclusive with [min_percentile][]
985 """
986 min_percentile: Optional[Annotated[float, Interval(ge=0, lt=100)]] = None
987 """Minimum percentile for clipping.
989 Exclusive with [min][].
991 In range [0, 100).
992 """
994 max: Optional[float] = None
995 """Maximum value for clipping.
997 Exclusive with `max_percentile`.
998 """
999 max_percentile: Optional[Annotated[float, Interval(gt=1, le=100)]] = None
1000 """Maximum percentile for clipping.
1002 Exclusive with `max`.
1004 In range (1, 100].
1005 """
1007 axes: Annotated[
1008 Optional[Sequence[AxisId]], Field(examples=[("batch", "x", "y")])
1009 ] = None
1010 """The subset of axes to determine percentiles jointly,
1012 i.e. axes to reduce to compute min/max from `min_percentile`/`max_percentile`.
1013 For example to clip 'batch', 'x' and 'y' jointly in a tensor ('batch', 'channel', 'y', 'x')
1014 resulting in a tensor of equal shape with clipped values per channel, specify `axes=('batch', 'x', 'y')`.
1015 To clip samples independently, leave out the 'batch' axis.
1017 Only valid if `min_percentile` and/or `max_percentile` are set.
1019 Default: Compute percentiles over all axes jointly."""
1021 @model_validator(mode="after")
1022 def _validate(self) -> Self:
1023 if (self.min is not None) and (self.min_percentile is not None):
1024 raise ValueError(
1025 "Only one of `min` and `min_percentile` may be set, not both."
1026 )
1027 if (self.max is not None) and (self.max_percentile is not None):
1028 raise ValueError(
1029 "Only one of `max` and `max_percentile` may be set, not both."
1030 )
1031 if (
1032 self.min is None
1033 and self.min_percentile is None
1034 and self.max is None
1035 and self.max_percentile is None
1036 ):
1037 raise ValueError(
1038 "At least one of `min`, `min_percentile`, `max`, or `max_percentile` must be set."
1039 )
1041 if (
1042 self.axes is not None
1043 and self.min_percentile is None
1044 and self.max_percentile is None
1045 ):
1046 raise ValueError(
1047 "If `axes` is set, at least one of `min_percentile` or `max_percentile` must be set."
1048 )
1050 return self
1053class ClipDescr(NodeWithExplicitlySetFields):
1054 """Set tensor values below min to min and above max to max.
1056 See `ScaleRangeDescr` for examples.
1057 """
1059 implemented_id: ClassVar[Literal["clip"]] = "clip"
1060 if TYPE_CHECKING:
1061 id: Literal["clip"] = "clip"
1062 else:
1063 id: Literal["clip"]
1065 kwargs: ClipKwargs
1068class EnsureDtypeKwargs(KwargsNode):
1069 """key word arguments for [EnsureDtypeDescr][]"""
1071 dtype: Literal[
1072 "float32",
1073 "float64",
1074 "uint8",
1075 "int8",
1076 "uint16",
1077 "int16",
1078 "uint32",
1079 "int32",
1080 "uint64",
1081 "int64",
1082 "bool",
1083 ]
1086class EnsureDtypeDescr(NodeWithExplicitlySetFields):
1087 """Cast the tensor data type to `EnsureDtypeKwargs.dtype` (if not matching).
1089 This can for example be used to ensure the inner neural network model gets a
1090 different input tensor data type than the fully described bioimage.io model does.
1092 Examples:
1093 The described bioimage.io model (incl. preprocessing) accepts any
1094 float32-compatible tensor, normalizes it with percentiles and clipping and then
1095 casts it to uint8, which is what the neural network in this example expects.
1096 - in YAML
1097 ```yaml
1098 inputs:
1099 - data:
1100 type: float32 # described bioimage.io model is compatible with any float32 input tensor
1101 preprocessing:
1102 - id: scale_range
1103 kwargs:
1104 axes: ['y', 'x']
1105 max_percentile: 99.8
1106 min_percentile: 5.0
1107 - id: clip
1108 kwargs:
1109 min: 0.0
1110 max: 1.0
1111 - id: ensure_dtype # the neural network of the model requires uint8
1112 kwargs:
1113 dtype: uint8
1114 ```
1115 - in Python:
1116 >>> preprocessing = [
1117 ... ScaleRangeDescr(
1118 ... kwargs=ScaleRangeKwargs(
1119 ... axes= (AxisId('y'), AxisId('x')),
1120 ... max_percentile= 99.8,
1121 ... min_percentile= 5.0,
1122 ... )
1123 ... ),
1124 ... ClipDescr(kwargs=ClipKwargs(min=0.0, max=1.0)),
1125 ... EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype="uint8")),
1126 ... ]
1127 """
1129 implemented_id: ClassVar[Literal["ensure_dtype"]] = "ensure_dtype"
1130 if TYPE_CHECKING:
1131 id: Literal["ensure_dtype"] = "ensure_dtype"
1132 else:
1133 id: Literal["ensure_dtype"]
1135 kwargs: EnsureDtypeKwargs
1138class ScaleLinearKwargs(KwargsNode):
1139 """Key word arguments for [ScaleLinearDescr][]"""
1141 gain: float = 1.0
1142 """multiplicative factor"""
1144 offset: float = 0.0
1145 """additive term"""
1147 @model_validator(mode="after")
1148 def _validate(self) -> Self:
1149 if self.gain == 1.0 and self.offset == 0.0:
1150 raise ValueError(
1151 "Redundant linear scaling not allowd. Set `gain` != 1.0 and/or `offset`"
1152 + " != 0.0."
1153 )
1155 return self
1158class ScaleLinearAlongAxisKwargs(KwargsNode):
1159 """Key word arguments for [ScaleLinearDescr][]"""
1161 axis: Annotated[NonBatchAxisId, Field(examples=["channel"])]
1162 """The axis of gain and offset values."""
1164 gain: Union[float, NotEmpty[List[float]]] = 1.0
1165 """multiplicative factor"""
1167 offset: Union[float, NotEmpty[List[float]]] = 0.0
1168 """additive term"""
1170 @model_validator(mode="after")
1171 def _validate(self) -> Self:
1172 if isinstance(self.gain, list):
1173 if isinstance(self.offset, list):
1174 if len(self.gain) != len(self.offset):
1175 raise ValueError(
1176 f"Size of `gain` ({len(self.gain)}) and `offset` ({len(self.offset)}) must match."
1177 )
1178 else:
1179 self.offset = [float(self.offset)] * len(self.gain)
1180 elif isinstance(self.offset, list):
1181 self.gain = [float(self.gain)] * len(self.offset)
1182 else:
1183 raise ValueError(
1184 "Do not specify an `axis` for scalar gain and offset values."
1185 )
1187 if all(g == 1.0 for g in self.gain) and all(off == 0.0 for off in self.offset):
1188 raise ValueError(
1189 "Redundant linear scaling not allowd. Set `gain` != 1.0 and/or `offset`"
1190 + " != 0.0."
1191 )
1193 return self
1196class ScaleLinearDescr(NodeWithExplicitlySetFields):
1197 """Fixed linear scaling.
1199 Examples:
1200 1. Scale with scalar gain and offset
1201 - in YAML
1202 ```yaml
1203 preprocessing:
1204 - id: scale_linear
1205 kwargs:
1206 gain: 2.0
1207 offset: 3.0
1208 ```
1209 - in Python:
1211 >>> preprocessing = [
1212 ... ScaleLinearDescr(kwargs=ScaleLinearKwargs(gain= 2.0, offset=3.0))
1213 ... ]
1215 2. Independent scaling along an axis
1216 - in YAML
1217 ```yaml
1218 preprocessing:
1219 - id: scale_linear
1220 kwargs:
1221 axis: 'channel'
1222 gain: [1.0, 2.0, 3.0]
1223 ```
1224 - in Python:
1226 >>> preprocessing = [
1227 ... ScaleLinearDescr(
1228 ... kwargs=ScaleLinearAlongAxisKwargs(
1229 ... axis=AxisId("channel"),
1230 ... gain=[1.0, 2.0, 3.0],
1231 ... )
1232 ... )
1233 ... ]
1235 """
1237 implemented_id: ClassVar[Literal["scale_linear"]] = "scale_linear"
1238 if TYPE_CHECKING:
1239 id: Literal["scale_linear"] = "scale_linear"
1240 else:
1241 id: Literal["scale_linear"]
1242 kwargs: Union[ScaleLinearKwargs, ScaleLinearAlongAxisKwargs]
1245class SigmoidDescr(NodeWithExplicitlySetFields):
1246 """The logistic sigmoid function, a.k.a. expit function.
1248 Examples:
1249 - in YAML
1250 ```yaml
1251 postprocessing:
1252 - id: sigmoid
1253 ```
1254 - in Python:
1256 >>> postprocessing = [SigmoidDescr()]
1257 """
1259 implemented_id: ClassVar[Literal["sigmoid"]] = "sigmoid"
1260 if TYPE_CHECKING:
1261 id: Literal["sigmoid"] = "sigmoid"
1262 else:
1263 id: Literal["sigmoid"]
1265 @property
1266 def kwargs(self) -> KwargsNode:
1267 """empty kwargs"""
1268 return KwargsNode()
1271class SoftmaxKwargs(KwargsNode):
1272 """key word arguments for [SoftmaxDescr][]"""
1274 axis: Annotated[NonBatchAxisId, Field(examples=["channel"])] = AxisId("channel")
1275 """The axis to apply the softmax function along.
1276 Note:
1277 Defaults to 'channel' axis
1278 (which may not exist, in which case
1279 a different axis id has to be specified).
1280 """
1283class SoftmaxDescr(NodeWithExplicitlySetFields):
1284 """The softmax function.
1286 Examples:
1287 - in YAML
1288 ```yaml
1289 postprocessing:
1290 - id: softmax
1291 kwargs:
1292 axis: channel
1293 ```
1294 - in Python:
1296 >>> postprocessing = [SoftmaxDescr(kwargs=SoftmaxKwargs(axis=AxisId("channel")))]
1297 """
1299 implemented_id: ClassVar[Literal["softmax"]] = "softmax"
1300 if TYPE_CHECKING:
1301 id: Literal["softmax"] = "softmax"
1302 else:
1303 id: Literal["softmax"]
1305 kwargs: SoftmaxKwargs = Field(default_factory=SoftmaxKwargs.model_construct)
1308class _StardistPostprocessingKwargsBase(KwargsNode):
1309 """key word arguments for [StardistPostprocessingDescr][]"""
1311 prob_threshold: float
1312 """The probability threshold for object candidate selection."""
1314 nms_threshold: float
1315 """The IoU threshold for non-maximum suppression."""
1317 n_rays: int
1318 """Number of radial lines (rays) cast from the center of an object to its boundary."""
1321class StardistPostprocessingKwargs2D(_StardistPostprocessingKwargsBase):
1322 grid: Tuple[int, int]
1323 """Grid size of network predictions."""
1325 b: Union[int, Tuple[Tuple[int, int], Tuple[int, int]]]
1326 """Border region in which object probability is set to zero."""
1329class StardistPostprocessingKwargs3D(_StardistPostprocessingKwargsBase):
1330 grid: Tuple[int, int, int]
1331 """Grid size of network predictions."""
1333 b: Union[int, Tuple[Tuple[int, int], Tuple[int, int], Tuple[int, int]]]
1334 """Border region in which object probability is set to zero."""
1336 anisotropy: Tuple[float, float, float]
1337 """Anisotropy factors for 3D star-convex polyhedra, i.e. the physical pixel size along each spatial axis."""
1339 overlap_label: Optional[int] = None
1340 """Optional label to apply to any area of overlapping predicted objects."""
1343class StardistPostprocessingDescr(NodeWithExplicitlySetFields):
1344 """Stardist postprocessing including non-maximum suppression and converting polygon representations to instance labels
1346 as described in:
1347 - Uwe Schmidt, Martin Weigert, Coleman Broaddus, and Gene Myers.
1348 [*Cell Detection with Star-convex Polygons*](https://arxiv.org/abs/1806.03535).
1349 International Conference on Medical Image Computing and Computer-Assisted Intervention (MICCAI), Granada, Spain, September 2018.
1350 - Martin Weigert, Uwe Schmidt, Robert Haase, Ko Sugawara, and Gene Myers.
1351 [*Star-convex Polyhedra for 3D Object Detection and Segmentation in Microscopy*](http://openaccess.thecvf.com/content_WACV_2020/papers/Weigert_Star-convex_Polyhedra_for_3D_Object_Detection_and_Segmentation_in_Microscopy_WACV_2020_paper.pdf).
1352 The IEEE Winter Conference on Applications of Computer Vision (WACV), Snowmass Village, Colorado, March 2020.
1354 Note: Only available if the `stardist` package is installed.
1355 """
1357 implemented_id: ClassVar[Literal["stardist_postprocessing"]] = (
1358 "stardist_postprocessing"
1359 )
1360 if TYPE_CHECKING:
1361 id: Literal["stardist_postprocessing"] = "stardist_postprocessing"
1362 else:
1363 id: Literal["stardist_postprocessing"]
1365 kwargs: Union[StardistPostprocessingKwargs2D, StardistPostprocessingKwargs3D]
1368class CellposeFlowDynamicsKwargs(KwargsNode):
1369 """key word arguments for [CellposeFlowDynamicsDescr][]"""
1371 cellprob_threshold: float
1372 flow_threshold: float
1373 do_3D: bool
1374 min_size: int = 15
1375 """Minimum size of objects to keep, in pixels. Default is 15, which is the default in Cellpose. Set to 0 to disable filtering by size."""
1376 output_dtype: Literal["uint16", "uint32"] = "uint16"
1379class CellposeFlowDynamicsDescr(NodeWithExplicitlySetFields):
1380 """Cellpose flow dynamics postprocessing as described in:
1381 - Carsen Stringer and Marius Pachitariu. [*Cellpose: a generalist algorithm for cellular segmentation*](https://www.nature.com/articles/s41592-020-01018-x). Nature Methods, 2021.
1383 Note: Only available if the `cellpose` package is installed.
1384 """
1386 implemented_id: ClassVar[Literal["cellpose_flow_dynamics"]] = (
1387 "cellpose_flow_dynamics"
1388 )
1389 if TYPE_CHECKING:
1390 id: Literal["cellpose_flow_dynamics"] = "cellpose_flow_dynamics"
1391 else:
1392 id: Literal["cellpose_flow_dynamics"]
1394 kwargs: CellposeFlowDynamicsKwargs
1397class CustomProcessingDescr(NodeWithExplicitlySetFields, FileDescr):
1398 """Custom (post)processing op — source file shipped inline with the model.
1400 Supports (post)processing that cannot be expressed by the built-in named
1401 operations (watershed, connected components, etc.)
1402 using a simple Python callable interface.
1404 The op is implemented in a ``.py`` file packaged alongside the model weights.
1405 Two styles are supported:
1407 *Callable class* — kwargs go to ``__init__``, tensors arrive in ``__call__``:
1409 .. code-block:: python
1411 # my_postprocess.py
1412 import numpy as np
1414 class my_postprocess:
1415 def __init__(self, threshold: float = 0.5) -> None:
1416 self.threshold = threshold
1417 def __call__(self, *arrays: np.ndarray) -> np.ndarray:
1418 # arrays = model output tensors in rdf.yaml declaration order
1419 return (arrays[0] > self.threshold).astype(np.uint8)
1421 *Factory function* — alternative closure style, identical runtime behaviour:
1423 .. code-block:: python
1425 # my_postprocess.py
1426 import numpy as np
1428 def my_postprocess(threshold: float = 0.5):
1429 def run(*arrays: np.ndarray) -> np.ndarray:
1430 return (arrays[0] > threshold).astype(np.uint8)
1431 return run
1433 Reference it in ``rdf.yaml`` with the source file included in the package:
1435 .. code-block:: yaml
1437 postprocessing:
1438 - id: custom
1439 callable: my_postprocess # class or function name in source
1440 source: my_postprocess.py # packaged alongside weights
1441 sha256: <hash> # sha256 of the source file
1442 kwargs: # forwarded to __init__ / factory
1443 threshold: 0.5
1445 **Security:** source files are SHA-256 verified before execution.
1446 Execution requires explicit opt-in in bioimageio.core and curator
1447 review before Zoo publication.
1448 """
1450 implemented_id: ClassVar[Literal["custom"]] = "custom"
1451 if TYPE_CHECKING:
1452 id: Literal["custom"] = "custom"
1453 else:
1454 id: Literal["custom"]
1456 callable: Annotated[
1457 str,
1458 Field(examples=["my_postprocess_factory", "MyPostprocessClass"]),
1459 ]
1460 """Name of the callable class or factory function defined in ``source``.
1462 At runtime: ``op = callable(**kwargs)``, then ``result = op(*output_tensors)``
1463 per image. Both a class with ``__call__`` and a factory function returning
1464 a callable satisfy this protocol."""
1466 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
1467 """Python source file (included when packaging the model)."""
1469 kwargs: Dict[str, YamlValue] = Field(
1470 default_factory=cast(Callable[[], Dict[str, YamlValue]], dict)
1471 )
1472 """Keyword arguments forwarded to the callable (``__init__`` or factory)."""
1474 @model_serializer(mode="wrap", when_used="unless-none")
1475 def _serialize(
1476 self, nxt: SerializerFunctionWrapHandler, info: SerializationInfo
1477 ) -> Dict[str, YamlValue]:
1478 return package_file_descr_serializer(self, nxt, info)
1481class FixedZeroMeanUnitVarianceKwargs(KwargsNode):
1482 """key word arguments for [FixedZeroMeanUnitVarianceDescr][]"""
1484 mean: float
1485 """The mean value to normalize with."""
1487 std: Annotated[float, Ge(1e-6)]
1488 """The standard deviation value to normalize with."""
1491class FixedZeroMeanUnitVarianceAlongAxisKwargs(KwargsNode):
1492 """key word arguments for [FixedZeroMeanUnitVarianceDescr][]"""
1494 mean: NotEmpty[List[float]]
1495 """The mean value(s) to normalize with."""
1497 std: NotEmpty[List[Annotated[float, Ge(1e-6)]]]
1498 """The standard deviation value(s) to normalize with.
1499 Size must match `mean` values."""
1501 axis: Annotated[NonBatchAxisId, Field(examples=["channel", "index"])]
1502 """The axis of the mean/std values to normalize each entry along that dimension
1503 separately."""
1505 @model_validator(mode="after")
1506 def _mean_and_std_match(self) -> Self:
1507 if len(self.mean) != len(self.std):
1508 raise ValueError(
1509 f"Size of `mean` ({len(self.mean)}) and `std` ({len(self.std)})"
1510 + " must match."
1511 )
1513 return self
1516class FixedZeroMeanUnitVarianceDescr(NodeWithExplicitlySetFields):
1517 """Subtract a given mean and divide by the standard deviation.
1519 Normalize with fixed, precomputed values for
1520 `FixedZeroMeanUnitVarianceKwargs.mean` and `FixedZeroMeanUnitVarianceKwargs.std`
1521 Use `FixedZeroMeanUnitVarianceAlongAxisKwargs` for independent scaling along given
1522 axes.
1524 Examples:
1525 1. scalar value for whole tensor
1526 - in YAML
1527 ```yaml
1528 preprocessing:
1529 - id: fixed_zero_mean_unit_variance
1530 kwargs:
1531 mean: 103.5
1532 std: 13.7
1533 ```
1534 - in Python
1535 >>> preprocessing = [FixedZeroMeanUnitVarianceDescr(
1536 ... kwargs=FixedZeroMeanUnitVarianceKwargs(mean=103.5, std=13.7)
1537 ... )]
1539 2. independently along an axis
1540 - in YAML
1541 ```yaml
1542 preprocessing:
1543 - id: fixed_zero_mean_unit_variance
1544 kwargs:
1545 axis: channel
1546 mean: [101.5, 102.5, 103.5]
1547 std: [11.7, 12.7, 13.7]
1548 ```
1549 - in Python
1550 >>> preprocessing = [FixedZeroMeanUnitVarianceDescr(
1551 ... kwargs=FixedZeroMeanUnitVarianceAlongAxisKwargs(
1552 ... axis=AxisId("channel"),
1553 ... mean=[101.5, 102.5, 103.5],
1554 ... std=[11.7, 12.7, 13.7],
1555 ... )
1556 ... )]
1557 """
1559 implemented_id: ClassVar[Literal["fixed_zero_mean_unit_variance"]] = (
1560 "fixed_zero_mean_unit_variance"
1561 )
1562 if TYPE_CHECKING:
1563 id: Literal["fixed_zero_mean_unit_variance"] = "fixed_zero_mean_unit_variance"
1564 else:
1565 id: Literal["fixed_zero_mean_unit_variance"]
1567 kwargs: Union[
1568 FixedZeroMeanUnitVarianceKwargs, FixedZeroMeanUnitVarianceAlongAxisKwargs
1569 ]
1572class ZeroMeanUnitVarianceKwargs(KwargsNode):
1573 """key word arguments for [ZeroMeanUnitVarianceDescr][]"""
1575 axes: Annotated[
1576 Optional[Sequence[AxisId]], Field(examples=[("batch", "x", "y")])
1577 ] = None
1578 """The subset of axes to normalize jointly, i.e. axes to reduce to compute mean/std.
1579 For example to normalize 'batch', 'x' and 'y' jointly in a tensor ('batch', 'channel', 'y', 'x')
1580 resulting in a tensor of equal shape normalized per channel, specify `axes=('batch', 'x', 'y')`.
1581 To normalize each sample independently leave out the 'batch' axis.
1582 Default: Scale all axes jointly."""
1584 eps: Annotated[float, Interval(gt=0, le=0.1)] = 1e-6
1585 """epsilon for numeric stability: `out = (tensor - mean) / (std + eps)`."""
1588class ZeroMeanUnitVarianceDescr(NodeWithExplicitlySetFields):
1589 """Subtract mean and divide by variance.
1591 Examples:
1592 Subtract tensor mean and variance
1593 - in YAML
1594 ```yaml
1595 preprocessing:
1596 - id: zero_mean_unit_variance
1597 ```
1598 - in Python
1599 >>> preprocessing = [ZeroMeanUnitVarianceDescr()]
1600 """
1602 implemented_id: ClassVar[Literal["zero_mean_unit_variance"]] = (
1603 "zero_mean_unit_variance"
1604 )
1605 if TYPE_CHECKING:
1606 id: Literal["zero_mean_unit_variance"] = "zero_mean_unit_variance"
1607 else:
1608 id: Literal["zero_mean_unit_variance"]
1610 kwargs: ZeroMeanUnitVarianceKwargs = Field(
1611 default_factory=ZeroMeanUnitVarianceKwargs.model_construct
1612 )
1615class ScaleRangeKwargs(KwargsNode):
1616 """key word arguments for [ScaleRangeDescr][]
1618 For `min_percentile`=0.0 (the default) and `max_percentile`=100 (the default)
1619 this processing step normalizes data to the [0, 1] intervall.
1620 For other percentiles the normalized values will partially be outside the [0, 1]
1621 intervall. Use `ScaleRange` followed by `ClipDescr` if you want to limit the
1622 normalized values to a range.
1623 """
1625 axes: Annotated[
1626 Optional[Sequence[AxisId]], Field(examples=[("batch", "x", "y")])
1627 ] = None
1628 """The subset of axes to normalize jointly, i.e. axes to reduce to compute the min/max percentile value.
1629 For example to normalize 'batch', 'x' and 'y' jointly in a tensor ('batch', 'channel', 'y', 'x')
1630 resulting in a tensor of equal shape normalized per channel, specify `axes=('batch', 'x', 'y')`.
1631 To normalize samples independently, leave out the "batch" axis.
1632 Default: Scale all axes jointly."""
1634 min_percentile: Annotated[float, Interval(ge=0, lt=100)] = 0.0
1635 """The lower percentile used to determine the value to align with zero."""
1637 max_percentile: Annotated[float, Interval(gt=1, le=100)] = 100.0
1638 """The upper percentile used to determine the value to align with one.
1639 Has to be bigger than `min_percentile`.
1640 The range is 1 to 100 instead of 0 to 100 to avoid mistakenly
1641 accepting percentiles specified in the range 0.0 to 1.0."""
1643 eps: Annotated[float, Interval(gt=0, le=0.1)] = 1e-6
1644 """Epsilon for numeric stability.
1645 `out = (tensor - v_lower) / (v_upper - v_lower + eps)`;
1646 with `v_lower,v_upper` values at the respective percentiles."""
1648 reference_tensor: Optional[TensorId] = None
1649 """ID of the unprocessed input tensor to compute the percentiles from.
1650 Default: The tensor itself.
1651 """
1653 @field_validator("max_percentile", mode="after")
1654 @classmethod
1655 def min_smaller_max(cls, value: float, info: ValidationInfo) -> float:
1656 if (min_p := info.data["min_percentile"]) >= value:
1657 raise ValueError(f"min_percentile {min_p} >= max_percentile {value}")
1659 return value
1662class ScaleRangeDescr(NodeWithExplicitlySetFields):
1663 """Scale with percentiles.
1665 Examples:
1666 1. Scale linearly to map 5th percentile to 0 and 99.8th percentile to 1.0
1667 - in YAML
1668 ```yaml
1669 preprocessing:
1670 - id: scale_range
1671 kwargs:
1672 axes: ['y', 'x']
1673 max_percentile: 99.8
1674 min_percentile: 5.0
1675 ```
1676 - in Python
1678 >>> preprocessing = [
1679 ... ScaleRangeDescr(
1680 ... kwargs=ScaleRangeKwargs(
1681 ... axes= (AxisId('y'), AxisId('x')),
1682 ... max_percentile= 99.8,
1683 ... min_percentile= 5.0,
1684 ... )
1685 ... )
1686 ... ]
1688 2. Combine the above scaling with additional clipping to clip values outside the range given by the percentiles.
1689 - in YAML
1690 ```yaml
1691 preprocessing:
1692 - id: scale_range
1693 kwargs:
1694 axes: ['y', 'x']
1695 max_percentile: 99.8
1696 min_percentile: 5.0
1697 - id: clip
1698 kwargs:
1699 min: 0.0
1700 max: 1.0
1701 ```
1702 - in Python
1704 >>> preprocessing = [
1705 ... ScaleRangeDescr(
1706 ... kwargs=ScaleRangeKwargs(
1707 ... axes= (AxisId('y'), AxisId('x')),
1708 ... max_percentile= 99.8,
1709 ... min_percentile= 5.0,
1710 ... )
1711 ... ),
1712 ... ClipDescr(
1713 ... kwargs=ClipKwargs(
1714 ... min=0.0,
1715 ... max=1.0,
1716 ... )
1717 ... ),
1718 ... ]
1720 """
1722 implemented_id: ClassVar[Literal["scale_range"]] = "scale_range"
1723 if TYPE_CHECKING:
1724 id: Literal["scale_range"] = "scale_range"
1725 else:
1726 id: Literal["scale_range"]
1727 kwargs: ScaleRangeKwargs = Field(default_factory=ScaleRangeKwargs.model_construct)
1730class ScaleMeanVarianceKwargs(KwargsNode):
1731 """key word arguments for [ScaleMeanVarianceKwargs][]"""
1733 reference_tensor: TensorId
1734 """ID of unprocessed input tensor to match."""
1736 axes: Annotated[
1737 Optional[Sequence[AxisId]], Field(examples=[("batch", "x", "y")])
1738 ] = None
1739 """The subset of axes to normalize jointly, i.e. axes to reduce to compute mean/std.
1740 For example to normalize 'batch', 'x' and 'y' jointly in a tensor ('batch', 'channel', 'y', 'x')
1741 resulting in a tensor of equal shape normalized per channel, specify `axes=('batch', 'x', 'y')`.
1742 To normalize samples independently, leave out the 'batch' axis.
1743 Default: Scale all axes jointly."""
1745 eps: Annotated[float, Interval(gt=0, le=0.1)] = 1e-6
1746 """Epsilon for numeric stability:
1747 `out = (tensor - mean) / (std + eps) * (ref_std + eps) + ref_mean.`"""
1750class ScaleMeanVarianceDescr(NodeWithExplicitlySetFields):
1751 """Scale a tensor's data distribution to match another tensor's mean/std.
1752 `out = (tensor - mean) / (std + eps) * (ref_std + eps) + ref_mean.`
1753 """
1755 implemented_id: ClassVar[Literal["scale_mean_variance"]] = "scale_mean_variance"
1756 if TYPE_CHECKING:
1757 id: Literal["scale_mean_variance"] = "scale_mean_variance"
1758 else:
1759 id: Literal["scale_mean_variance"]
1760 kwargs: ScaleMeanVarianceKwargs
1763PreprocessingDescr = Annotated[
1764 Union[
1765 BinarizeDescr,
1766 ClipDescr,
1767 EnsureDtypeDescr,
1768 FixedZeroMeanUnitVarianceDescr,
1769 ScaleLinearDescr,
1770 ScaleRangeDescr,
1771 SigmoidDescr,
1772 SoftmaxDescr,
1773 ZeroMeanUnitVarianceDescr,
1774 ],
1775 Discriminator("id"),
1776]
1777PostprocessingDescr = Annotated[
1778 Union[
1779 BinarizeDescr,
1780 CellposeFlowDynamicsDescr,
1781 ClipDescr,
1782 CustomProcessingDescr,
1783 EnsureDtypeDescr,
1784 FixedZeroMeanUnitVarianceDescr,
1785 ScaleLinearDescr,
1786 ScaleMeanVarianceDescr,
1787 ScaleRangeDescr,
1788 SigmoidDescr,
1789 SoftmaxDescr,
1790 StardistPostprocessingDescr,
1791 ZeroMeanUnitVarianceDescr,
1792 ],
1793 Discriminator("id"),
1794]
1796IO_AxisT = TypeVar("IO_AxisT", InputAxis, OutputAxis)
1799class TensorDescrBase(Node, Generic[IO_AxisT]):
1800 id: TensorId
1801 """Tensor id. No duplicates are allowed."""
1803 description: Annotated[str, MaxLen(128)] = ""
1804 """free text description"""
1806 axes: NotEmpty[Sequence[IO_AxisT]]
1807 """tensor axes"""
1809 @property
1810 def shape(self):
1811 return tuple(a.size for a in self.axes)
1813 @field_validator("axes", mode="after", check_fields=False)
1814 @classmethod
1815 def _validate_axes(cls, axes: Sequence[AnyAxis]) -> Sequence[AnyAxis]:
1816 batch_axes = [a for a in axes if a.type == "batch"]
1817 if len(batch_axes) > 1:
1818 raise ValueError(
1819 f"Only one batch axis (per tensor) allowed, but got {batch_axes}"
1820 )
1822 seen_ids: Set[AxisId] = set()
1823 duplicate_axes_ids: Set[AxisId] = set()
1824 for a in axes:
1825 (duplicate_axes_ids if a.id in seen_ids else seen_ids).add(a.id)
1827 if duplicate_axes_ids:
1828 raise ValueError(f"Duplicate axis ids: {duplicate_axes_ids}")
1830 return axes
1832 test_tensor: FAIR[Optional[FileDescr_]] = None
1833 """An example tensor to use for testing.
1834 Using the model with the test input tensors is expected to yield the test output tensors.
1835 Each test tensor has be a an ndarray in the
1836 [numpy.lib file format](https://numpy.org/doc/stable/reference/generated/numpy.lib.format.html#module-numpy.lib.format).
1837 The file extension must be '.npy'."""
1839 sample_tensor: FAIR[Optional[FileDescr_]] = None
1840 """A sample tensor to illustrate a possible input/output for the model,
1841 The sample image primarily serves to inform a human user about an example use case
1842 and is typically stored as .hdf5, .png or .tiff.
1843 It has to be readable by the [imageio library](https://imageio.readthedocs.io/en/stable/formats/index.html#supported-formats)
1844 (numpy's `.npy` format is not supported).
1845 The image dimensionality has to match the number of axes specified in this tensor description.
1846 """
1848 @model_validator(mode="after")
1849 def _validate_sample_tensor(self) -> Self:
1850 if self.sample_tensor is None or not get_validation_context().perform_io_checks:
1851 return self
1853 reader = get_reader(self.sample_tensor.source, sha256=self.sample_tensor.sha256)
1854 tensor: NDArray[Any] = imread( # pyright: ignore[reportUnknownVariableType]
1855 reader.read(),
1856 extension=PurePosixPath(reader.original_file_name).suffix,
1857 )
1858 n_dims = len(tensor.squeeze().shape)
1859 n_dims_min = n_dims_max = len(self.axes)
1861 for a in self.axes:
1862 if isinstance(a, BatchAxis):
1863 n_dims_min -= 1
1864 elif isinstance(a.size, int):
1865 if a.size == 1:
1866 n_dims_min -= 1
1867 elif isinstance(a.size, (ParameterizedSize, DataDependentSize)):
1868 if a.size.min == 1:
1869 n_dims_min -= 1
1870 elif isinstance(a.size, SizeReference):
1871 if a.size.offset < 2:
1872 # size reference may result in singleton axis
1873 n_dims_min -= 1
1874 else:
1875 assert_never(a.size)
1877 n_dims_min = max(0, n_dims_min)
1878 if n_dims < n_dims_min or n_dims > n_dims_max:
1879 raise ValueError(
1880 f"Expected sample tensor to have {n_dims_min} to"
1881 + f" {n_dims_max} dimensions, but found {n_dims} (shape: {tensor.shape})."
1882 )
1884 return self
1886 data: Union[TensorDataDescr, NotEmpty[Sequence[TensorDataDescr]]] = (
1887 IntervalOrRatioDataDescr()
1888 )
1889 """Description of the tensor's data values, optionally per channel.
1890 If specified per channel, the data `type` needs to match across channels."""
1892 @property
1893 def dtype(
1894 self,
1895 ) -> Literal[
1896 "float32",
1897 "float64",
1898 "uint8",
1899 "int8",
1900 "uint16",
1901 "int16",
1902 "uint32",
1903 "int32",
1904 "uint64",
1905 "int64",
1906 "bool",
1907 ]:
1908 """dtype as specified under `data.type` or `data[i].type`"""
1909 if isinstance(self.data, collections.abc.Sequence):
1910 return self.data[0].type
1911 else:
1912 return self.data.type
1914 @field_validator("data", mode="after")
1915 @classmethod
1916 def _check_data_type_across_channels(
1917 cls, value: Union[TensorDataDescr, NotEmpty[Sequence[TensorDataDescr]]]
1918 ) -> Union[TensorDataDescr, NotEmpty[Sequence[TensorDataDescr]]]:
1919 if not isinstance(value, list):
1920 return value
1922 dtypes = {t.type for t in value}
1923 if len(dtypes) > 1:
1924 raise ValueError(
1925 "Tensor data descriptions per channel need to agree in their data"
1926 + f" `type`, but found {dtypes}."
1927 )
1929 return value
1931 @model_validator(mode="after")
1932 def _check_data_matches_channelaxis(self) -> Self:
1933 if not isinstance(self.data, (list, tuple)):
1934 return self
1936 for a in self.axes:
1937 if isinstance(a, ChannelAxis):
1938 size = a.size
1939 assert isinstance(size, int)
1940 break
1941 else:
1942 return self
1944 if len(self.data) != size:
1945 raise ValueError(
1946 f"Got tensor data descriptions for {len(self.data)} channels, but"
1947 + f" '{a.id}' axis has size {size}."
1948 )
1950 return self
1952 def get_axis_sizes_for_array(self, array: NDArray[Any]) -> Dict[AxisId, int]:
1953 if len(array.shape) != len(self.axes):
1954 raise ValueError(
1955 f"Dimension mismatch: array shape {array.shape} (#{len(array.shape)})"
1956 + f" incompatible with {len(self.axes)} axes."
1957 )
1958 return {a.id: array.shape[i] for i, a in enumerate(self.axes)}
1961class ConstantPadding(Node):
1962 mode: Literal["constant"] = "constant"
1963 value: Union[int, float] = 0
1966class EdgePadding(Node):
1967 mode: Literal["edge"] = "edge"
1970class ReflectPadding(Node):
1971 mode: Literal["reflect"] = "reflect"
1974class SymmetricPadding(Node):
1975 mode: Literal["symmetric"] = "symmetric"
1978Padding = Union[ConstantPadding, EdgePadding, ReflectPadding, SymmetricPadding]
1981class InputTensorDescr(TensorDescrBase[InputAxis]):
1982 id: TensorId = TensorId("input")
1983 """Input tensor id.
1984 No duplicates are allowed across all inputs and outputs."""
1986 optional: bool = False
1987 """indicates that this tensor may be `None`"""
1989 pad: Optional[Padding] = None
1990 """Explicitly specify how to pad this input tensor.
1992 Use `axes[i].pad` to specify padding width.
1994 Note:
1995 Non-blockwise sample prediction only applies padding for axes with a `pad` specification.
1996 """
1998 preprocessing: List[PreprocessingDescr] = Field(
1999 default_factory=cast(Callable[[], List[PreprocessingDescr]], list)
2000 )
2001 """Description of how this input should be preprocessed.
2003 notes:
2004 - If preprocessing does not start with an 'ensure_dtype' entry, it is added
2005 to ensure an input tensor's data type matches the input tensor's data description.
2006 - If preprocessing does not end with an 'ensure_dtype' or 'binarize' entry, an
2007 'ensure_dtype' step is added to ensure preprocessing steps are not unintentionally
2008 changing the data type.
2009 """
2011 @model_validator(mode="after")
2012 def _validate_preprocessing_kwargs(self) -> Self:
2013 axes_ids = [a.id for a in self.axes]
2014 for p in self.preprocessing:
2015 kwargs_axes: Optional[Sequence[Any]] = p.kwargs.get("axes")
2016 if kwargs_axes is None:
2017 continue
2019 if not isinstance(kwargs_axes, collections.abc.Sequence):
2020 raise ValueError(
2021 f"Expected `preprocessing.i.kwargs.axes` to be a sequence, but got {type(kwargs_axes)}"
2022 )
2024 if any(a not in axes_ids for a in kwargs_axes):
2025 raise ValueError(
2026 "`preprocessing.i.kwargs.axes` needs to be subset of axes ids"
2027 )
2029 if isinstance(self.data, (NominalOrOrdinalDataDescr, IntervalOrRatioDataDescr)):
2030 dtype = self.data.type
2031 else:
2032 dtype = self.data[0].type
2034 # ensure `preprocessing` begins with `EnsureDtypeDescr`
2035 if not self.preprocessing or not isinstance(
2036 self.preprocessing[0], EnsureDtypeDescr
2037 ):
2038 self.preprocessing.insert(
2039 0, EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype=dtype))
2040 )
2042 # ensure `preprocessing` ends with `EnsureDtypeDescr` or `BinarizeDescr`
2043 if not isinstance(self.preprocessing[-1], (EnsureDtypeDescr, BinarizeDescr)):
2044 self.preprocessing.append(
2045 EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype=dtype))
2046 )
2048 return self
2051def convert_axes(
2052 axes: str,
2053 *,
2054 shape: Union[
2055 Sequence[int], _ParameterizedInputShape_v0_4, _ImplicitOutputShape_v0_4
2056 ],
2057 tensor_type: Literal["input", "output"],
2058 halo: Optional[Sequence[int]],
2059 size_refs: Mapping[_TensorName_v0_4, Mapping[str, int]],
2060):
2061 ret: List[AnyAxis] = []
2062 for i, a in enumerate(axes):
2063 axis_type = _AXIS_TYPE_MAP.get(a, a)
2064 if axis_type == "batch":
2065 ret.append(BatchAxis())
2066 continue
2068 scale = 1.0
2069 if isinstance(shape, _ParameterizedInputShape_v0_4):
2070 if shape.step[i] == 0:
2071 size = shape.min[i]
2072 else:
2073 size = ParameterizedSize(min=shape.min[i], step=shape.step[i])
2074 elif isinstance(shape, _ImplicitOutputShape_v0_4):
2075 ref_t = str(shape.reference_tensor)
2076 if ref_t.count(".") == 1:
2077 t_id, orig_a_id = ref_t.split(".")
2078 else:
2079 t_id = ref_t
2080 orig_a_id = a
2082 a_id = _AXIS_ID_MAP.get(orig_a_id, a)
2083 if not (orig_scale := shape.scale[i]):
2084 # old way to insert a new axis dimension
2085 size = int(2 * shape.offset[i])
2086 else:
2087 scale = 1 / orig_scale
2088 if axis_type in ("channel", "index"):
2089 # these axes no longer have a scale
2090 offset_from_scale = orig_scale * size_refs.get(
2091 _TensorName_v0_4(t_id), {}
2092 ).get(orig_a_id, 0)
2093 else:
2094 offset_from_scale = 0
2095 size = SizeReference(
2096 tensor_id=TensorId(t_id),
2097 axis_id=AxisId(a_id),
2098 offset=int(offset_from_scale + 2 * shape.offset[i]),
2099 )
2100 else:
2101 size = shape[i]
2103 if axis_type == "time":
2104 if tensor_type == "input":
2105 ret.append(TimeInputAxis(size=size, scale=scale))
2106 else:
2107 assert not isinstance(size, ParameterizedSize)
2108 if halo is None:
2109 ret.append(TimeOutputAxis(size=size, scale=scale))
2110 else:
2111 assert not isinstance(size, int)
2112 ret.append(
2113 TimeOutputAxisWithHalo(size=size, scale=scale, halo=halo[i])
2114 )
2116 elif axis_type == "index":
2117 if tensor_type == "input":
2118 ret.append(IndexInputAxis(size=size))
2119 else:
2120 if isinstance(size, ParameterizedSize):
2121 size = DataDependentSize(min=size.min)
2123 ret.append(IndexOutputAxis(size=size))
2124 elif axis_type == "channel":
2125 assert not isinstance(size, ParameterizedSize)
2126 if isinstance(size, SizeReference):
2127 warnings.warn(
2128 "Conversion of channel size from an implicit output shape may be"
2129 + " wrong"
2130 )
2131 ret.append(
2132 ChannelAxis(
2133 channel_names=[f"channel{i}" for i in range(size.offset)]
2134 )
2135 )
2136 else:
2137 ret.append(
2138 ChannelAxis(channel_names=[f"channel{i}" for i in range(size)])
2139 )
2140 elif axis_type == "space":
2141 if tensor_type == "input":
2142 ret.append(SpaceInputAxis(id=AxisId(a), size=size, scale=scale))
2143 else:
2144 assert not isinstance(size, ParameterizedSize)
2145 if halo is None or halo[i] == 0:
2146 ret.append(SpaceOutputAxis(id=AxisId(a), size=size, scale=scale))
2147 elif isinstance(size, int):
2148 raise NotImplementedError(
2149 f"output axis with halo and fixed size (here {size}) not allowed"
2150 )
2151 else:
2152 ret.append(
2153 SpaceOutputAxisWithHalo(
2154 id=AxisId(a), size=size, scale=scale, halo=halo[i]
2155 )
2156 )
2158 return ret
2161def _axes_letters_to_ids(
2162 axes: Optional[str],
2163) -> Optional[List[AxisId]]:
2164 if axes is None:
2165 return None
2167 return [AxisId(a) for a in axes]
2170def _get_complement_v04_axis(
2171 tensor_axes: Sequence[str], axes: Optional[Sequence[str]]
2172) -> Optional[AxisId]:
2173 if axes is None:
2174 return None
2176 non_complement_axes = set(axes) | {"b"}
2177 complement_axes = [a for a in tensor_axes if a not in non_complement_axes]
2178 if len(complement_axes) > 1:
2179 raise ValueError(
2180 f"Expected none or a single complement axis, but axes '{axes}' "
2181 + f"for tensor dims '{tensor_axes}' leave '{complement_axes}'."
2182 )
2184 return None if not complement_axes else AxisId(complement_axes[0])
2187def _convert_proc(
2188 p: Union[_PreprocessingDescr_v0_4, _PostprocessingDescr_v0_4],
2189 tensor_axes: Sequence[str],
2190) -> Union[PreprocessingDescr, PostprocessingDescr]:
2191 if isinstance(p, _BinarizeDescr_v0_4):
2192 return BinarizeDescr(kwargs=BinarizeKwargs(threshold=p.kwargs.threshold))
2193 elif isinstance(p, _ClipDescr_v0_4):
2194 return ClipDescr(kwargs=ClipKwargs(min=p.kwargs.min, max=p.kwargs.max))
2195 elif isinstance(p, _SigmoidDescr_v0_4):
2196 return SigmoidDescr()
2197 elif isinstance(p, _ScaleLinearDescr_v0_4):
2198 axes = _axes_letters_to_ids(p.kwargs.axes)
2199 if p.kwargs.axes is None:
2200 axis = None
2201 else:
2202 axis = _get_complement_v04_axis(tensor_axes, p.kwargs.axes)
2204 if axis is None:
2205 assert not isinstance(p.kwargs.gain, list)
2206 assert not isinstance(p.kwargs.offset, list)
2207 kwargs = ScaleLinearKwargs(gain=p.kwargs.gain, offset=p.kwargs.offset)
2208 else:
2209 kwargs = ScaleLinearAlongAxisKwargs(
2210 axis=axis, gain=p.kwargs.gain, offset=p.kwargs.offset
2211 )
2212 return ScaleLinearDescr(kwargs=kwargs)
2213 elif isinstance(p, _ScaleMeanVarianceDescr_v0_4):
2214 return ScaleMeanVarianceDescr(
2215 kwargs=ScaleMeanVarianceKwargs(
2216 axes=_axes_letters_to_ids(p.kwargs.axes),
2217 reference_tensor=TensorId(str(p.kwargs.reference_tensor)),
2218 eps=p.kwargs.eps,
2219 )
2220 )
2221 elif isinstance(p, _ZeroMeanUnitVarianceDescr_v0_4):
2222 if p.kwargs.mode == "fixed":
2223 mean = p.kwargs.mean
2224 std = p.kwargs.std
2225 assert mean is not None
2226 assert std is not None
2228 axis = _get_complement_v04_axis(tensor_axes, p.kwargs.axes)
2230 if axis is None:
2231 if isinstance(mean, list):
2232 raise ValueError("Expected single float value for mean, not <list>")
2233 if isinstance(std, list):
2234 raise ValueError("Expected single float value for std, not <list>")
2235 return FixedZeroMeanUnitVarianceDescr(
2236 kwargs=FixedZeroMeanUnitVarianceKwargs.model_construct(
2237 mean=mean,
2238 std=std,
2239 )
2240 )
2241 else:
2242 if not isinstance(mean, list):
2243 mean = [float(mean)]
2244 if not isinstance(std, list):
2245 std = [float(std)]
2247 return FixedZeroMeanUnitVarianceDescr(
2248 kwargs=FixedZeroMeanUnitVarianceAlongAxisKwargs(
2249 axis=axis, mean=mean, std=std
2250 )
2251 )
2253 else:
2254 axes = _axes_letters_to_ids(p.kwargs.axes) or []
2255 if p.kwargs.mode == "per_dataset":
2256 axes = [AxisId("batch")] + axes
2257 if not axes:
2258 axes = None
2259 return ZeroMeanUnitVarianceDescr(
2260 kwargs=ZeroMeanUnitVarianceKwargs(axes=axes, eps=p.kwargs.eps)
2261 )
2263 elif isinstance(p, _ScaleRangeDescr_v0_4):
2264 return ScaleRangeDescr(
2265 kwargs=ScaleRangeKwargs(
2266 axes=_axes_letters_to_ids(p.kwargs.axes),
2267 min_percentile=p.kwargs.min_percentile,
2268 max_percentile=p.kwargs.max_percentile,
2269 eps=p.kwargs.eps,
2270 )
2271 )
2272 else:
2273 assert_never(p)
2276class _InputTensorConv(
2277 Converter[
2278 _InputTensorDescr_v0_4,
2279 InputTensorDescr,
2280 FileSource_,
2281 Optional[FileSource_],
2282 Mapping[_TensorName_v0_4, Mapping[str, int]],
2283 ]
2284):
2285 def _convert(
2286 self,
2287 src: _InputTensorDescr_v0_4,
2288 tgt: "type[InputTensorDescr] | type[dict[str, Any]]",
2289 test_tensor: FileSource_,
2290 sample_tensor: Optional[FileSource_],
2291 size_refs: Mapping[_TensorName_v0_4, Mapping[str, int]],
2292 ) -> "InputTensorDescr | dict[str, Any]":
2293 axes: List[InputAxis] = convert_axes( # pyright: ignore[reportAssignmentType]
2294 src.axes,
2295 shape=src.shape,
2296 tensor_type="input",
2297 halo=None,
2298 size_refs=size_refs,
2299 )
2300 prep: List[PreprocessingDescr] = []
2301 for p in src.preprocessing:
2302 cp = _convert_proc(p, src.axes)
2303 assert not isinstance(
2304 cp,
2305 (
2306 CellposeFlowDynamicsDescr,
2307 CustomProcessingDescr,
2308 ScaleMeanVarianceDescr,
2309 StardistPostprocessingDescr,
2310 ),
2311 )
2312 prep.append(cp)
2314 prep.append(EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype="float32")))
2316 return tgt(
2317 axes=axes,
2318 id=TensorId(str(src.name)),
2319 test_tensor=FileDescr(source=test_tensor),
2320 sample_tensor=(
2321 None if sample_tensor is None else FileDescr(source=sample_tensor)
2322 ),
2323 data=dict(type=src.data_type), # pyright: ignore[reportArgumentType]
2324 preprocessing=prep,
2325 )
2328_input_tensor_conv = _InputTensorConv(_InputTensorDescr_v0_4, InputTensorDescr)
2331class OutputTensorDescr(TensorDescrBase[OutputAxis]):
2332 id: TensorId = TensorId("output")
2333 """Output tensor id.
2334 No duplicates are allowed across all inputs and outputs."""
2336 postprocessing: List[PostprocessingDescr] = Field(
2337 default_factory=cast(Callable[[], List[PostprocessingDescr]], list)
2338 )
2339 """Description of how this output should be postprocessed.
2341 note: `postprocessing` always ends with an 'ensure_dtype' operation.
2342 If not given this is added to cast to this tensor's `data.type`.
2343 """
2345 @model_validator(mode="after")
2346 def _validate_postprocessing_kwargs(self) -> Self:
2347 axes_ids = [a.id for a in self.axes]
2348 for p in self.postprocessing:
2349 kwargs_axes = p.kwargs.get("axes")
2350 if kwargs_axes is None:
2351 continue
2353 if not isinstance(kwargs_axes, collections.abc.Sequence):
2354 raise ValueError(
2355 f"expected `axes` sequence, but got {type(kwargs_axes)}"
2356 )
2358 kwargs_axes_seq: Sequence[Any] = cast(Sequence[Any], kwargs_axes)
2359 if any(a not in axes_ids for a in kwargs_axes_seq):
2360 raise ValueError("`kwargs.axes` needs to be subset of axes ids")
2362 if isinstance(self.data, (NominalOrOrdinalDataDescr, IntervalOrRatioDataDescr)):
2363 dtype = self.data.type
2364 else:
2365 dtype = self.data[0].type
2367 # ensure `postprocessing` ends with `EnsureDtypeDescr` or `BinarizeDescr`
2368 if not self.postprocessing or not isinstance(
2369 self.postprocessing[-1], (EnsureDtypeDescr, BinarizeDescr)
2370 ):
2371 self.postprocessing.append(
2372 EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype=dtype))
2373 )
2374 return self
2377class _OutputTensorConv(
2378 Converter[
2379 _OutputTensorDescr_v0_4,
2380 OutputTensorDescr,
2381 FileSource_,
2382 Optional[FileSource_],
2383 Mapping[_TensorName_v0_4, Mapping[str, int]],
2384 ]
2385):
2386 def _convert(
2387 self,
2388 src: _OutputTensorDescr_v0_4,
2389 tgt: "type[OutputTensorDescr] | type[dict[str, Any]]",
2390 test_tensor: FileSource_,
2391 sample_tensor: Optional[FileSource_],
2392 size_refs: Mapping[_TensorName_v0_4, Mapping[str, int]],
2393 ) -> "OutputTensorDescr | dict[str, Any]":
2394 # TODO: split convert_axes into convert_output_axes and convert_input_axes
2395 axes: List[OutputAxis] = convert_axes( # pyright: ignore[reportAssignmentType]
2396 src.axes,
2397 shape=src.shape,
2398 tensor_type="output",
2399 halo=src.halo,
2400 size_refs=size_refs,
2401 )
2402 data_descr: Dict[str, Any] = dict(type=src.data_type)
2403 if data_descr["type"] == "bool":
2404 data_descr["values"] = [False, True]
2406 return tgt(
2407 axes=axes,
2408 id=TensorId(str(src.name)),
2409 test_tensor=FileDescr(source=test_tensor),
2410 sample_tensor=(
2411 None if sample_tensor is None else FileDescr(source=sample_tensor)
2412 ),
2413 data=data_descr, # pyright: ignore[reportArgumentType]
2414 postprocessing=[_convert_proc(p, src.axes) for p in src.postprocessing],
2415 )
2418_output_tensor_conv = _OutputTensorConv(_OutputTensorDescr_v0_4, OutputTensorDescr)
2421TensorDescr = Union[InputTensorDescr, OutputTensorDescr]
2424def get_halos(
2425 tensors: Mapping[TensorId, TensorDescr],
2426 /,
2427) -> Dict[TensorId, Dict[AxisId, Tuple[int, int]]]:
2428 """Get all input and output halos from tensor descriptions.
2430 Note:
2431 - Input halos are to be padded
2432 - Output halos are to be cropped
2433 """
2434 halos: Dict[TensorId, Dict[AxisId, Tuple[int, int]]] = {}
2435 for descr in tensors.values():
2436 if isinstance(descr, InputTensorDescr):
2437 continue
2438 for axis in descr.axes:
2439 if not isinstance(axis, WithHalo):
2440 continue
2442 ref_scale = next(
2443 a
2444 for a in tensors[axis.size.tensor_id].axes
2445 if a.id == axis.size.axis_id
2446 ).scale
2448 # set output halo (to be cropped)
2449 halos.setdefault(descr.id, {})[axis.id] = (axis.halo, axis.halo)
2450 # set input halo (to be padded)
2451 pad_width = int(axis.halo / axis.scale * ref_scale)
2452 halos.setdefault(axis.size.tensor_id, {})[axis.size.axis_id] = (
2453 pad_width,
2454 pad_width,
2455 )
2457 return halos
2460def validate_tensors(
2461 tensors: Mapping[TensorId, Tuple[TensorDescr, Optional[NDArray[Any]]]],
2462 tensor_origin: Literal[
2463 "source", "test_tensor"
2464 ] = "source", # for more precise error messages
2465 *,
2466 pad_inputs: Union[bool, Literal["allow"]] = True,
2467 crop_outputs: Union[bool, Literal["allow"]] = True,
2468):
2469 """Validate all inputs (and optionally output tensors) against their tensor descriptions.
2471 Args:
2472 tensors: Mapping of tensor id to a tuple of tensor description and optional numpy array.
2473 tensor_origin: String to use in error messages to indicate the origin of the tensors being validated.
2474 pad_inputs: Wether to apply/allow padding of inputs before shape comparison
2475 crop_outputs: Wether to apply/allow cropping of outputs before shape comparison.
2476 """
2477 all_tensor_axes: Dict[TensorId, Dict[AxisId, Tuple[AnyAxis, Optional[int]]]] = {}
2479 def e_msg_location(d: TensorDescr):
2480 return f"{'inputs' if isinstance(d, InputTensorDescr) else 'outputs'}[{d.id}]"
2482 for descr, array in tensors.values():
2483 if array is None:
2484 axis_sizes = {a.id: None for a in descr.axes}
2485 else:
2486 try:
2487 axis_sizes = descr.get_axis_sizes_for_array(array)
2488 except ValueError as e:
2489 raise ValueError(f"{e_msg_location(descr)} {e}")
2491 all_tensor_axes[descr.id] = {a.id: (a, axis_sizes[a.id]) for a in descr.axes}
2493 # get halos to be padded/cropped to validate against halo-adjusted sizes
2494 io_halos = get_halos({k: v[0] for k, v in tensors.items()})
2496 for descr, array in tensors.values():
2497 if array is None:
2498 continue
2500 if descr.dtype in ("float32", "float64"):
2501 invalid_test_tensor_dtype = array.dtype.name not in (
2502 "float32",
2503 "float64",
2504 "uint8",
2505 "int8",
2506 "uint16",
2507 "int16",
2508 "uint32",
2509 "int32",
2510 "uint64",
2511 "int64",
2512 )
2513 else:
2514 invalid_test_tensor_dtype = array.dtype.name != descr.dtype
2516 if invalid_test_tensor_dtype:
2517 raise ValueError(
2518 f"{tensor_origin} data type '{array.dtype.name}' does not"
2519 + f" match described {e_msg_location(descr)}.dtype '{descr.dtype}'"
2520 )
2522 if array.min() > -1e-4 and array.max() < 1e-4:
2523 raise ValueError(
2524 "Output values are too small for reliable testing."
2525 + f" Values <-1e5 or >=1e5 must be present in {tensor_origin}"
2526 )
2528 for a in descr.axes:
2529 actual_size = all_tensor_axes[descr.id][a.id][1]
2531 if actual_size is None:
2532 continue
2534 if a.size is None:
2535 continue
2537 # add padding width to actual tensor size
2538 total_axis_halo = sum(io_halos.get(descr.id, {}).get(a.id, (0, 0)))
2539 if isinstance(descr, InputTensorDescr):
2540 # pad input halos
2541 actual_size_with_halo = actual_size + total_axis_halo
2542 if pad_inputs is True:
2543 check_sizes = {actual_size_with_halo}
2544 size_hint = " (after padding input halo)"
2545 elif pad_inputs == "allow":
2546 check_sizes = {actual_size, actual_size_with_halo}
2547 size_hint = " (with or without padding input halo)"
2548 elif pad_inputs is False:
2549 check_sizes = {actual_size}
2550 size_hint = ""
2551 else:
2552 assert_never(pad_inputs)
2554 elif isinstance(descr, OutputTensorDescr):
2555 # crop output halos
2556 actual_size_with_halo = max(0, actual_size - total_axis_halo)
2557 if crop_outputs is True:
2558 check_sizes = {actual_size_with_halo}
2559 size_hint = " (after cropping output halo)"
2560 elif crop_outputs == "allow":
2561 check_sizes = {actual_size, actual_size_with_halo}
2562 size_hint = " (with or without cropping output halo)"
2563 elif crop_outputs is False:
2564 check_sizes = {actual_size}
2565 size_hint = ""
2566 else:
2567 assert_never(crop_outputs)
2568 else:
2569 assert_never(descr)
2571 del actual_size # make sure we explicitly use unchanged or halo-adjusted size from here on
2573 if isinstance(a.size, int):
2574 if a.size not in check_sizes:
2575 raise ValueError(
2576 f"{e_msg_location(descr)}.axes[{a.id}]: {tensor_origin} axis "
2577 + f"has incompatible size {check_sizes}{size_hint}, expected {a.size}"
2578 )
2579 elif isinstance(a.size, (ParameterizedSize, DataDependentSize)):
2580 _ = try_all_raise_last(
2581 (partial(a.size.validate_size, s) for s in check_sizes),
2582 f"{e_msg_location(descr)}.axes[{a.id}]: {tensor_origin} axis ",
2583 )
2584 elif isinstance(a.size, SizeReference):
2585 ref_tensor_axes = all_tensor_axes.get(a.size.tensor_id)
2586 if ref_tensor_axes is None:
2587 raise ValueError(
2588 f"{e_msg_location(descr)}.axes[{a.id}].size.tensor_id: Unknown tensor"
2589 + f" reference '{a.size.tensor_id}', available: {list(all_tensor_axes)}"
2590 )
2592 ref_axis, ref_size = ref_tensor_axes.get(a.size.axis_id, (None, None))
2593 if ref_axis is None or ref_size is None:
2594 raise ValueError(
2595 f"{e_msg_location(descr)}.axes[{a.id}].size.axis_id: Unknown tensor axis"
2596 + f" reference '{a.size.tensor_id}.{a.size.axis_id}, available: {list(ref_tensor_axes)}"
2597 )
2599 if a.unit != ref_axis.unit:
2600 raise ValueError(
2601 f"{e_msg_location(descr)}.axes[{a.id}].size: `SizeReference` requires"
2602 + " axis and reference axis to have the same `unit`, but"
2603 + f" {a.unit}!={ref_axis.unit}"
2604 )
2606 if (
2607 expected_size := (
2608 ref_size * ref_axis.scale / a.scale + a.size.offset
2609 )
2610 ) not in check_sizes:
2611 raise ValueError(
2612 f"{e_msg_location(descr)}.{tensor_origin}: axis '{a.id}' of size"
2613 + f" {check_sizes} invalid for referenced size {ref_size};"
2614 + f" expected {expected_size}"
2615 )
2616 else:
2617 assert_never(a.size)
2620FileDescr_dependencies = Annotated[
2621 FileDescr_,
2622 WithSuffix((".yaml", ".yml"), case_sensitive=True),
2623 Field(examples=[dict(source="environment.yaml")]),
2624]
2627class _ArchitectureCallableDescr(Node):
2628 callable: Annotated[Identifier, Field(examples=["MyNetworkClass", "get_my_model"])]
2629 """Identifier of the callable that returns a torch.nn.Module instance."""
2631 kwargs: Dict[str, YamlValue] = Field(
2632 default_factory=cast(Callable[[], Dict[str, YamlValue]], dict)
2633 )
2634 """key word arguments for the `callable`"""
2637class ArchitectureFromFileDescr(_ArchitectureCallableDescr, FileDescr):
2638 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
2639 """Architecture source file"""
2641 @model_serializer(mode="wrap", when_used="unless-none")
2642 def _serialize(self, nxt: SerializerFunctionWrapHandler, info: SerializationInfo):
2643 return package_file_descr_serializer(self, nxt, info)
2646class ArchitectureFromLibraryDescr(_ArchitectureCallableDescr):
2647 import_from: str
2648 """Where to import the callable from, i.e. `from <import_from> import <callable>`"""
2651class _ArchFileConv(
2652 Converter[
2653 _CallableFromFile_v0_4,
2654 ArchitectureFromFileDescr,
2655 Optional[Sha256],
2656 Dict[str, Any],
2657 ]
2658):
2659 def _convert(
2660 self,
2661 src: _CallableFromFile_v0_4,
2662 tgt: "type[ArchitectureFromFileDescr | dict[str, Any]]",
2663 sha256: Optional[Sha256],
2664 kwargs: Dict[str, Any],
2665 ) -> "ArchitectureFromFileDescr | dict[str, Any]":
2666 if src.startswith("http") and src.count(":") == 2:
2667 http, source, callable_ = src.split(":")
2668 source = ":".join((http, source))
2669 elif not src.startswith("http") and src.count(":") == 1:
2670 source, callable_ = src.split(":")
2671 else:
2672 source = str(src)
2673 callable_ = str(src)
2674 return tgt(
2675 callable=Identifier(callable_),
2676 source=cast(FileSource_, source),
2677 sha256=sha256,
2678 kwargs=kwargs,
2679 )
2682_arch_file_conv = _ArchFileConv(_CallableFromFile_v0_4, ArchitectureFromFileDescr)
2685class _ArchLibConv(
2686 Converter[
2687 _CallableFromDepencency_v0_4, ArchitectureFromLibraryDescr, Dict[str, Any]
2688 ]
2689):
2690 def _convert(
2691 self,
2692 src: _CallableFromDepencency_v0_4,
2693 tgt: "type[ArchitectureFromLibraryDescr | dict[str, Any]]",
2694 kwargs: Dict[str, Any],
2695 ) -> "ArchitectureFromLibraryDescr | dict[str, Any]":
2696 *mods, callable_ = src.split(".")
2697 import_from = ".".join(mods)
2698 return tgt(
2699 import_from=import_from, callable=Identifier(callable_), kwargs=kwargs
2700 )
2703_arch_lib_conv = _ArchLibConv(
2704 _CallableFromDepencency_v0_4, ArchitectureFromLibraryDescr
2705)
2708class WeightsEntryDescrBase(FileDescr):
2709 type: ClassVar[WeightsFormat]
2710 weights_format_name: ClassVar[str] # human readable
2712 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
2713 """Source of the weights file."""
2715 authors: Optional[List[Author]] = None
2716 """Authors
2717 Either the person(s) that have trained this model resulting in the original weights file.
2718 (If this is the initial weights entry, i.e. it does not have a `parent`)
2719 Or the person(s) who have converted the weights to this weights format.
2720 (If this is a child weight, i.e. it has a `parent` field)
2721 """
2723 parent: Annotated[
2724 Optional[WeightsFormat], Field(examples=["pytorch_state_dict"])
2725 ] = None
2726 """The source weights these weights were converted from.
2727 For example, if a model's weights were converted from the `pytorch_state_dict` format to `torchscript`,
2728 The `pytorch_state_dict` weights entry has no `parent` and is the parent of the `torchscript` weights.
2729 All weight entries except one (the initial set of weights resulting from training the model),
2730 need to have this field."""
2732 comment: str = ""
2733 """A comment about this weights entry, for example how these weights were created."""
2735 @model_validator(mode="after")
2736 def _validate(self) -> Self:
2737 if self.type == self.parent:
2738 raise ValueError("Weights entry can't be it's own parent.")
2740 return self
2742 @model_serializer(mode="wrap", when_used="unless-none")
2743 def _serialize(self, nxt: SerializerFunctionWrapHandler, info: SerializationInfo):
2744 return package_file_descr_serializer(self, nxt, info)
2747class KerasHdf5WeightsDescr(WeightsEntryDescrBase):
2748 type: ClassVar[WeightsFormat] = "keras_hdf5"
2749 weights_format_name: ClassVar[str] = "Keras HDF5"
2750 tensorflow_version: Version
2751 """TensorFlow version used to create these weights."""
2754class KerasV3WeightsDescr(WeightsEntryDescrBase):
2755 type: ClassVar[WeightsFormat] = "keras_v3"
2756 weights_format_name: ClassVar[str] = "Keras v3"
2757 keras_version: Annotated[Version, Ge(Version(3))]
2758 """Keras version used to create these weights."""
2759 backend: Tuple[Literal["tensorflow", "jax", "torch"], Version]
2760 """Keras backend used to create these weights."""
2761 source: Annotated[
2762 FileSource,
2763 AfterValidator(wo_special_file_name),
2764 WithSuffix(".keras", case_sensitive=True),
2765 ]
2766 """Source of the .keras weights file."""
2769FileDescr_external_data = Annotated[
2770 FileDescr_,
2771 WithSuffix(".data", case_sensitive=True),
2772 Field(examples=[dict(source="weights.onnx.data")]),
2773]
2776class OnnxWeightsDescr(WeightsEntryDescrBase):
2777 type: ClassVar[WeightsFormat] = "onnx"
2778 weights_format_name: ClassVar[str] = "ONNX"
2779 opset_version: Annotated[int, Ge(7)]
2780 """ONNX opset version"""
2782 external_data: Optional[FileDescr_external_data] = None
2783 """Source of the external ONNX data file holding the weights.
2784 (If present **source** holds the ONNX architecture without weights)."""
2786 @model_validator(mode="after")
2787 def _validate_external_data_unique_file_name(self) -> Self:
2788 if self.external_data is not None and (
2789 extract_file_name(self.source)
2790 == extract_file_name(self.external_data.source)
2791 ):
2792 raise ValueError(
2793 f"ONNX `external_data` file name '{extract_file_name(self.external_data.source)}'"
2794 + " must be different from ONNX `source` file name."
2795 )
2797 return self
2800class PytorchStateDictWeightsDescr(WeightsEntryDescrBase):
2801 type: ClassVar[WeightsFormat] = "pytorch_state_dict"
2802 weights_format_name: ClassVar[str] = "Pytorch State Dict"
2803 architecture: Union[ArchitectureFromFileDescr, ArchitectureFromLibraryDescr]
2804 pytorch_version: Version
2805 """Version of the PyTorch library used.
2806 If `architecture.depencencies` is specified it has to include pytorch and any version pinning has to be compatible.
2807 """
2808 dependencies: Optional[FileDescr_dependencies] = None
2809 """Custom depencies beyond pytorch described in a Conda environment file.
2810 Allows to specify custom dependencies, see conda docs:
2811 - [Exporting an environment file across platforms](https://conda.io/projects/conda/en/latest/user-guide/tasks/manage-environments.html#exporting-an-environment-file-across-platforms)
2812 - [Creating an environment file manually](https://conda.io/projects/conda/en/latest/user-guide/tasks/manage-environments.html#creating-an-environment-file-manually)
2814 The conda environment file should include pytorch and any version pinning has to be compatible with
2815 **pytorch_version**.
2816 """
2817 strict: bool = True
2818 """Whether to allow missing or unexpected keys or to be strict about the architecture matching the state dict weights."""
2821class TensorflowJsWeightsDescr(WeightsEntryDescrBase):
2822 type: ClassVar[WeightsFormat] = "tensorflow_js"
2823 weights_format_name: ClassVar[str] = "Tensorflow.js"
2824 tensorflow_version: Version
2825 """Version of the TensorFlow library used."""
2827 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
2828 """The multi-file weights.
2829 All required files/folders should be a zip archive."""
2832class TensorflowSavedModelBundleWeightsDescr(WeightsEntryDescrBase):
2833 type: ClassVar[WeightsFormat] = "tensorflow_saved_model_bundle"
2834 weights_format_name: ClassVar[str] = "Tensorflow Saved Model"
2835 tensorflow_version: Version
2836 """Version of the TensorFlow library used."""
2838 dependencies: Optional[FileDescr_dependencies] = None
2839 """Custom dependencies beyond tensorflow.
2840 Should include tensorflow and any version pinning has to be compatible with **tensorflow_version**."""
2842 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
2843 """The multi-file weights.
2844 All required files/folders should be a zip archive."""
2847class TorchscriptWeightsDescr(WeightsEntryDescrBase):
2848 type: ClassVar[WeightsFormat] = "torchscript"
2849 weights_format_name: ClassVar[str] = "TorchScript"
2850 pytorch_version: Version
2851 """Version of the PyTorch library used."""
2854SpecificWeightsDescr = Union[
2855 KerasHdf5WeightsDescr,
2856 KerasV3WeightsDescr,
2857 OnnxWeightsDescr,
2858 PytorchStateDictWeightsDescr,
2859 TensorflowJsWeightsDescr,
2860 TensorflowSavedModelBundleWeightsDescr,
2861 TorchscriptWeightsDescr,
2862]
2865class WeightsDescr(Node):
2866 keras_hdf5: Optional[KerasHdf5WeightsDescr] = None
2867 keras_v3: Optional[KerasV3WeightsDescr] = None
2868 onnx: Optional[OnnxWeightsDescr] = None
2869 pytorch_state_dict: Optional[PytorchStateDictWeightsDescr] = None
2870 tensorflow_js: Optional[TensorflowJsWeightsDescr] = None
2871 tensorflow_saved_model_bundle: Optional[TensorflowSavedModelBundleWeightsDescr] = (
2872 None
2873 )
2874 torchscript: Optional[TorchscriptWeightsDescr] = None
2876 @model_validator(mode="after")
2877 def check_entries(self) -> Self:
2878 entries = {wtype for wtype, entry in self if entry is not None}
2880 if not entries:
2881 raise ValueError("Missing weights entry")
2883 entries_wo_parent = {
2884 wtype
2885 for wtype, entry in self
2886 if entry is not None and hasattr(entry, "parent") and entry.parent is None
2887 }
2888 if len(entries_wo_parent) != 1:
2889 issue_warning(
2890 "Exactly one weights entry may not specify the `parent` field (got"
2891 + " {value}). That entry is considered the original set of model weights."
2892 + " Other weight formats are created through conversion of the orignal or"
2893 + " already converted weights. They have to reference the weights format"
2894 + " they were converted from as their `parent`.",
2895 value=len(entries_wo_parent),
2896 field="weights",
2897 )
2899 for wtype, entry in self:
2900 if entry is None:
2901 continue
2903 assert hasattr(entry, "type")
2904 assert hasattr(entry, "parent")
2905 assert wtype == entry.type
2906 if (
2907 entry.parent is not None and entry.parent not in entries
2908 ): # self reference checked for `parent` field
2909 raise ValueError(
2910 f"`weights.{wtype}.parent={entry.parent} not in specified weight"
2911 + f" formats: {entries}"
2912 )
2914 return self
2916 def __getitem__(
2917 self,
2918 key: WeightsFormat,
2919 ):
2920 if key == "keras_hdf5":
2921 ret = self.keras_hdf5
2922 elif key == "keras_v3":
2923 ret = self.keras_v3
2924 elif key == "onnx":
2925 ret = self.onnx
2926 elif key == "pytorch_state_dict":
2927 ret = self.pytorch_state_dict
2928 elif key == "tensorflow_js":
2929 ret = self.tensorflow_js
2930 elif key == "tensorflow_saved_model_bundle":
2931 ret = self.tensorflow_saved_model_bundle
2932 elif key == "torchscript":
2933 ret = self.torchscript
2934 else:
2935 raise KeyError(key)
2937 if ret is None:
2938 raise KeyError(key)
2940 return ret
2942 @overload
2943 def __setitem__(
2944 self, key: Literal["keras_hdf5"], value: Optional[KerasHdf5WeightsDescr]
2945 ) -> None: ...
2946 @overload
2947 def __setitem__(
2948 self, key: Literal["keras_v3"], value: Optional[KerasV3WeightsDescr]
2949 ) -> None: ...
2950 @overload
2951 def __setitem__(
2952 self, key: Literal["onnx"], value: Optional[OnnxWeightsDescr]
2953 ) -> None: ...
2954 @overload
2955 def __setitem__(
2956 self,
2957 key: Literal["pytorch_state_dict"],
2958 value: Optional[PytorchStateDictWeightsDescr],
2959 ) -> None: ...
2960 @overload
2961 def __setitem__(
2962 self, key: Literal["tensorflow_js"], value: Optional[TensorflowJsWeightsDescr]
2963 ) -> None: ...
2964 @overload
2965 def __setitem__(
2966 self,
2967 key: Literal["tensorflow_saved_model_bundle"],
2968 value: Optional[TensorflowSavedModelBundleWeightsDescr],
2969 ) -> None: ...
2970 @overload
2971 def __setitem__(
2972 self, key: Literal["torchscript"], value: Optional[TorchscriptWeightsDescr]
2973 ) -> None: ...
2975 def __setitem__(
2976 self,
2977 key: WeightsFormat,
2978 value: Optional[SpecificWeightsDescr],
2979 ):
2980 if key == "keras_hdf5":
2981 if value is not None and not isinstance(value, KerasHdf5WeightsDescr):
2982 raise TypeError(
2983 f"Expected KerasHdf5WeightsDescr or None for key 'keras_hdf5', got {type(value)}"
2984 )
2985 self.keras_hdf5 = value
2986 elif key == "keras_v3":
2987 if value is not None and not isinstance(value, KerasV3WeightsDescr):
2988 raise TypeError(
2989 f"Expected KerasV3WeightsDescr or None for key 'keras_v3', got {type(value)}"
2990 )
2991 self.keras_v3 = value
2992 elif key == "onnx":
2993 if value is not None and not isinstance(value, OnnxWeightsDescr):
2994 raise TypeError(
2995 f"Expected OnnxWeightsDescr or None for key 'onnx', got {type(value)}"
2996 )
2997 self.onnx = value
2998 elif key == "pytorch_state_dict":
2999 if value is not None and not isinstance(
3000 value, PytorchStateDictWeightsDescr
3001 ):
3002 raise TypeError(
3003 f"Expected PytorchStateDictWeightsDescr or None for key 'pytorch_state_dict', got {type(value)}"
3004 )
3005 self.pytorch_state_dict = value
3006 elif key == "tensorflow_js":
3007 if value is not None and not isinstance(value, TensorflowJsWeightsDescr):
3008 raise TypeError(
3009 f"Expected TensorflowJsWeightsDescr or None for key 'tensorflow_js', got {type(value)}"
3010 )
3011 self.tensorflow_js = value
3012 elif key == "tensorflow_saved_model_bundle":
3013 if value is not None and not isinstance(
3014 value, TensorflowSavedModelBundleWeightsDescr
3015 ):
3016 raise TypeError(
3017 f"Expected TensorflowSavedModelBundleWeightsDescr or None for key 'tensorflow_saved_model_bundle', got {type(value)}"
3018 )
3019 self.tensorflow_saved_model_bundle = value
3020 elif key == "torchscript":
3021 if value is not None and not isinstance(value, TorchscriptWeightsDescr):
3022 raise TypeError(
3023 f"Expected TorchscriptWeightsDescr or None for key 'torchscript', got {type(value)}"
3024 )
3025 self.torchscript = value
3026 else:
3027 raise KeyError(key)
3029 @property
3030 def available_formats(self) -> Dict[WeightsFormat, SpecificWeightsDescr]:
3031 return {
3032 **({} if self.keras_hdf5 is None else {"keras_hdf5": self.keras_hdf5}),
3033 **({} if self.keras_v3 is None else {"keras_v3": self.keras_v3}),
3034 **({} if self.onnx is None else {"onnx": self.onnx}),
3035 **(
3036 {}
3037 if self.pytorch_state_dict is None
3038 else {"pytorch_state_dict": self.pytorch_state_dict}
3039 ),
3040 **(
3041 {}
3042 if self.tensorflow_js is None
3043 else {"tensorflow_js": self.tensorflow_js}
3044 ),
3045 **(
3046 {}
3047 if self.tensorflow_saved_model_bundle is None
3048 else {
3049 "tensorflow_saved_model_bundle": self.tensorflow_saved_model_bundle
3050 }
3051 ),
3052 **({} if self.torchscript is None else {"torchscript": self.torchscript}),
3053 }
3055 @property
3056 def missing_formats(self) -> Set[WeightsFormat]:
3057 return {
3058 wf for wf in get_args(WeightsFormat) if wf not in self.available_formats
3059 }
3062class ModelId(ResourceId):
3063 pass
3066class LinkedModel(LinkedResourceBase):
3067 """Reference to a bioimage.io model."""
3069 id: ModelId
3070 """A valid model `id` from the bioimage.io collection."""
3073class _DataDepSize(NamedTuple):
3074 min: StrictInt
3075 max: Optional[StrictInt]
3078class _AxisSizes(NamedTuple):
3079 """the lenghts of all axes of model inputs and outputs"""
3081 inputs: Dict[Tuple[TensorId, AxisId], int]
3082 outputs: Dict[Tuple[TensorId, AxisId], Union[int, _DataDepSize]]
3085class _TensorSizes(NamedTuple):
3086 """_AxisSizes as nested dicts"""
3088 inputs: Dict[TensorId, Dict[AxisId, int]]
3089 outputs: Dict[TensorId, Dict[AxisId, Union[int, _DataDepSize]]]
3092class ReproducibilityTolerance(Node, extra="allow"):
3093 """Describes what small numerical differences -- if any -- may be tolerated
3094 in the generated output when executing in different environments.
3096 A tensor element *output* is considered mismatched to the **test_tensor** if
3097 abs(*output* - **test_tensor**) > **absolute_tolerance** + **relative_tolerance** * abs(**test_tensor**).
3098 (Internally we call [numpy.testing.assert_allclose](https://numpy.org/doc/stable/reference/generated/numpy.testing.assert_allclose.html).)
3100 Motivation:
3101 For testing we can request the respective deep learning frameworks to be as
3102 reproducible as possible by setting seeds and chosing deterministic algorithms,
3103 but differences in operating systems, available hardware and installed drivers
3104 may still lead to numerical differences.
3105 """
3107 relative_tolerance: RelativeTolerance = 1e-3
3108 """Maximum relative tolerance of reproduced test tensor."""
3110 absolute_tolerance: AbsoluteTolerance = 1e-3
3111 """Maximum absolute tolerance of reproduced test tensor."""
3113 mismatched_elements_per_million: MismatchedElementsPerMillion = 100
3114 """Maximum number of mismatched elements/pixels per million to tolerate."""
3116 output_ids: Sequence[TensorId] = ()
3117 """Limits the output tensor IDs these reproducibility details apply to."""
3119 weights_formats: Sequence[WeightsFormat] = ()
3120 """Limits the weights formats these details apply to."""
3123class BiasRisksLimitations(Node, extra="allow"):
3124 """Known biases, risks, technical limitations, and recommendations for model use."""
3126 known_biases: str = dedent("""\
3127 In general bioimage models may suffer from biases caused by:
3129 - Imaging protocol dependencies
3130 - Use of a specific cell type
3131 - Species-specific training data limitations
3133 """)
3134 """Biases in training data or model behavior."""
3136 risks: str = dedent("""\
3137 Common risks in bioimage analysis include:
3139 - Erroneously assuming generalization to unseen experimental conditions
3140 - Trusting (overconfident) model outputs without validation
3141 - Misinterpretation of results
3143 """)
3144 """Potential risks in the context of bioimage analysis."""
3146 limitations: Optional[str] = None
3147 """Technical limitations and failure modes."""
3149 recommendations: str = "Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model."
3150 """Mitigation strategies regarding `known_biases`, `risks`, and `limitations`, as well as applicable best practices.
3152 Consider:
3153 - How to use a validation dataset?
3154 - How to manually validate?
3155 - Feasibility of domain adaptation for different experimental setups?
3157 """
3159 def format_md(self) -> str:
3160 if self.limitations is None:
3161 limitations_header = ""
3162 else:
3163 limitations_header = "## Limitations\n\n"
3165 return f"""# Bias, Risks, and Limitations
3167{self.known_biases}
3169{self.risks}
3171{limitations_header}{self.limitations or ""}
3173## Recommendations
3175{self.recommendations}
3177"""
3180class TrainingDetails(Node, extra="allow"):
3181 training_preprocessing: Optional[str] = None
3182 """Detailed image preprocessing steps during model training:
3184 Mention:
3185 - *Normalization methods*
3186 - *Augmentation strategies*
3187 - *Resizing/resampling procedures*
3188 - *Artifact handling*
3190 """
3192 training_epochs: Optional[float] = None
3193 """Number of training epochs."""
3195 training_batch_size: Optional[float] = None
3196 """Batch size used in training."""
3198 initial_learning_rate: Optional[float] = None
3199 """Initial learning rate used in training."""
3201 learning_rate_schedule: Optional[str] = None
3202 """Learning rate schedule used in training."""
3204 loss_function: Optional[str] = None
3205 """Loss function used in training, e.g. nn.MSELoss."""
3207 loss_function_kwargs: Dict[str, YamlValue] = Field(
3208 default_factory=cast(Callable[[], Dict[str, YamlValue]], dict)
3209 )
3210 """key word arguments for the `loss_function`"""
3212 optimizer: Optional[str] = None
3213 """optimizer, e.g. torch.optim.Adam"""
3215 optimizer_kwargs: Dict[str, YamlValue] = Field(
3216 default_factory=cast(Callable[[], Dict[str, YamlValue]], dict)
3217 )
3218 """key word arguments for the `optimizer`"""
3220 regularization: Optional[str] = None
3221 """Regularization techniques used during training, e.g. drop-out or weight decay."""
3223 training_duration: Optional[float] = None
3224 """Total training duration in hours."""
3227class Evaluation(Node, extra="allow"):
3228 model_id: Optional[ModelId] = None
3229 """Model being evaluated."""
3231 dataset_id: DatasetId
3232 """Dataset used for evaluation."""
3234 dataset_source: HttpUrl
3235 """Source of the dataset."""
3237 dataset_role: Literal["train", "validation", "test", "independent", "unknown"]
3238 """Role of the dataset used for evaluation.
3240 - `train`: dataset was (part of) the training data
3241 - `validation`: dataset was (part of) the validation data used during training, e.g. used for model selection or hyperparameter tuning
3242 - `test`: dataset was (part of) the designated test data; not used during training or validation, but acquired from the same source/distribution as training data
3243 - `independent`: dataset is entirely independent test data; not used during training or validation, and acquired from a different source/distribution than training data
3244 - `unknown`: role of the dataset is unknown; choose this if you are not certain if (a subset) of the data was seen by the model during training.
3245 """
3247 sample_count: int
3248 """Number of evaluated samples."""
3250 evaluation_factors: List[Annotated[str, MaxLen(16)]]
3251 """(Abbreviations of) each evaluation factor.
3253 Evaluation factors are criteria along which model performance is evaluated, e.g. different image conditions
3254 like 'low SNR', 'high cell density', or different biological conditions like 'cell type A', 'cell type B'.
3255 An 'overall' factor may be included to summarize performance across all conditions.
3256 """
3258 evaluation_factors_long: List[str]
3259 """Descriptions (long form) of each evaluation factor."""
3261 metrics: List[Annotated[str, MaxLen(16)]]
3262 """(Abbreviations of) metrics used for evaluation."""
3264 metrics_long: List[str]
3265 """Description of each metric used."""
3267 @model_validator(mode="after")
3268 def _validate_list_lengths(self) -> Self:
3269 if len(self.evaluation_factors) != len(self.evaluation_factors_long):
3270 raise ValueError(
3271 "`evaluation_factors` and `evaluation_factors_long` must have the same length"
3272 )
3274 if len(self.metrics) != len(self.metrics_long):
3275 raise ValueError("`metrics` and `metrics_long` must have the same length")
3277 if len(self.results) != len(self.metrics):
3278 raise ValueError("`results` must have the same number of rows as `metrics`")
3280 for row in self.results:
3281 if len(row) != len(self.evaluation_factors):
3282 raise ValueError(
3283 "`results` must have the same number of columns (in every row) as `evaluation_factors`"
3284 )
3286 return self
3288 results: List[List[Union[str, float, int]]]
3289 """Results for each metric (rows; outer list) and each evaluation factor (columns; inner list)."""
3291 results_summary: Optional[str] = None
3292 """Interpretation of results for general audience.
3294 Consider:
3295 - Overall model performance
3296 - Comparison to existing methods
3297 - Limitations and areas for improvement
3299"""
3301 def format_md(self):
3302 results_header = ["Metric"] + self.evaluation_factors
3303 results_table_cells = [results_header, ["---"] * len(results_header)] + [
3304 [metric] + [str(r) for r in row]
3305 for metric, row in zip(self.metrics, self.results)
3306 ]
3308 results_table = "".join(
3309 "| " + " | ".join(row) + " |\n" for row in results_table_cells
3310 )
3311 factors = "".join(
3312 f"\n - {ef}: {efl}"
3313 for ef, efl in zip(self.evaluation_factors, self.evaluation_factors_long)
3314 )
3315 metrics = "".join(
3316 f"\n - {em}: {eml}" for em, eml in zip(self.metrics, self.metrics_long)
3317 )
3319 return f"""## Testing Data, Factors & Metrics
3321Evaluation of {self.model_id or "this"} model on the {self.dataset_id} dataset (dataset role: {self.dataset_role}).
3323### Testing Data
3325- **Source:** [{self.dataset_id}]({self.dataset_source})
3326- **Size:** {self.sample_count} evaluated samples
3328### Factors
3329{factors}
3331### Metrics
3332{metrics}
3334## Results
3336### Quantitative Results
3338{results_table}
3340### Summary
3342{self.results_summary or "missing"}
3344"""
3347class EnvironmentalImpact(Node, extra="allow"):
3348 """Environmental considerations for model training and deployment.
3350 Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
3351 """
3353 hardware_type: Optional[str] = None
3354 """GPU/CPU specifications"""
3356 hours_used: Optional[float] = None
3357 """Total compute hours"""
3359 cloud_provider: Optional[str] = None
3360 """If applicable"""
3362 compute_region: Optional[str] = None
3363 """Geographic location"""
3365 co2_emitted: Optional[float] = None
3366 """kg CO2 equivalent
3368 Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
3369 """
3371 def format_md(self):
3372 """Filled Markdown template section following [Hugging Face Model Card Template](https://huggingface.co/docs/hub/en/model-card-annotated)."""
3373 if self == self.__class__():
3374 return ""
3376 ret = "# Environmental Impact\n\n"
3377 if self.hardware_type is not None:
3378 ret += f"- **Hardware Type:** {self.hardware_type}\n"
3379 if self.hours_used is not None:
3380 ret += f"- **Hours used:** {self.hours_used}\n"
3381 if self.cloud_provider is not None:
3382 ret += f"- **Cloud Provider:** {self.cloud_provider}\n"
3383 if self.compute_region is not None:
3384 ret += f"- **Compute Region:** {self.compute_region}\n"
3385 if self.co2_emitted is not None:
3386 ret += f"- **Carbon Emitted:** {self.co2_emitted} kg CO2e\n"
3388 return ret + "\n"
3391class BioimageioConfig(Node, extra="allow"):
3392 reproducibility_tolerance: Sequence[ReproducibilityTolerance] = ()
3393 """Tolerances to allow when reproducing the model's test outputs
3394 from the model's test inputs.
3395 Only the first entry matching tensor id and weights format is considered.
3396 """
3398 funded_by: Optional[str] = None
3399 """Funding agency, grant number if applicable"""
3401 architecture_type: Optional[Annotated[str, MaxLen(32)]] = (
3402 None # TODO: add to differentiated tags
3403 )
3404 """Model architecture type, e.g., 3D U-Net, ResNet, transformer"""
3406 architecture_description: Optional[str] = None
3407 """Text description of model architecture."""
3409 modality: Optional[str] = None # TODO: add to differentiated tags
3410 """Input modality, e.g., fluorescence microscopy, electron microscopy"""
3412 target_structure: List[str] = Field( # TODO: add to differentiated tags
3413 default_factory=cast(Callable[[], List[str]], list)
3414 )
3415 """Biological structure(s) the model is designed to analyze, e.g., nuclei, mitochondria, cells"""
3417 task: Optional[str] = None # TODO: add to differentiated tags
3418 """Bioimage-specific task type, e.g., segmentation, classification, detection, denoising"""
3420 new_version: Optional[ModelId] = None
3421 """A new version of this model exists with a different model id."""
3423 out_of_scope_use: Optional[str] = None
3424 """Describe how the model may be misused in bioimage analysis contexts and what users should **not** do with the model."""
3426 bias_risks_limitations: BiasRisksLimitations = Field(
3427 default_factory=BiasRisksLimitations.model_construct
3428 )
3429 """Description of known bias, risks, and technical limitations for in-scope model use."""
3431 model_parameter_count: Optional[int] = None
3432 """Total number of model parameters."""
3434 training: TrainingDetails = Field(default_factory=TrainingDetails.model_construct)
3435 """Details on how the model was trained."""
3437 inference_time: Optional[str] = None
3438 """Average inference time per image/tile. Specify hardware and image size. Multiple examples can be given."""
3440 memory_requirements_inference: Optional[str] = None
3441 """GPU memory needed for inference. Multiple examples with different image size can be given."""
3443 memory_requirements_training: Optional[str] = None
3444 """GPU memory needed for training. Multiple examples with different image/batch sizes can be given."""
3446 evaluations: List[Evaluation] = Field(
3447 default_factory=cast(Callable[[], List[Evaluation]], list)
3448 )
3449 """Quantitative model evaluations.
3451 Note:
3452 At the moment we recommend to include only a single test dataset
3453 (with evaluation factors that may mark subsets of the dataset)
3454 to avoid confusion and make the presentation of results cleaner.
3455 """
3457 environmental_impact: EnvironmentalImpact = Field(
3458 default_factory=EnvironmentalImpact.model_construct
3459 )
3460 """Environmental considerations for model training and deployment"""
3463class Config(Node, extra="allow"):
3464 bioimageio: BioimageioConfig = Field(
3465 default_factory=BioimageioConfig.model_construct
3466 )
3467 stardist: YamlValue = None
3470class ModelDescr(GenericModelDescrBase):
3471 """Specification of the fields used in a bioimage.io-compliant RDF to describe AI models with pretrained weights.
3472 These fields are typically stored in a YAML file which we call a model resource description file (model RDF).
3473 """
3475 implemented_format_version: ClassVar[Literal["0.5.10"]] = "0.5.10"
3476 if TYPE_CHECKING:
3477 format_version: Literal["0.5.10"] = "0.5.10"
3478 else:
3479 format_version: Literal["0.5.10"]
3480 """Version of the bioimage.io model description specification used.
3481 When creating a new model always use the latest micro/patch version described here.
3482 The `format_version` is important for any consumer software to understand how to parse the fields.
3483 """
3485 implemented_type: ClassVar[Literal["model"]] = "model"
3486 if TYPE_CHECKING:
3487 type: Literal["model"] = "model"
3488 else:
3489 type: Literal["model"]
3490 """Specialized resource type 'model'"""
3492 id: Optional[ModelId] = None
3493 """bioimage.io-wide unique resource identifier
3494 assigned by bioimage.io; version **un**specific."""
3496 authors: FAIR[List[Author]] = Field(
3497 default_factory=cast(Callable[[], List[Author]], list)
3498 )
3499 """The authors are the creators of the model RDF and the primary points of contact."""
3501 documentation: FAIR[Optional[FileSource_documentation]] = None
3502 """URL or relative path to a markdown file with additional documentation.
3503 The recommended documentation file name is `README.md`. An `.md` suffix is mandatory.
3504 The documentation should include a '#[#] Validation' (sub)section
3505 with details on how to quantitatively validate the model on unseen data."""
3507 @field_validator("documentation", mode="after")
3508 @classmethod
3509 def _validate_documentation(
3510 cls, value: Optional[FileSource_documentation]
3511 ) -> Optional[FileSource_documentation]:
3512 if not get_validation_context().perform_io_checks or value is None:
3513 return value
3515 doc_reader = get_reader(value)
3516 doc_content = doc_reader.read().decode(encoding="utf-8")
3517 if not re.search("#.*[vV]alidation", doc_content):
3518 issue_warning(
3519 "No '# Validation' (sub)section found in {value}.",
3520 value=value,
3521 field="documentation",
3522 )
3524 return value
3526 inputs: NotEmpty[Sequence[InputTensorDescr]]
3527 """Describes the input tensors expected by this model."""
3529 @field_validator("inputs", mode="after")
3530 @classmethod
3531 def _validate_input_axes(
3532 cls, inputs: Sequence[InputTensorDescr]
3533 ) -> Sequence[InputTensorDescr]:
3534 input_size_refs = cls._get_axes_with_independent_size(inputs)
3536 for i, ipt in enumerate(inputs):
3537 valid_independent_refs: Dict[
3538 Tuple[TensorId, AxisId],
3539 Tuple[TensorDescr, AnyAxis, Union[int, ParameterizedSize]],
3540 ] = {
3541 **{
3542 (ipt.id, a.id): (ipt, a, a.size)
3543 for a in ipt.axes
3544 if not isinstance(a, BatchAxis)
3545 and isinstance(a.size, (int, ParameterizedSize))
3546 },
3547 **input_size_refs,
3548 }
3549 for a, ax in enumerate(ipt.axes):
3550 cls._validate_axis(
3551 "inputs",
3552 i=i,
3553 tensor_id=ipt.id,
3554 a=a,
3555 axis=ax,
3556 valid_independent_refs=valid_independent_refs,
3557 )
3558 return inputs
3560 @staticmethod
3561 def _validate_axis(
3562 field_name: str,
3563 i: int,
3564 tensor_id: TensorId,
3565 a: int,
3566 axis: AnyAxis,
3567 valid_independent_refs: Dict[
3568 Tuple[TensorId, AxisId],
3569 Tuple[TensorDescr, AnyAxis, Union[int, ParameterizedSize]],
3570 ],
3571 ):
3572 if isinstance(axis, BatchAxis) or isinstance(
3573 axis.size, (int, ParameterizedSize, DataDependentSize)
3574 ):
3575 return
3576 elif not isinstance(axis.size, SizeReference):
3577 assert_never(axis.size)
3579 # validate axis.size SizeReference
3580 ref = (axis.size.tensor_id, axis.size.axis_id)
3581 if ref not in valid_independent_refs:
3582 raise ValueError(
3583 "Invalid tensor axis reference at"
3584 + f" {field_name}[{i}].axes[{a}].size: {axis.size}."
3585 )
3586 if ref == (tensor_id, axis.id):
3587 raise ValueError(
3588 "Self-referencing not allowed for"
3589 + f" {field_name}[{i}].axes[{a}].size: {axis.size}"
3590 )
3591 if axis.type == "channel":
3592 if valid_independent_refs[ref][1].type != "channel":
3593 raise ValueError(
3594 "A channel axis' size may only reference another fixed size"
3595 + " channel axis."
3596 )
3597 if isinstance(axis.channel_names, str) and "{i}" in axis.channel_names:
3598 ref_size = valid_independent_refs[ref][2]
3599 assert isinstance(ref_size, int), (
3600 "channel axis ref (another channel axis) has to specify fixed"
3601 + " size"
3602 )
3603 generated_channel_names = [
3604 axis.channel_names.format(i=i) for i in range(1, ref_size + 1)
3605 ]
3606 axis.channel_names = generated_channel_names
3608 if (ax_unit := getattr(axis, "unit", None)) != (
3609 ref_unit := getattr(valid_independent_refs[ref][1], "unit", None)
3610 ):
3611 raise ValueError(
3612 "The units of an axis and its reference axis need to match, but"
3613 + f" '{ax_unit}' != '{ref_unit}'."
3614 )
3615 ref_axis = valid_independent_refs[ref][1]
3616 if isinstance(ref_axis, BatchAxis):
3617 raise ValueError(
3618 f"Invalid reference axis '{ref_axis.id}' for {tensor_id}.{axis.id}"
3619 + " (a batch axis is not allowed as reference)."
3620 )
3622 if isinstance(axis, WithHalo):
3623 min_size = axis.size.get_size(axis, ref_axis, n=0)
3624 if (min_size - 2 * axis.halo) < 1:
3625 raise ValueError(
3626 f"axis {axis.id} with minimum size {min_size} is too small for halo"
3627 + f" {axis.halo}."
3628 )
3630 ref_halo = axis.halo * axis.scale / ref_axis.scale
3631 if ref_halo != int(ref_halo):
3632 raise ValueError(
3633 f"Inferred halo for {'.'.join(ref)} is not an integer ({ref_halo} ="
3634 + f" {tensor_id}.{axis.id}.halo {axis.halo}"
3635 + f" * {tensor_id}.{axis.id}.scale {axis.scale}"
3636 + f" / {'.'.join(ref)}.scale {ref_axis.scale})."
3637 )
3639 def validate_input_tensors(
3640 self,
3641 sources: Union[
3642 Sequence[NDArray[Any]], Mapping[TensorId, Optional[NDArray[Any]]]
3643 ],
3644 *,
3645 pad_inputs: Union[bool, Literal["allow"]] = True,
3646 crop_outputs: Union[bool, Literal["allow"]] = True,
3647 ) -> Mapping[TensorId, Optional[NDArray[Any]]]:
3648 """Check if the given input tensors match the model's input tensor descriptions.
3649 This includes checks of tensor shapes and dtypes, but not of the actual values.
3650 """
3651 if not isinstance(sources, collections.abc.Mapping):
3652 sources = {descr.id: tensor for descr, tensor in zip(self.inputs, sources)}
3654 tensors = {
3655 **{descr.id: (descr, sources.get(descr.id)) for descr in self.inputs},
3656 **{ # outputs are required for halo
3657 descr.id: (descr, None) for descr in self.outputs
3658 },
3659 }
3660 validate_tensors(tensors, pad_inputs=pad_inputs, crop_outputs=crop_outputs)
3662 return sources
3664 @model_validator(mode="after")
3665 def _validate_test_tensors(self) -> Self:
3666 if not get_validation_context().perform_io_checks:
3667 return self
3669 test_inputs = {
3670 descr.id: (
3671 descr,
3672 None if descr.test_tensor is None else load_array(descr.test_tensor),
3673 )
3674 for descr in self.inputs
3675 }
3676 test_outputs = {
3677 descr.id: (
3678 descr,
3679 None if descr.test_tensor is None else load_array(descr.test_tensor),
3680 )
3681 for descr in self.outputs
3682 }
3684 validate_tensors(
3685 {**test_inputs, **test_outputs},
3686 tensor_origin="test_tensor",
3687 pad_inputs="allow",
3688 crop_outputs="allow",
3689 )
3691 for rep_tol in self.config.bioimageio.reproducibility_tolerance:
3692 if not rep_tol.absolute_tolerance:
3693 continue
3695 if rep_tol.output_ids:
3696 out_arrays = {
3697 k: v[1] for k, v in test_outputs.items() if k in rep_tol.output_ids
3698 }
3699 else:
3700 out_arrays = {k: v[1] for k, v in test_outputs.items()}
3702 for out_id, array in out_arrays.items():
3703 if array is None:
3704 continue
3706 if rep_tol.absolute_tolerance > (max_test_value := array.max()) * 0.01:
3707 raise ValueError(
3708 "config.bioimageio.reproducibility_tolerance.absolute_tolerance="
3709 + f"{rep_tol.absolute_tolerance} > 0.01*{max_test_value}"
3710 + f" (1% of the maximum value of the test tensor '{out_id}')"
3711 )
3713 return self
3715 @model_validator(mode="after")
3716 def _validate_tensor_references_in_proc_kwargs(self, info: ValidationInfo) -> Self:
3717 ipt_refs = {t.id for t in self.inputs}
3718 missing_refs = [
3719 k["reference_tensor"]
3720 for k in [p.kwargs for ipt in self.inputs for p in ipt.preprocessing]
3721 + [p.kwargs for out in self.outputs for p in out.postprocessing]
3722 if "reference_tensor" in k
3723 and k["reference_tensor"] is not None
3724 and k["reference_tensor"] not in ipt_refs
3725 ]
3727 if missing_refs:
3728 raise ValueError(
3729 f"`reference_tensor`s {missing_refs} not found. Valid input tensor"
3730 + f" references are: {ipt_refs}."
3731 )
3733 return self
3735 name: Annotated[
3736 str,
3737 RestrictCharacters(string.ascii_letters + string.digits + "_+- ()"),
3738 MinLen(5),
3739 MaxLen(128),
3740 warn(MaxLen(64), "Name longer than 64 characters.", INFO),
3741 ]
3742 """A human-readable name of this model.
3743 It should be no longer than 64 characters
3744 and may only contain letter, number, underscore, minus, parentheses and spaces.
3745 We recommend to chose a name that refers to the model's task and image modality.
3746 """
3748 outputs: NotEmpty[Sequence[OutputTensorDescr]]
3749 """Describes the output tensors."""
3751 @field_validator("outputs", mode="after")
3752 @classmethod
3753 def _validate_tensor_ids(
3754 cls, outputs: Sequence[OutputTensorDescr], info: ValidationInfo
3755 ) -> Sequence[OutputTensorDescr]:
3756 tensor_ids = [
3757 t.id for t in info.data.get("inputs", []) + info.data.get("outputs", [])
3758 ]
3759 duplicate_tensor_ids: List[str] = []
3760 seen: Set[str] = set()
3761 for t in tensor_ids:
3762 if t in seen:
3763 duplicate_tensor_ids.append(t)
3765 seen.add(t)
3767 if duplicate_tensor_ids:
3768 raise ValueError(f"Duplicate tensor ids: {duplicate_tensor_ids}")
3770 return outputs
3772 @staticmethod
3773 def _get_axes_with_parameterized_size(
3774 io: Union[Sequence[InputTensorDescr], Sequence[OutputTensorDescr]],
3775 ):
3776 return {
3777 f"{t.id}.{a.id}": (t, a, a.size)
3778 for t in io
3779 for a in t.axes
3780 if not isinstance(a, BatchAxis) and isinstance(a.size, ParameterizedSize)
3781 }
3783 @staticmethod
3784 def _get_axes_with_independent_size(
3785 io: Union[Sequence[InputTensorDescr], Sequence[OutputTensorDescr]],
3786 ):
3787 return {
3788 (t.id, a.id): (t, a, a.size)
3789 for t in io
3790 for a in t.axes
3791 if not isinstance(a, BatchAxis)
3792 and isinstance(a.size, (int, ParameterizedSize))
3793 }
3795 @field_validator("outputs", mode="after")
3796 @classmethod
3797 def _validate_output_axes(
3798 cls, outputs: List[OutputTensorDescr], info: ValidationInfo
3799 ) -> List[OutputTensorDescr]:
3800 input_size_refs = cls._get_axes_with_independent_size(
3801 info.data.get("inputs", [])
3802 )
3803 output_size_refs = cls._get_axes_with_independent_size(outputs)
3805 for i, out in enumerate(outputs):
3806 valid_independent_refs: Dict[
3807 Tuple[TensorId, AxisId],
3808 Tuple[TensorDescr, AnyAxis, Union[int, ParameterizedSize]],
3809 ] = {
3810 **{
3811 (out.id, a.id): (out, a, a.size)
3812 for a in out.axes
3813 if not isinstance(a, BatchAxis)
3814 and isinstance(a.size, (int, ParameterizedSize))
3815 },
3816 **input_size_refs,
3817 **output_size_refs,
3818 }
3819 for a, ax in enumerate(out.axes):
3820 cls._validate_axis(
3821 "outputs",
3822 i,
3823 out.id,
3824 a,
3825 ax,
3826 valid_independent_refs=valid_independent_refs,
3827 )
3829 return outputs
3831 packaged_by: List[Author] = Field(
3832 default_factory=cast(Callable[[], List[Author]], list)
3833 )
3834 """The persons that have packaged and uploaded this model.
3835 Only required if those persons differ from the `authors`."""
3837 parent: Optional[LinkedModel] = None
3838 """The model from which this model is derived, e.g. by fine-tuning the weights."""
3840 @model_validator(mode="after")
3841 def _validate_parent_is_not_self(self) -> Self:
3842 if self.parent is not None and self.parent.id == self.id:
3843 raise ValueError("A model description may not reference itself as parent.")
3845 return self
3847 run_mode: Annotated[
3848 Optional[RunMode],
3849 warn(None, "Run mode '{value}' has limited support across consumer softwares."),
3850 ] = None
3851 """Custom run mode for this model: for more complex prediction procedures like test time
3852 data augmentation that currently cannot be expressed in the specification.
3853 No standard run modes are defined yet."""
3855 timestamp: Datetime = Field(default_factory=Datetime.now)
3856 """Timestamp in [ISO 8601](#https://en.wikipedia.org/wiki/ISO_8601) format
3857 with a few restrictions listed [here](https://docs.python.org/3/library/datetime.html#datetime.datetime.fromisoformat).
3858 (In Python a datetime object is valid, too)."""
3860 training_data: Annotated[
3861 Union[None, LinkedDataset, DatasetDescr, DatasetDescr02],
3862 Field(union_mode="left_to_right"),
3863 ] = None
3864 """The dataset used to train this model"""
3866 weights: Annotated[WeightsDescr, WrapSerializer(package_weights)]
3867 """The weights for this model.
3868 Weights can be given for different formats, but should otherwise be equivalent.
3869 The available weight formats determine which consumers can use this model."""
3871 config: Config = Field(default_factory=Config.model_construct)
3873 @model_validator(mode="after")
3874 def _add_default_cover(self) -> Self:
3875 if not get_validation_context().perform_io_checks or self.covers:
3876 return self
3878 try:
3879 generated_covers = generate_covers(
3880 [
3881 (t, load_array(t.test_tensor))
3882 for t in self.inputs
3883 if t.test_tensor is not None
3884 ],
3885 [
3886 (t, load_array(t.test_tensor))
3887 for t in self.outputs
3888 if t.test_tensor is not None
3889 ],
3890 )
3891 except Exception as e:
3892 issue_warning(
3893 "Failed to generate cover image(s): {e}",
3894 value=self.covers,
3895 msg_context=dict(e=e),
3896 field="covers",
3897 )
3898 else:
3899 self.covers.extend(generated_covers)
3901 return self
3903 def get_input_test_arrays(self) -> List[NDArray[Any]]:
3904 return self._get_test_arrays(self.inputs)
3906 def get_output_test_arrays(self) -> List[NDArray[Any]]:
3907 return self._get_test_arrays(self.outputs)
3909 @staticmethod
3910 def _get_test_arrays(
3911 io_descr: Union[Sequence[InputTensorDescr], Sequence[OutputTensorDescr]],
3912 ):
3913 ts: List[FileDescr] = []
3914 for d in io_descr:
3915 if d.test_tensor is None:
3916 raise ValueError(
3917 f"Failed to get test arrays: description of '{d.id}' is missing a `test_tensor`."
3918 )
3919 ts.append(d.test_tensor)
3921 data = [load_array(t) for t in ts]
3922 assert all(isinstance(d, np.ndarray) for d in data)
3923 return data
3925 @staticmethod
3926 def get_batch_size(tensor_sizes: Mapping[TensorId, Mapping[AxisId, int]]) -> int:
3927 batch_size = 1
3928 tensor_with_batchsize: Optional[TensorId] = None
3929 for tid in tensor_sizes:
3930 for aid, s in tensor_sizes[tid].items():
3931 if aid != BATCH_AXIS_ID or s == 1 or s == batch_size:
3932 continue
3934 if batch_size != 1:
3935 assert tensor_with_batchsize is not None
3936 raise ValueError(
3937 f"batch size mismatch for tensors '{tensor_with_batchsize}' ({batch_size}) and '{tid}' ({s})"
3938 )
3940 batch_size = s
3941 tensor_with_batchsize = tid
3943 return batch_size
3945 def get_output_tensor_sizes(
3946 self, input_sizes: Mapping[TensorId, Mapping[AxisId, int]]
3947 ) -> Dict[TensorId, Dict[AxisId, Union[int, _DataDepSize]]]:
3948 """Returns the tensor output sizes for given **input_sizes**.
3949 Only if **input_sizes** has a valid input shape, the tensor output size is exact.
3950 Otherwise it might be larger than the actual (valid) output"""
3951 batch_size = self.get_batch_size(input_sizes)
3952 ns = self.get_ns(input_sizes)
3954 tensor_sizes = self.get_tensor_sizes(ns, batch_size=batch_size)
3955 return tensor_sizes.outputs
3957 def get_ns(self, input_sizes: Mapping[TensorId, Mapping[AxisId, int]]):
3958 """get parameter `n` for each parameterized axis
3959 such that the valid input size is >= the given input size"""
3960 ret: Dict[Tuple[TensorId, AxisId], ParameterizedSize_N] = {}
3961 axes = {t.id: {a.id: a for a in t.axes} for t in self.inputs}
3962 for tid in input_sizes:
3963 for aid, s in input_sizes[tid].items():
3964 size_descr = axes[tid][aid].size
3965 if isinstance(size_descr, ParameterizedSize):
3966 ret[(tid, aid)] = size_descr.get_n(s)
3967 elif size_descr is None or isinstance(size_descr, (int, SizeReference)):
3968 pass
3969 else:
3970 assert_never(size_descr)
3972 return ret
3974 def get_tensor_sizes(
3975 self, ns: Mapping[Tuple[TensorId, AxisId], ParameterizedSize_N], batch_size: int
3976 ) -> _TensorSizes:
3977 axis_sizes = self.get_axis_sizes(ns, batch_size=batch_size)
3978 return _TensorSizes(
3979 {
3980 t: {
3981 aa: axis_sizes.inputs[(tt, aa)]
3982 for tt, aa in axis_sizes.inputs
3983 if tt == t
3984 }
3985 for t in {tt for tt, _ in axis_sizes.inputs}
3986 },
3987 {
3988 t: {
3989 aa: axis_sizes.outputs[(tt, aa)]
3990 for tt, aa in axis_sizes.outputs
3991 if tt == t
3992 }
3993 for t in {tt for tt, _ in axis_sizes.outputs}
3994 },
3995 )
3997 def get_axis_sizes(
3998 self,
3999 ns: Mapping[Tuple[TensorId, AxisId], ParameterizedSize_N],
4000 batch_size: Optional[int] = None,
4001 *,
4002 max_input_shape: Optional[Mapping[Tuple[TensorId, AxisId], int]] = None,
4003 ) -> _AxisSizes:
4004 """Determine input and output block shape for scale factors **ns**
4005 of parameterized input sizes.
4007 Args:
4008 ns: Scale factor `n` for each axis (keyed by (tensor_id, axis_id))
4009 that is parameterized as `size = min + n * step`.
4010 batch_size: The desired size of the batch dimension.
4011 If given **batch_size** overwrites any batch size present in
4012 **max_input_shape**. Default 1.
4013 max_input_shape: Limits the derived block shapes.
4014 Each axis for which the input size, parameterized by `n`, is larger
4015 than **max_input_shape** is set to the minimal value `n_min` for which
4016 this is still true.
4017 Use this for small input samples or large values of **ns**.
4018 Or simply whenever you know the full input shape.
4020 Returns:
4021 Resolved axis sizes for model inputs and outputs.
4022 """
4023 max_input_shape = max_input_shape or {}
4024 if batch_size is None:
4025 for (_t_id, a_id), s in max_input_shape.items():
4026 if a_id == BATCH_AXIS_ID:
4027 batch_size = s
4028 break
4029 else:
4030 batch_size = 1
4032 all_axes = {
4033 t.id: {a.id: a for a in t.axes} for t in chain(self.inputs, self.outputs)
4034 }
4036 inputs: Dict[Tuple[TensorId, AxisId], int] = {}
4037 outputs: Dict[Tuple[TensorId, AxisId], Union[int, _DataDepSize]] = {}
4039 def get_axis_size(a: Union[InputAxis, OutputAxis]):
4040 if isinstance(a, BatchAxis):
4041 if (t_descr.id, a.id) in ns:
4042 logger.warning(
4043 "Ignoring unexpected size increment factor (n) for batch axis"
4044 + " of tensor '{}'.",
4045 t_descr.id,
4046 )
4047 return batch_size
4048 elif isinstance(a.size, int):
4049 if (t_descr.id, a.id) in ns:
4050 logger.warning(
4051 "Ignoring unexpected size increment factor (n) for fixed size"
4052 + " axis '{}' of tensor '{}'.",
4053 a.id,
4054 t_descr.id,
4055 )
4056 return a.size
4057 elif isinstance(a.size, ParameterizedSize):
4058 if (t_descr.id, a.id) not in ns:
4059 raise ValueError(
4060 "Size increment factor (n) missing for parametrized axis"
4061 + f" '{a.id}' of tensor '{t_descr.id}'."
4062 )
4063 n = ns[(t_descr.id, a.id)]
4064 s_max = max_input_shape.get((t_descr.id, a.id))
4065 if s_max is not None:
4066 n = min(n, a.size.get_n(s_max))
4068 return a.size.get_size(n)
4070 elif isinstance(a.size, SizeReference):
4071 if (t_descr.id, a.id) in ns:
4072 logger.warning(
4073 "Ignoring unexpected size increment factor (n) for axis '{}'"
4074 + " of tensor '{}' with size reference.",
4075 a.id,
4076 t_descr.id,
4077 )
4078 assert not isinstance(a, BatchAxis)
4079 ref_axis = all_axes[a.size.tensor_id][a.size.axis_id]
4080 assert not isinstance(ref_axis, BatchAxis)
4081 ref_key = (a.size.tensor_id, a.size.axis_id)
4082 ref_size = inputs.get(ref_key, outputs.get(ref_key))
4083 assert ref_size is not None, ref_key
4084 assert not isinstance(ref_size, _DataDepSize), ref_key
4085 return a.size.get_size(
4086 axis=a,
4087 ref_axis=ref_axis,
4088 ref_size=ref_size,
4089 )
4090 elif isinstance(a.size, DataDependentSize):
4091 if (t_descr.id, a.id) in ns:
4092 logger.warning(
4093 "Ignoring unexpected increment factor (n) for data dependent"
4094 + " size axis '{}' of tensor '{}'.",
4095 a.id,
4096 t_descr.id,
4097 )
4098 return _DataDepSize(a.size.min, a.size.max)
4099 else:
4100 assert_never(a.size)
4102 # first resolve all , but the `SizeReference` input sizes
4103 for t_descr in self.inputs:
4104 for a in t_descr.axes:
4105 if not isinstance(a.size, SizeReference):
4106 s = get_axis_size(a)
4107 assert not isinstance(s, _DataDepSize)
4108 inputs[t_descr.id, a.id] = s
4110 # resolve all other input axis sizes
4111 for t_descr in self.inputs:
4112 for a in t_descr.axes:
4113 if isinstance(a.size, SizeReference):
4114 s = get_axis_size(a)
4115 assert not isinstance(s, _DataDepSize)
4116 inputs[t_descr.id, a.id] = s
4118 # resolve all output axis sizes
4119 for t_descr in self.outputs:
4120 for a in t_descr.axes:
4121 assert not isinstance(a.size, ParameterizedSize)
4122 s = get_axis_size(a)
4123 outputs[t_descr.id, a.id] = s
4125 return _AxisSizes(inputs=inputs, outputs=outputs)
4127 @model_validator(mode="before")
4128 @classmethod
4129 def _convert(cls, data: Dict[str, Any]) -> Dict[str, Any]:
4130 cls.convert_from_old_format_wo_validation(data)
4131 return data
4133 @classmethod
4134 def convert_from_old_format_wo_validation(cls, data: Dict[str, Any]) -> None:
4135 """Convert metadata following an older format version to this classes' format
4136 without validating the result.
4137 """
4138 if (
4139 data.get("type") == "model"
4140 and isinstance(fv := data.get("format_version"), str)
4141 and fv.count(".") == 2
4142 ):
4143 fv_parts = fv.split(".")
4144 if any(not p.isdigit() for p in fv_parts):
4145 return
4147 fv_tuple = tuple(map(int, fv_parts))
4149 assert cls.implemented_format_version_tuple[0:2] == (0, 5)
4150 if fv_tuple[:2] in ((0, 3), (0, 4)):
4151 m04 = _ModelDescr_v0_4.load(data)
4152 if isinstance(m04, InvalidDescr):
4153 try:
4154 updated = _model_conv.convert_as_dict(
4155 m04 # pyright: ignore[reportArgumentType]
4156 )
4157 except Exception as e:
4158 logger.error(
4159 "Failed to convert from invalid model 0.4 description."
4160 + f"\nerror: {e}"
4161 + "\nProceeding with model 0.5 validation without conversion."
4162 )
4163 updated = None
4164 else:
4165 updated = _model_conv.convert_as_dict(m04)
4167 if updated is not None:
4168 data.clear()
4169 data.update(updated)
4171 elif fv_tuple[:2] == (0, 5):
4172 # bump patch version
4173 data["format_version"] = cls.implemented_format_version
4176class _ModelConv(Converter[_ModelDescr_v0_4, ModelDescr]):
4177 def _convert(
4178 self, src: _ModelDescr_v0_4, tgt: "type[ModelDescr] | type[dict[str, Any]]"
4179 ) -> "ModelDescr | dict[str, Any]":
4180 name = "".join(
4181 c if c in string.ascii_letters + string.digits + "_+- ()" else " "
4182 for c in src.name
4183 )
4185 def conv_authors(auths: Optional[Sequence[_Author_v0_4]]):
4186 conv = (
4187 _author_conv.convert if TYPE_CHECKING else _author_conv.convert_as_dict
4188 )
4189 return None if auths is None else [conv(a) for a in auths]
4191 if TYPE_CHECKING:
4192 arch_file_conv = _arch_file_conv.convert
4193 arch_lib_conv = _arch_lib_conv.convert
4194 else:
4195 arch_file_conv = _arch_file_conv.convert_as_dict
4196 arch_lib_conv = _arch_lib_conv.convert_as_dict
4198 input_size_refs = {
4199 ipt.name: {
4200 a: s
4201 for a, s in zip(
4202 ipt.axes,
4203 (
4204 ipt.shape.min
4205 if isinstance(ipt.shape, _ParameterizedInputShape_v0_4)
4206 else ipt.shape
4207 ),
4208 )
4209 }
4210 for ipt in src.inputs
4211 if ipt.shape
4212 }
4213 output_size_refs = {
4214 **{
4215 out.name: {a: s for a, s in zip(out.axes, out.shape)}
4216 for out in src.outputs
4217 if not isinstance(out.shape, _ImplicitOutputShape_v0_4)
4218 },
4219 **input_size_refs,
4220 }
4222 return tgt(
4223 attachments=(
4224 []
4225 if src.attachments is None
4226 else [FileDescr(source=f) for f in src.attachments.files]
4227 ),
4228 authors=[_author_conv.convert_as_dict(a) for a in src.authors], # pyright: ignore[reportArgumentType]
4229 cite=[{"text": c.text, "doi": c.doi, "url": c.url} for c in src.cite], # pyright: ignore[reportArgumentType]
4230 config=src.config, # pyright: ignore[reportArgumentType]
4231 covers=src.covers,
4232 description=src.description,
4233 documentation=src.documentation,
4234 format_version="0.5.10",
4235 git_repo=src.git_repo, # pyright: ignore[reportArgumentType]
4236 icon=src.icon,
4237 id=None if src.id is None else ModelId(src.id),
4238 id_emoji=src.id_emoji,
4239 license=src.license, # type: ignore
4240 links=src.links,
4241 maintainers=[_maintainer_conv.convert_as_dict(m) for m in src.maintainers], # pyright: ignore[reportArgumentType]
4242 name=name,
4243 tags=src.tags,
4244 type=src.type,
4245 uploader=src.uploader,
4246 version=src.version,
4247 inputs=[ # pyright: ignore[reportArgumentType]
4248 _input_tensor_conv.convert_as_dict(ipt, tt, st, input_size_refs)
4249 for ipt, tt, st in zip(
4250 src.inputs,
4251 src.test_inputs,
4252 src.sample_inputs or [None] * len(src.test_inputs),
4253 )
4254 ],
4255 outputs=[ # pyright: ignore[reportArgumentType]
4256 _output_tensor_conv.convert_as_dict(out, tt, st, output_size_refs)
4257 for out, tt, st in zip(
4258 src.outputs,
4259 src.test_outputs,
4260 src.sample_outputs or [None] * len(src.test_outputs),
4261 )
4262 ],
4263 parent=(
4264 None
4265 if src.parent is None
4266 else LinkedModel(
4267 id=ModelId(
4268 str(src.parent.id)
4269 + (
4270 ""
4271 if src.parent.version_number is None
4272 else f"/{src.parent.version_number}"
4273 )
4274 )
4275 )
4276 ),
4277 training_data=(
4278 None
4279 if src.training_data is None
4280 else (
4281 LinkedDataset(
4282 id=DatasetId(
4283 str(src.training_data.id)
4284 + (
4285 ""
4286 if src.training_data.version_number is None
4287 else f"/{src.training_data.version_number}"
4288 )
4289 )
4290 )
4291 if isinstance(src.training_data, LinkedDataset02)
4292 else src.training_data
4293 )
4294 ),
4295 packaged_by=[_author_conv.convert_as_dict(a) for a in src.packaged_by], # pyright: ignore[reportArgumentType]
4296 run_mode=src.run_mode,
4297 timestamp=src.timestamp,
4298 weights=(WeightsDescr if TYPE_CHECKING else dict)(
4299 keras_hdf5=(w := src.weights.keras_hdf5)
4300 and (KerasHdf5WeightsDescr if TYPE_CHECKING else dict)(
4301 authors=conv_authors(w.authors),
4302 source=w.source,
4303 tensorflow_version=w.tensorflow_version or Version("1.15"),
4304 parent=w.parent,
4305 ),
4306 onnx=(w := src.weights.onnx)
4307 and (OnnxWeightsDescr if TYPE_CHECKING else dict)(
4308 source=w.source,
4309 authors=conv_authors(w.authors),
4310 parent=w.parent,
4311 opset_version=w.opset_version or 15,
4312 ),
4313 pytorch_state_dict=(w := src.weights.pytorch_state_dict)
4314 and (PytorchStateDictWeightsDescr if TYPE_CHECKING else dict)(
4315 source=w.source,
4316 authors=conv_authors(w.authors),
4317 parent=w.parent,
4318 architecture=(
4319 arch_file_conv(
4320 w.architecture,
4321 w.architecture_sha256,
4322 w.kwargs,
4323 )
4324 if isinstance(w.architecture, _CallableFromFile_v0_4)
4325 else arch_lib_conv(w.architecture, w.kwargs)
4326 ),
4327 pytorch_version=w.pytorch_version or Version("1.10"),
4328 dependencies=(
4329 None
4330 if w.dependencies is None
4331 else (FileDescr if TYPE_CHECKING else dict)(
4332 source=cast(
4333 FileSource,
4334 str(deps := w.dependencies)[
4335 (
4336 len("conda:")
4337 if str(deps).startswith("conda:")
4338 else 0
4339 ) :
4340 ],
4341 )
4342 )
4343 ),
4344 ),
4345 tensorflow_js=(w := src.weights.tensorflow_js)
4346 and (TensorflowJsWeightsDescr if TYPE_CHECKING else dict)(
4347 source=w.source,
4348 authors=conv_authors(w.authors),
4349 parent=w.parent,
4350 tensorflow_version=w.tensorflow_version or Version("1.15"),
4351 ),
4352 tensorflow_saved_model_bundle=(
4353 w := src.weights.tensorflow_saved_model_bundle
4354 )
4355 and (TensorflowSavedModelBundleWeightsDescr if TYPE_CHECKING else dict)(
4356 authors=conv_authors(w.authors),
4357 parent=w.parent,
4358 source=w.source,
4359 tensorflow_version=w.tensorflow_version or Version("1.15"),
4360 dependencies=(
4361 None
4362 if w.dependencies is None
4363 else (FileDescr if TYPE_CHECKING else dict)(
4364 source=cast(
4365 FileSource,
4366 (
4367 str(w.dependencies)[len("conda:") :]
4368 if str(w.dependencies).startswith("conda:")
4369 else str(w.dependencies)
4370 ),
4371 )
4372 )
4373 ),
4374 ),
4375 torchscript=(w := src.weights.torchscript)
4376 and (TorchscriptWeightsDescr if TYPE_CHECKING else dict)(
4377 source=w.source,
4378 authors=conv_authors(w.authors),
4379 parent=w.parent,
4380 pytorch_version=w.pytorch_version or Version("1.10"),
4381 ),
4382 ),
4383 )
4386_model_conv = _ModelConv(_ModelDescr_v0_4, ModelDescr)
4389# create better cover images for 3d data and non-image outputs
4390def generate_covers(
4391 inputs: Sequence[Tuple[InputTensorDescr, NDArray[Any]]],
4392 outputs: Sequence[Tuple[OutputTensorDescr, NDArray[Any]]],
4393) -> List[Path]:
4394 def squeeze(
4395 data: NDArray[Any], axes: Sequence[AnyAxis]
4396 ) -> Tuple[NDArray[Any], List[AnyAxis]]:
4397 """apply numpy.ndarray.squeeze while keeping track of the axis descriptions remaining"""
4398 if data.ndim != len(axes):
4399 raise ValueError(
4400 f"tensor shape {data.shape} does not match described axes"
4401 + f" {[a.id for a in axes]}"
4402 )
4404 axes = [deepcopy(a) for a, s in zip(axes, data.shape) if s != 1]
4405 return data.squeeze(), axes
4407 def normalize(
4408 data: NDArray[Any], axis: Optional[Tuple[int, ...]], eps: float = 1e-7
4409 ) -> NDArray[np.float32]:
4410 data = data.astype("float32")
4411 data -= data.min(axis=axis, keepdims=True)
4412 data /= data.max(axis=axis, keepdims=True) + eps
4413 return data
4415 def to_2d_image(data: NDArray[Any], axes: Sequence[AnyAxis]):
4416 original_shape = data.shape
4417 original_axes = list(axes)
4418 data, axes = squeeze(data, axes)
4420 # take slice fom any batch or index axis if needed
4421 # and convert the first channel axis and take a slice from any additional channel axes
4422 slices: Tuple[slice, ...] = ()
4423 ndim = data.ndim
4424 ndim_need = 3 if any(isinstance(a, ChannelAxis) for a in axes) else 2
4425 has_c_axis = False
4426 for i, a in enumerate(axes):
4427 s = data.shape[i]
4428 assert s > 1
4429 if (
4430 isinstance(a, (BatchAxis, IndexInputAxis, IndexOutputAxis))
4431 and ndim > ndim_need
4432 ):
4433 data = data[slices + (slice(s // 2 - 1, s // 2),)]
4434 ndim -= 1
4435 elif isinstance(a, ChannelAxis):
4436 if has_c_axis:
4437 # second channel axis
4438 data = data[slices + (slice(0, 1),)]
4439 ndim -= 1
4440 else:
4441 has_c_axis = True
4442 if s == 2:
4443 # visualize two channels with cyan and magenta
4444 data = np.concatenate(
4445 [
4446 data[slices + (slice(1, 2),)],
4447 data[slices + (slice(0, 1),)],
4448 (
4449 data[slices + (slice(0, 1),)]
4450 + data[slices + (slice(1, 2),)]
4451 )
4452 / 2, # TODO: take maximum instead?
4453 ],
4454 axis=i,
4455 )
4456 elif data.shape[i] == 3:
4457 pass # visualize 3 channels as RGB
4458 else:
4459 # visualize first 3 channels as RGB
4460 data = data[slices + (slice(3),)]
4462 assert data.shape[i] == 3
4464 slices += (slice(None),)
4466 data, axes = squeeze(data, axes)
4467 assert len(axes) == ndim
4468 # take slice from z axis if needed
4469 slices = ()
4470 if ndim > ndim_need:
4471 for i, a in enumerate(axes):
4472 s = data.shape[i]
4473 if a.id == AxisId("z"):
4474 data = data[slices + (slice(s // 2 - 1, s // 2),)]
4475 data, axes = squeeze(data, axes)
4476 ndim -= 1
4477 break
4479 slices += (slice(None),)
4481 # take slice from any space or time axis
4482 slices = ()
4484 for i, a in enumerate(axes):
4485 if ndim <= ndim_need:
4486 break
4488 s = data.shape[i]
4489 assert s > 1
4490 if isinstance(
4491 a, (SpaceInputAxis, SpaceOutputAxis, TimeInputAxis, TimeOutputAxis)
4492 ):
4493 data = data[slices + (slice(s // 2 - 1, s // 2),)]
4494 ndim -= 1
4496 slices += (slice(None),)
4498 del slices
4499 data, axes = squeeze(data, axes)
4500 assert len(axes) == ndim
4502 if (has_c_axis and ndim != 3) or (not has_c_axis and ndim != 2):
4503 raise ValueError(
4504 f"Failed to construct cover image from shape {original_shape} with axes {[a.id for a in original_axes]}."
4505 )
4507 if not has_c_axis:
4508 assert ndim == 2
4509 data = np.repeat(data[:, :, None], 3, axis=2)
4510 axes.append(ChannelAxis(channel_names=list("RGB")))
4511 ndim += 1
4513 assert ndim == 3
4515 # transpose axis order such that longest axis comes first...
4516 axis_order: List[int] = list(np.argsort(list(data.shape)))
4517 axis_order.reverse()
4518 # ... and channel axis is last
4519 c = [i for i in range(3) if isinstance(axes[i], ChannelAxis)][0]
4520 axis_order.append(axis_order.pop(c))
4521 axes = [axes[ao] for ao in axis_order]
4522 data = data.transpose(axis_order)
4524 # h, w = data.shape[:2]
4525 # if h / w in (1.0 or 2.0):
4526 # pass
4527 # elif h / w < 2:
4528 # TODO: enforce 2:1 or 1:1 aspect ratio for generated cover images
4530 norm_along = (
4531 tuple(i for i, a in enumerate(axes) if a.type in ("space", "time")) or None
4532 )
4533 # normalize the data and map to 8 bit
4534 data = normalize(data, norm_along)
4535 data = (data * 255).astype("uint8")
4537 return data
4539 def create_diagonal_split_image(im0: NDArray[Any], im1: NDArray[Any]):
4540 assert im0.dtype == im1.dtype == np.uint8
4541 assert im0.shape == im1.shape
4542 assert im0.ndim == 3
4543 N, M, C = im0.shape
4544 assert C == 3
4545 out = np.ones((N, M, C), dtype="uint8")
4546 for c in range(C):
4547 outc = np.tril(im0[..., c])
4548 mask = outc == 0
4549 outc[mask] = np.triu(im1[..., c])[mask]
4550 out[..., c] = outc
4552 return out
4554 if not inputs:
4555 raise ValueError("Missing test input tensor for cover generation.")
4557 if not outputs:
4558 raise ValueError("Missing test output tensor for cover generation.")
4560 ipt_descr, ipt = inputs[0]
4561 out_descr, out = outputs[0]
4563 ipt_img = to_2d_image(ipt, ipt_descr.axes)
4564 out_img = to_2d_image(out, out_descr.axes)
4566 cover_folder = Path(mkdtemp())
4567 if ipt_img.shape == out_img.shape:
4568 covers = [cover_folder / "cover.png"]
4569 imwrite(covers[0], create_diagonal_split_image(ipt_img, out_img))
4570 else:
4571 covers = [cover_folder / "input.png", cover_folder / "output.png"]
4572 imwrite(covers[0], ipt_img)
4573 imwrite(covers[1], out_img)
4575 return covers