Coverage for src/bioimageio/spec/model/v0_5.py: 71%
1684 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-08 12:40 +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, TypeAlias, 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_package,
79 package_file_descr_serializer,
80)
81from .._internal.io_utils import load_array, open_bioimageio_yaml
82from .._internal.node_converter import Converter
83from .._internal.type_guards import is_dict, is_sequence
84from .._internal.types import (
85 FAIR,
86 AbsoluteTolerance,
87 LowerCaseIdentifier,
88 LowerCaseIdentifierAnno,
89 MismatchedElementsPerMillion,
90 RelativeTolerance,
91)
92from .._internal.types import Datetime as Datetime
93from .._internal.types import Identifier as Identifier
94from .._internal.types import NotEmpty as NotEmpty
95from .._internal.types import SiUnit as SiUnit
96from .._internal.url import HttpUrl as HttpUrl
97from .._internal.utils import try_all_raise_last
98from .._internal.validation_context import get_validation_context
99from .._internal.validator_annotations import RestrictCharacters
100from .._internal.version_type import Version as Version
101from .._internal.warning_levels import INFO
102from ..dataset.v0_2 import DatasetDescr as DatasetDescr02
103from ..dataset.v0_2 import LinkedDataset as LinkedDataset02
104from ..dataset.v0_3 import DatasetDescr as DatasetDescr
105from ..dataset.v0_3 import DatasetId as DatasetId
106from ..dataset.v0_3 import LinkedDataset as LinkedDataset
107from ..dataset.v0_3 import Uploader as Uploader
108from ..generic._v0_3_converter import convert_plain_covers_and_docs_and_icon
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 FileDescr_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 ModelDescr as _ModelDescr04
131from .v0_4 import Author as _Author_v0_4
132from .v0_4 import BinarizeDescr as _BinarizeDescr_v0_4
133from .v0_4 import CallableFromDepencency as CallableFromDepencency
134from .v0_4 import CallableFromDepencency as _CallableFromDepencency_v0_4
135from .v0_4 import CallableFromFile as _CallableFromFile_v0_4
136from .v0_4 import ClipDescr as _ClipDescr_v0_4
137from .v0_4 import ImplicitOutputShape as _ImplicitOutputShape_v0_4
138from .v0_4 import InputTensorDescr as _InputTensorDescr_v0_4
139from .v0_4 import KnownRunMode as KnownRunMode
140from .v0_4 import ModelDescr as _ModelDescr_v0_4
141from .v0_4 import OutputTensorDescr as _OutputTensorDescr_v0_4
142from .v0_4 import ParameterizedInputShape as _ParameterizedInputShape_v0_4
143from .v0_4 import PostprocessingDescr as _PostprocessingDescr_v0_4
144from .v0_4 import PreprocessingDescr as _PreprocessingDescr_v0_4
145from .v0_4 import RunMode as RunMode
146from .v0_4 import ScaleLinearDescr as _ScaleLinearDescr_v0_4
147from .v0_4 import ScaleMeanVarianceDescr as _ScaleMeanVarianceDescr_v0_4
148from .v0_4 import ScaleRangeDescr as _ScaleRangeDescr_v0_4
149from .v0_4 import SigmoidDescr as _SigmoidDescr_v0_4
150from .v0_4 import TensorName as _TensorName_v0_4
151from .v0_4 import ZeroMeanUnitVarianceDescr as _ZeroMeanUnitVarianceDescr_v0_4
152from .v0_4 import package_weights
154SpaceUnit = Literal[
155 "attometer",
156 "angstrom",
157 "centimeter",
158 "decimeter",
159 "exameter",
160 "femtometer",
161 "foot",
162 "gigameter",
163 "hectometer",
164 "inch",
165 "kilometer",
166 "megameter",
167 "meter",
168 "micrometer",
169 "mile",
170 "millimeter",
171 "nanometer",
172 "parsec",
173 "petameter",
174 "picometer",
175 "terameter",
176 "yard",
177 "yoctometer",
178 "yottameter",
179 "zeptometer",
180 "zettameter",
181]
182"""Space unit compatible to the [OME-Zarr axes specification 0.5](https://ngff.openmicroscopy.org/0.5/#axes-md)"""
184TimeUnit = Literal[
185 "attosecond",
186 "centisecond",
187 "day",
188 "decisecond",
189 "exasecond",
190 "femtosecond",
191 "gigasecond",
192 "hectosecond",
193 "hour",
194 "kilosecond",
195 "megasecond",
196 "microsecond",
197 "millisecond",
198 "minute",
199 "nanosecond",
200 "petasecond",
201 "picosecond",
202 "second",
203 "terasecond",
204 "yoctosecond",
205 "yottasecond",
206 "zeptosecond",
207 "zettasecond",
208]
209"""Time unit compatible to the [OME-Zarr axes specification 0.5](https://ngff.openmicroscopy.org/0.5/#axes-md)"""
211AxisType = Literal["batch", "channel", "index", "time", "space"]
213_AXIS_TYPE_MAP: Mapping[str, AxisType] = {
214 "b": "batch",
215 "t": "time",
216 "i": "index",
217 "c": "channel",
218 "x": "space",
219 "y": "space",
220 "z": "space",
221}
223_AXIS_ID_MAP = {
224 "b": "batch",
225 "t": "time",
226 "i": "index",
227 "c": "channel",
228}
230WeightsFormat = Literal[
231 "keras_hdf5",
232 "keras_v3",
233 "onnx",
234 "pytorch_state_dict",
235 "tensorflow_js",
236 "tensorflow_saved_model_bundle",
237 "torchscript",
238]
241class TensorId(LowerCaseIdentifier):
242 root_model: ClassVar[Type[RootModel[Any]]] = RootModel[
243 Annotated[LowerCaseIdentifierAnno, MaxLen(32)]
244 ]
247def _normalize_axis_id(a: str):
248 a = str(a)
249 normalized = _AXIS_ID_MAP.get(a, a)
250 if a != normalized:
251 logger.opt(depth=3).warning(
252 "Normalized axis id from '{}' to '{}'.", a, normalized
253 )
254 return normalized
257class AxisId(LowerCaseIdentifier):
258 root_model: ClassVar[Type[RootModel[Any]]] = RootModel[
259 Annotated[
260 LowerCaseIdentifierAnno,
261 MaxLen(16),
262 AfterValidator(_normalize_axis_id),
263 ]
264 ]
267def _is_batch(a: str) -> bool:
268 return str(a) == "batch"
271def _is_not_batch(a: str) -> bool:
272 return not _is_batch(a)
275NonBatchAxisId = Annotated[AxisId, Predicate(_is_not_batch)]
277PreprocessingId = Literal[
278 "binarize",
279 "clip",
280 "ensure_dtype",
281 "fixed_zero_mean_unit_variance",
282 "scale_linear",
283 "scale_range",
284 "sigmoid",
285 "softmax",
286]
287PostprocessingId = Literal[
288 "binarize",
289 "clip",
290 "custom",
291 "ensure_dtype",
292 "fixed_zero_mean_unit_variance",
293 "scale_linear",
294 "scale_mean_variance",
295 "scale_range",
296 "sigmoid",
297 "softmax",
298 "zero_mean_unit_variance",
299]
302SAME_AS_TYPE = "<same as type>"
305ParameterizedSize_N: TypeAlias = int
306"""
307Annotates an integer to calculate a concrete axis size from a `ParameterizedSize`.
308"""
311class ParameterizedSize(Node):
312 """Describes a range of valid tensor axis sizes as `size = min + n*step`.
314 - **min** and **step** are given by the model description.
315 - All blocksize paramters n = 0,1,2,... yield a valid `size`.
316 - A greater blocksize paramter n = 0,1,2,... results in a greater **size**.
317 This allows to adjust the axis size more generically.
318 """
320 N: ClassVar[Type[int]] = ParameterizedSize_N
321 """Positive integer to parameterize this axis"""
323 min: Annotated[int, Gt(0)]
324 step: Annotated[int, Gt(0)]
326 def validate_size(self, size: int, msg_prefix: str = "") -> int:
327 if size < self.min:
328 raise ValueError(
329 f"{msg_prefix}size {size} < {self.min} (minimum axis size)"
330 )
331 if (size - self.min) % self.step != 0:
332 raise ValueError(
333 f"{msg_prefix}size {size} is not parameterized by `min + n*step` ="
334 + f" `{self.min} + n*{self.step}`"
335 )
337 return size
339 def get_size(self, n: ParameterizedSize_N) -> int:
340 return self.min + self.step * n
342 def get_n(self, s: int) -> ParameterizedSize_N:
343 """return smallest n parameterizing a size greater or equal than `s`"""
344 return ceil((s - self.min) / self.step)
347class DataDependentSize(Node):
348 min: Annotated[int, Gt(0)] = 1
349 max: Annotated[Optional[int], Gt(1)] = None
351 @model_validator(mode="after")
352 def _validate_max_gt_min(self):
353 if self.max is not None and self.min >= self.max:
354 raise ValueError(f"expected `min` < `max`, but got {self.min}, {self.max}")
356 return self
358 def validate_size(self, size: int, msg_prefix: str = "") -> int:
359 if size < self.min:
360 raise ValueError(f"{msg_prefix}size {size} < {self.min}")
362 if self.max is not None and size > self.max:
363 raise ValueError(f"{msg_prefix}size {size} > {self.max}")
365 return size
368class SizeReference(Node):
369 """A tensor axis size (extent in pixels/frames) defined in relation to a reference axis.
371 `axis.size = reference.size * reference.scale / axis.scale + offset`
373 Note:
374 1. The axis and the referenced axis need to have the same unit (or no unit).
375 2. Batch axes may not be referenced.
376 3. Fractions are rounded down.
377 4. If the reference axis is `concatenable` the referencing axis is assumed to be
378 `concatenable` as well with the same block order.
380 Example:
381 An unisotropic input image of w*h=100*49 pixels depicts a phsical space of 200*196mm².
382 Let's assume that we want to express the image height h in relation to its width w
383 instead of only accepting input images of exactly 100*49 pixels
384 (for example to express a range of valid image shapes by parametrizing w, see `ParameterizedSize`).
386 >>> w = SpaceInputAxis(id=AxisId("w"), size=100, unit="millimeter", scale=2)
387 >>> h = SpaceInputAxis(
388 ... id=AxisId("h"),
389 ... size=SizeReference(tensor_id=TensorId("input"), axis_id=AxisId("w"), offset=-1),
390 ... unit="millimeter",
391 ... scale=4,
392 ... )
393 >>> print(h.size.get_size(h, w))
394 49
396 ⇒ h = w * w.scale / h.scale + offset = 100 * 2mm / 4mm - 1 = 49
397 """
399 tensor_id: TensorId
400 """tensor id of the reference axis"""
402 axis_id: AxisId
403 """axis id of the reference axis"""
405 offset: StrictInt = 0
407 def get_size(
408 self,
409 axis: Union[
410 ChannelAxis,
411 IndexInputAxis,
412 IndexOutputAxis,
413 TimeInputAxis,
414 SpaceInputAxis,
415 TimeOutputAxis,
416 TimeOutputAxisWithHalo,
417 SpaceOutputAxis,
418 SpaceOutputAxisWithHalo,
419 ],
420 ref_axis: Union[
421 ChannelAxis,
422 IndexInputAxis,
423 IndexOutputAxis,
424 TimeInputAxis,
425 SpaceInputAxis,
426 TimeOutputAxis,
427 TimeOutputAxisWithHalo,
428 SpaceOutputAxis,
429 SpaceOutputAxisWithHalo,
430 ],
431 n: ParameterizedSize_N = 0,
432 ref_size: Optional[int] = None,
433 ):
434 """Compute the concrete size for a given axis and its reference axis.
436 Args:
437 axis: The axis this [SizeReference][] is the size of.
438 ref_axis: The reference axis to compute the size from.
439 n: If the **ref_axis** is parameterized (of type `ParameterizedSize`)
440 and no fixed **ref_size** is given,
441 **n** is used to compute the size of the parameterized **ref_axis**.
442 ref_size: Overwrite the reference size instead of deriving it from
443 **ref_axis**
444 (**ref_axis.scale** is still used; any given **n** is ignored).
445 """
446 assert axis.size == self, (
447 "Given `axis.size` is not defined by this `SizeReference`"
448 )
450 assert ref_axis.id == self.axis_id, (
451 f"Expected `ref_axis.id` to be {self.axis_id}, but got {ref_axis.id}."
452 )
454 assert axis.unit == ref_axis.unit, (
455 "`SizeReference` requires `axis` and `ref_axis` to have the same `unit`,"
456 f" but {axis.unit}!={ref_axis.unit}"
457 )
458 if ref_size is None:
459 if isinstance(ref_axis.size, (int, float)):
460 ref_size = ref_axis.size
461 elif isinstance(ref_axis.size, ParameterizedSize):
462 ref_size = ref_axis.size.get_size(n)
463 elif isinstance(ref_axis.size, DataDependentSize):
464 raise ValueError(
465 "Reference axis referenced in `SizeReference` may not be a `DataDependentSize`."
466 )
467 elif isinstance(ref_axis.size, SizeReference):
468 raise ValueError(
469 "Reference axis referenced in `SizeReference` may not be sized by a"
470 + " `SizeReference` itself."
471 )
472 else:
473 assert_never(ref_axis.size)
475 return int(ref_size * ref_axis.scale / axis.scale + self.offset)
477 @staticmethod
478 def _get_unit(
479 axis: Union[
480 ChannelAxis,
481 IndexInputAxis,
482 IndexOutputAxis,
483 TimeInputAxis,
484 SpaceInputAxis,
485 TimeOutputAxis,
486 TimeOutputAxisWithHalo,
487 SpaceOutputAxis,
488 SpaceOutputAxisWithHalo,
489 ],
490 ):
491 return axis.unit
494class AxisBase(NodeWithExplicitlySetFields):
495 id: AxisId
496 """An axis id unique across all axes of one tensor."""
498 description: Annotated[str, MaxLen(128)] = ""
499 """A short description of this axis beyond its type and id."""
502class WithHalo(Node):
503 halo: Annotated[int, Ge(1)]
504 """The halo should be cropped from the output tensor to avoid boundary effects.
505 It is to be cropped from both sides, i.e. `size_after_crop = size - 2 * halo`.
506 To document a halo that is already cropped by the model use `size.offset` instead."""
508 size: Annotated[
509 SizeReference,
510 Field(
511 examples=[
512 10,
513 SizeReference(
514 tensor_id=TensorId("t"), axis_id=AxisId("a"), offset=5
515 ).model_dump(mode="json"),
516 ]
517 ),
518 ]
519 """reference to another axis with an optional offset (see [SizeReference][])"""
522BATCH_AXIS_ID = AxisId("batch")
525class BatchAxis(AxisBase):
526 implemented_type: ClassVar[Literal["batch"]] = "batch"
527 if TYPE_CHECKING:
528 type: Literal["batch"] = "batch"
529 else:
530 type: Literal["batch"]
532 id: Annotated[AxisId, Predicate(_is_batch)] = BATCH_AXIS_ID
533 size: Optional[Literal[1]] = None
534 """The batch size may be fixed to 1,
535 otherwise (the default) it may be chosen arbitrarily depending on available memory"""
537 @property
538 def scale(self):
539 return 1.0
541 @property
542 def concatenable(self):
543 return True
545 @property
546 def unit(self):
547 return None
550class ChannelAxis(AxisBase):
551 implemented_type: ClassVar[Literal["channel"]] = "channel"
552 if TYPE_CHECKING:
553 type: Literal["channel"] = "channel"
554 else:
555 type: Literal["channel"]
557 id: NonBatchAxisId = AxisId("channel")
559 channel_names: NotEmpty[List[str]]
561 @property
562 def size(self) -> int:
563 return len(self.channel_names)
565 @property
566 def concatenable(self):
567 return False
569 @property
570 def scale(self) -> float:
571 return 1.0
573 @property
574 def unit(self):
575 return None
578class _WithInputAxisSize(Node):
579 size: Annotated[
580 Union[Annotated[int, Gt(0)], ParameterizedSize, SizeReference],
581 Field(
582 examples=[
583 10,
584 ParameterizedSize(min=32, step=16).model_dump(mode="json"),
585 SizeReference(
586 tensor_id=TensorId("t"), axis_id=AxisId("a"), offset=5
587 ).model_dump(mode="json"),
588 ]
589 ),
590 ]
591 """The size/length of this axis can be specified as
592 - fixed integer
593 - parameterized series of valid sizes ([ParameterizedSize][])
594 - reference to another axis with an optional offset ([SizeReference][])
595 """
598class IndexAxisBase(AxisBase):
599 implemented_type: ClassVar[Literal["index"]] = "index"
600 if TYPE_CHECKING:
601 type: Literal["index"] = "index"
602 else:
603 type: Literal["index"]
605 id: NonBatchAxisId = AxisId("index")
607 @property
608 def scale(self) -> float:
609 return 1.0
611 @property
612 def unit(self):
613 return None
616class IndexInputAxis(IndexAxisBase, _WithInputAxisSize):
617 concatenable: bool = False
618 """If a model has a `concatenable` input axis, it can be processed blockwise,
619 splitting a longer sample axis into blocks matching its input tensor description.
620 Output axes are concatenable if they have a [SizeReference][] to a concatenable
621 input axis.
622 """
625class IndexOutputAxis(IndexAxisBase):
626 size: Annotated[
627 Union[Annotated[int, Gt(0)], SizeReference, DataDependentSize],
628 Field(
629 examples=[
630 10,
631 SizeReference(
632 tensor_id=TensorId("t"), axis_id=AxisId("a"), offset=5
633 ).model_dump(mode="json"),
634 ]
635 ),
636 ]
637 """The size/length of this axis can be specified as
638 - fixed integer
639 - reference to another axis with an optional offset ([SizeReference][])
640 - data dependent size using [DataDependentSize][] (size is only known after model inference)
641 """
644class TimeAxisBase(AxisBase):
645 implemented_type: ClassVar[Literal["time"]] = "time"
646 if TYPE_CHECKING:
647 type: Literal["time"] = "time"
648 else:
649 type: Literal["time"]
651 id: NonBatchAxisId = AxisId("time")
652 unit: Optional[TimeUnit] = None
653 scale: Annotated[float, Gt(0)] = 1.0
656class TimeInputAxis(TimeAxisBase, _WithInputAxisSize):
657 concatenable: bool = False
658 """If a model has a `concatenable` input axis, it can be processed blockwise,
659 splitting a longer sample axis into blocks matching its input tensor description.
660 Output axes are concatenable if they have a [SizeReference][] to a concatenable
661 input axis.
662 """
665class SpaceAxisBase(AxisBase):
666 implemented_type: ClassVar[Literal["space"]] = "space"
667 if TYPE_CHECKING:
668 type: Literal["space"] = "space"
669 else:
670 type: Literal["space"]
672 id: Annotated[NonBatchAxisId, Field(examples=["x", "y", "z"])] = AxisId("x")
673 unit: Optional[SpaceUnit] = None
674 scale: Annotated[float, Gt(0)] = 1.0
677class SpaceInputAxis(SpaceAxisBase, _WithInputAxisSize):
678 concatenable: bool = False
679 """If a model has a `concatenable` input axis, it can be processed blockwise,
680 splitting a longer sample axis into blocks matching its input tensor description.
681 Output axes are concatenable if they have a [SizeReference][] to a concatenable
682 input axis.
683 """
686INPUT_AXIS_TYPES = (
687 BatchAxis,
688 ChannelAxis,
689 IndexInputAxis,
690 TimeInputAxis,
691 SpaceInputAxis,
692)
693"""intended for isinstance comparisons in py<3.10"""
695_InputAxisUnion = Union[
696 BatchAxis, ChannelAxis, IndexInputAxis, TimeInputAxis, SpaceInputAxis
697]
698InputAxis = Annotated[_InputAxisUnion, Discriminator("type")]
701class _WithOutputAxisSize(Node):
702 size: Annotated[
703 Union[Annotated[int, Gt(0)], SizeReference],
704 Field(
705 examples=[
706 10,
707 SizeReference(
708 tensor_id=TensorId("t"), axis_id=AxisId("a"), offset=5
709 ).model_dump(mode="json"),
710 ]
711 ),
712 ]
713 """The size/length of this axis can be specified as
714 - fixed integer
715 - reference to another axis with an optional offset (see [SizeReference][])
716 """
719class TimeOutputAxis(TimeAxisBase, _WithOutputAxisSize):
720 pass
723class TimeOutputAxisWithHalo(TimeAxisBase, WithHalo):
724 pass
727def _get_halo_axis_discriminator_value(v: Any) -> Literal["with_halo", "wo_halo"]:
728 if isinstance(v, dict):
729 return "with_halo" if "halo" in v else "wo_halo"
730 else:
731 return "with_halo" if hasattr(v, "halo") else "wo_halo"
734_TimeOutputAxisUnion = Annotated[
735 Union[
736 Annotated[TimeOutputAxis, Tag("wo_halo")],
737 Annotated[TimeOutputAxisWithHalo, Tag("with_halo")],
738 ],
739 Discriminator(_get_halo_axis_discriminator_value),
740]
743class SpaceOutputAxis(SpaceAxisBase, _WithOutputAxisSize):
744 pass
747class SpaceOutputAxisWithHalo(SpaceAxisBase, WithHalo):
748 pass
751_SpaceOutputAxisUnion = Annotated[
752 Union[
753 Annotated[SpaceOutputAxis, Tag("wo_halo")],
754 Annotated[SpaceOutputAxisWithHalo, Tag("with_halo")],
755 ],
756 Discriminator(_get_halo_axis_discriminator_value),
757]
760_OutputAxisUnion = Union[
761 BatchAxis, ChannelAxis, IndexOutputAxis, _TimeOutputAxisUnion, _SpaceOutputAxisUnion
762]
763OutputAxis = Annotated[_OutputAxisUnion, Discriminator("type")]
765OUTPUT_AXIS_TYPES = (
766 BatchAxis,
767 ChannelAxis,
768 IndexOutputAxis,
769 TimeOutputAxis,
770 TimeOutputAxisWithHalo,
771 SpaceOutputAxis,
772 SpaceOutputAxisWithHalo,
773)
774"""intended for isinstance comparisons in py<3.10"""
777AnyAxis = Union[InputAxis, OutputAxis]
779ANY_AXIS_TYPES = INPUT_AXIS_TYPES + OUTPUT_AXIS_TYPES
780"""intended for isinstance comparisons in py<3.10"""
782TVs = Union[
783 NotEmpty[List[int]],
784 NotEmpty[List[float]],
785 NotEmpty[List[bool]],
786 NotEmpty[List[str]],
787]
790NominalOrOrdinalDType = Literal[
791 "float32",
792 "float64",
793 "uint8",
794 "int8",
795 "uint16",
796 "int16",
797 "uint32",
798 "int32",
799 "uint64",
800 "int64",
801 "bool",
802]
805class NominalOrOrdinalDataDescr(Node):
806 values: TVs
807 """A fixed set of nominal or an ascending sequence of ordinal values.
808 In this case `data.type` is required to be an unsigend integer type, e.g. 'uint8'.
809 String `values` are interpreted as labels for tensor values 0, ..., N.
810 Note: as YAML 1.2 does not natively support a "set" datatype,
811 nominal values should be given as a sequence (aka list/array) as well.
812 """
814 type: Annotated[
815 NominalOrOrdinalDType,
816 Field(
817 examples=[
818 "float32",
819 "uint8",
820 "uint16",
821 "int64",
822 "bool",
823 ],
824 ),
825 ] = "uint8"
827 @model_validator(mode="after")
828 def _validate_values_match_type(
829 self,
830 ) -> Self:
831 incompatible: List[Any] = []
832 for v in self.values:
833 if self.type == "bool":
834 if not isinstance(v, bool):
835 incompatible.append(v)
836 elif self.type in DTYPE_LIMITS:
837 if (
838 isinstance(v, (int, float))
839 and (
840 v < DTYPE_LIMITS[self.type].min
841 or v > DTYPE_LIMITS[self.type].max
842 )
843 or (isinstance(v, str) and "uint" not in self.type)
844 or (isinstance(v, float) and "int" in self.type)
845 ):
846 incompatible.append(v)
847 else:
848 incompatible.append(v)
850 if len(incompatible) == 5:
851 incompatible.append("...")
852 break
854 if incompatible:
855 raise ValueError(
856 f"data type '{self.type}' incompatible with values {incompatible}"
857 )
859 return self
861 unit: Optional[Union[Literal["arbitrary unit"], SiUnit]] = None
863 @property
864 def range(self):
865 if isinstance(self.values[0], str):
866 return 0, len(self.values) - 1
867 else:
868 return min(self.values), max(self.values)
871IntervalOrRatioDType = Literal[
872 "float32",
873 "float64",
874 "uint8",
875 "int8",
876 "uint16",
877 "int16",
878 "uint32",
879 "int32",
880 "uint64",
881 "int64",
882]
885class IntervalOrRatioDataDescr(Node):
886 type: Annotated[ # TODO: rename to dtype
887 IntervalOrRatioDType,
888 Field(
889 examples=["float32", "float64", "uint8", "uint16"],
890 ),
891 ] = "float32"
892 range: Tuple[Optional[float], Optional[float]] = (
893 None,
894 None,
895 )
896 """Tuple `(minimum, maximum)` specifying the allowed range of the data in this tensor.
897 `None` corresponds to min/max of what can be expressed by **type**."""
898 unit: Union[Literal["arbitrary unit"], SiUnit] = "arbitrary unit"
899 scale: float = 1.0
900 """Scale for data on an interval (or ratio) scale."""
901 offset: Optional[float] = None
902 """Offset for data on a ratio scale."""
904 @model_validator(mode="before")
905 def _replace_inf(cls, data: Any):
906 if is_dict(data):
907 if "range" in data and is_sequence(data["range"]):
908 forbidden = (
909 "inf",
910 "-inf",
911 ".inf",
912 "-.inf",
913 float("inf"),
914 float("-inf"),
915 )
916 if any(v in forbidden for v in data["range"]):
917 issue_warning("replaced 'inf' value", value=data["range"])
919 data["range"] = tuple(
920 (None if v in forbidden else v) for v in data["range"]
921 )
923 return data
926TensorDataDescr = Union[NominalOrOrdinalDataDescr, IntervalOrRatioDataDescr]
929class BinarizeKwargs(KwargsNode):
930 """key word arguments for [BinarizeDescr][]"""
932 threshold: float
933 """The fixed threshold"""
936class BinarizeAlongAxisKwargs(KwargsNode):
937 """key word arguments for [BinarizeDescr][]"""
939 threshold: NotEmpty[List[float]]
940 """The fixed threshold values along `axis`"""
942 axis: Annotated[NonBatchAxisId, Field(examples=["channel"])]
943 """The `threshold` axis"""
946class BinarizeDescr(NodeWithExplicitlySetFields):
947 """Binarize the tensor with a fixed threshold.
949 Values above [BinarizeKwargs.threshold][]/[BinarizeAlongAxisKwargs.threshold][]
950 will be set to one, values below the threshold to zero.
952 Examples:
953 - in YAML
954 ```yaml
955 postprocessing:
956 - id: binarize
957 kwargs:
958 axis: 'channel'
959 threshold: [0.25, 0.5, 0.75]
960 ```
961 - in Python:
963 >>> postprocessing = [BinarizeDescr(
964 ... kwargs=BinarizeAlongAxisKwargs(
965 ... axis=AxisId('channel'),
966 ... threshold=[0.25, 0.5, 0.75],
967 ... )
968 ... )]
969 """
971 implemented_id: ClassVar[Literal["binarize"]] = "binarize"
972 if TYPE_CHECKING:
973 id: Literal["binarize"] = "binarize"
974 else:
975 id: Literal["binarize"]
976 kwargs: Union[BinarizeKwargs, BinarizeAlongAxisKwargs]
979class ClipKwargs(KwargsNode):
980 """key word arguments for [ClipDescr][]"""
982 min: Optional[float] = None
983 """Minimum value for clipping.
985 Exclusive with [min_percentile][]
986 """
987 min_percentile: Optional[Annotated[float, Interval(ge=0, lt=100)]] = None
988 """Minimum percentile for clipping.
990 Exclusive with [min][].
992 In range [0, 100).
993 """
995 max: Optional[float] = None
996 """Maximum value for clipping.
998 Exclusive with `max_percentile`.
999 """
1000 max_percentile: Optional[Annotated[float, Interval(gt=1, le=100)]] = None
1001 """Maximum percentile for clipping.
1003 Exclusive with `max`.
1005 In range (1, 100].
1006 """
1008 axes: Annotated[
1009 Optional[Sequence[AxisId]], Field(examples=[("batch", "x", "y")])
1010 ] = None
1011 """The subset of axes to determine percentiles jointly,
1013 i.e. axes to reduce to compute min/max from `min_percentile`/`max_percentile`.
1014 For example to clip 'batch', 'x' and 'y' jointly in a tensor ('batch', 'channel', 'y', 'x')
1015 resulting in a tensor of equal shape with clipped values per channel, specify `axes=('batch', 'x', 'y')`.
1016 To clip samples independently, leave out the 'batch' axis.
1018 Only valid if `min_percentile` and/or `max_percentile` are set.
1020 Default: Compute percentiles over all axes jointly."""
1022 @model_validator(mode="after")
1023 def _validate(self) -> Self:
1024 if (self.min is not None) and (self.min_percentile is not None):
1025 raise ValueError(
1026 "Only one of `min` and `min_percentile` may be set, not both."
1027 )
1028 if (self.max is not None) and (self.max_percentile is not None):
1029 raise ValueError(
1030 "Only one of `max` and `max_percentile` may be set, not both."
1031 )
1032 if (
1033 self.min is None
1034 and self.min_percentile is None
1035 and self.max is None
1036 and self.max_percentile is None
1037 ):
1038 raise ValueError(
1039 "At least one of `min`, `min_percentile`, `max`, or `max_percentile` must be set."
1040 )
1042 if (
1043 self.axes is not None
1044 and self.min_percentile is None
1045 and self.max_percentile is None
1046 ):
1047 raise ValueError(
1048 "If `axes` is set, at least one of `min_percentile` or `max_percentile` must be set."
1049 )
1051 return self
1054class ClipDescr(NodeWithExplicitlySetFields):
1055 """Set tensor values below min to min and above max to max.
1057 See `ScaleRangeDescr` for examples.
1058 """
1060 implemented_id: ClassVar[Literal["clip"]] = "clip"
1061 if TYPE_CHECKING:
1062 id: Literal["clip"] = "clip"
1063 else:
1064 id: Literal["clip"]
1066 kwargs: ClipKwargs
1069class EnsureDtypeKwargs(KwargsNode):
1070 """key word arguments for [EnsureDtypeDescr][]"""
1072 dtype: Literal[
1073 "float32",
1074 "float64",
1075 "uint8",
1076 "int8",
1077 "uint16",
1078 "int16",
1079 "uint32",
1080 "int32",
1081 "uint64",
1082 "int64",
1083 "bool",
1084 ]
1087class EnsureDtypeDescr(NodeWithExplicitlySetFields):
1088 """Cast the tensor data type to `EnsureDtypeKwargs.dtype` (if not matching).
1090 This can for example be used to ensure the inner neural network model gets a
1091 different input tensor data type than the fully described bioimage.io model does.
1093 Examples:
1094 The described bioimage.io model (incl. preprocessing) accepts any
1095 float32-compatible tensor, normalizes it with percentiles and clipping and then
1096 casts it to uint8, which is what the neural network in this example expects.
1097 - in YAML
1098 ```yaml
1099 inputs:
1100 - data:
1101 type: float32 # described bioimage.io model is compatible with any float32 input tensor
1102 preprocessing:
1103 - id: scale_range
1104 kwargs:
1105 axes: ['y', 'x']
1106 max_percentile: 99.8
1107 min_percentile: 5.0
1108 - id: clip
1109 kwargs:
1110 min: 0.0
1111 max: 1.0
1112 - id: ensure_dtype # the neural network of the model requires uint8
1113 kwargs:
1114 dtype: uint8
1115 ```
1116 - in Python:
1117 >>> preprocessing = [
1118 ... ScaleRangeDescr(
1119 ... kwargs=ScaleRangeKwargs(
1120 ... axes= (AxisId('y'), AxisId('x')),
1121 ... max_percentile= 99.8,
1122 ... min_percentile= 5.0,
1123 ... )
1124 ... ),
1125 ... ClipDescr(kwargs=ClipKwargs(min=0.0, max=1.0)),
1126 ... EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype="uint8")),
1127 ... ]
1128 """
1130 implemented_id: ClassVar[Literal["ensure_dtype"]] = "ensure_dtype"
1131 if TYPE_CHECKING:
1132 id: Literal["ensure_dtype"] = "ensure_dtype"
1133 else:
1134 id: Literal["ensure_dtype"]
1136 kwargs: EnsureDtypeKwargs
1139class ScaleLinearKwargs(KwargsNode):
1140 """Key word arguments for [ScaleLinearDescr][]"""
1142 gain: float = 1.0
1143 """multiplicative factor"""
1145 offset: float = 0.0
1146 """additive term"""
1148 @model_validator(mode="after")
1149 def _validate(self) -> Self:
1150 if self.gain == 1.0 and self.offset == 0.0:
1151 raise ValueError(
1152 "Redundant linear scaling not allowd. Set `gain` != 1.0 and/or `offset`"
1153 + " != 0.0."
1154 )
1156 return self
1159class ScaleLinearAlongAxisKwargs(KwargsNode):
1160 """Key word arguments for [ScaleLinearDescr][]"""
1162 axis: Annotated[NonBatchAxisId, Field(examples=["channel"])]
1163 """The axis of gain and offset values."""
1165 gain: Union[float, NotEmpty[List[float]]] = 1.0
1166 """multiplicative factor"""
1168 offset: Union[float, NotEmpty[List[float]]] = 0.0
1169 """additive term"""
1171 @model_validator(mode="after")
1172 def _validate(self) -> Self:
1173 if isinstance(self.gain, list):
1174 if isinstance(self.offset, list):
1175 if len(self.gain) != len(self.offset):
1176 raise ValueError(
1177 f"Size of `gain` ({len(self.gain)}) and `offset` ({len(self.offset)}) must match."
1178 )
1179 else:
1180 self.offset = [float(self.offset)] * len(self.gain)
1181 elif isinstance(self.offset, list):
1182 self.gain = [float(self.gain)] * len(self.offset)
1183 else:
1184 raise ValueError(
1185 "Do not specify an `axis` for scalar gain and offset values."
1186 )
1188 if all(g == 1.0 for g in self.gain) and all(off == 0.0 for off in self.offset):
1189 raise ValueError(
1190 "Redundant linear scaling not allowd. Set `gain` != 1.0 and/or `offset`"
1191 + " != 0.0."
1192 )
1194 return self
1197class ScaleLinearDescr(NodeWithExplicitlySetFields):
1198 """Fixed linear scaling.
1200 Examples:
1201 1. Scale with scalar gain and offset
1202 - in YAML
1203 ```yaml
1204 preprocessing:
1205 - id: scale_linear
1206 kwargs:
1207 gain: 2.0
1208 offset: 3.0
1209 ```
1210 - in Python:
1212 >>> preprocessing = [
1213 ... ScaleLinearDescr(kwargs=ScaleLinearKwargs(gain= 2.0, offset=3.0))
1214 ... ]
1216 2. Independent scaling along an axis
1217 - in YAML
1218 ```yaml
1219 preprocessing:
1220 - id: scale_linear
1221 kwargs:
1222 axis: 'channel'
1223 gain: [1.0, 2.0, 3.0]
1224 ```
1225 - in Python:
1227 >>> preprocessing = [
1228 ... ScaleLinearDescr(
1229 ... kwargs=ScaleLinearAlongAxisKwargs(
1230 ... axis=AxisId("channel"),
1231 ... gain=[1.0, 2.0, 3.0],
1232 ... )
1233 ... )
1234 ... ]
1236 """
1238 implemented_id: ClassVar[Literal["scale_linear"]] = "scale_linear"
1239 if TYPE_CHECKING:
1240 id: Literal["scale_linear"] = "scale_linear"
1241 else:
1242 id: Literal["scale_linear"]
1243 kwargs: Union[ScaleLinearKwargs, ScaleLinearAlongAxisKwargs]
1246class SigmoidDescr(NodeWithExplicitlySetFields):
1247 """The logistic sigmoid function, a.k.a. expit function.
1249 Examples:
1250 - in YAML
1251 ```yaml
1252 postprocessing:
1253 - id: sigmoid
1254 ```
1255 - in Python:
1257 >>> postprocessing = [SigmoidDescr()]
1258 """
1260 implemented_id: ClassVar[Literal["sigmoid"]] = "sigmoid"
1261 if TYPE_CHECKING:
1262 id: Literal["sigmoid"] = "sigmoid"
1263 else:
1264 id: Literal["sigmoid"]
1266 @property
1267 def kwargs(self) -> KwargsNode:
1268 """empty kwargs"""
1269 return KwargsNode()
1272class SoftmaxKwargs(KwargsNode):
1273 """key word arguments for [SoftmaxDescr][]"""
1275 axis: Annotated[NonBatchAxisId, Field(examples=["channel"])] = AxisId("channel")
1276 """The axis to apply the softmax function along.
1277 Note:
1278 Defaults to 'channel' axis
1279 (which may not exist, in which case
1280 a different axis id has to be specified).
1281 """
1284class SoftmaxDescr(NodeWithExplicitlySetFields):
1285 """The softmax function.
1287 Examples:
1288 - in YAML
1289 ```yaml
1290 postprocessing:
1291 - id: softmax
1292 kwargs:
1293 axis: channel
1294 ```
1295 - in Python:
1297 >>> postprocessing = [SoftmaxDescr(kwargs=SoftmaxKwargs(axis=AxisId("channel")))]
1298 """
1300 implemented_id: ClassVar[Literal["softmax"]] = "softmax"
1301 if TYPE_CHECKING:
1302 id: Literal["softmax"] = "softmax"
1303 else:
1304 id: Literal["softmax"]
1306 kwargs: SoftmaxKwargs = Field(default_factory=SoftmaxKwargs.model_construct)
1309class _StardistPostprocessingKwargsBase(KwargsNode):
1310 """key word arguments for [StardistPostprocessingDescr][]"""
1312 prob_threshold: float
1313 """The probability threshold for object candidate selection."""
1315 nms_threshold: float
1316 """The IoU threshold for non-maximum suppression."""
1318 n_rays: int
1319 """Number of radial lines (rays) cast from the center of an object to its boundary."""
1322class StardistPostprocessingKwargs2D(_StardistPostprocessingKwargsBase):
1323 grid: Tuple[int, int]
1324 """Grid size of network predictions."""
1326 b: Union[int, Tuple[Tuple[int, int], Tuple[int, int]]]
1327 """Border region in which object probability is set to zero."""
1330class StardistPostprocessingKwargs3D(_StardistPostprocessingKwargsBase):
1331 grid: Tuple[int, int, int]
1332 """Grid size of network predictions."""
1334 b: Union[int, Tuple[Tuple[int, int], Tuple[int, int], Tuple[int, int]]]
1335 """Border region in which object probability is set to zero."""
1337 anisotropy: Tuple[float, float, float]
1338 """Anisotropy factors for 3D star-convex polyhedra, i.e. the physical pixel size along each spatial axis."""
1340 overlap_label: Optional[int] = None
1341 """Optional label to apply to any area of overlapping predicted objects."""
1344class StardistPostprocessingDescr(NodeWithExplicitlySetFields):
1345 """Stardist postprocessing including non-maximum suppression and converting polygon representations to instance labels
1347 as described in:
1348 - Uwe Schmidt, Martin Weigert, Coleman Broaddus, and Gene Myers.
1349 [*Cell Detection with Star-convex Polygons*](https://arxiv.org/abs/1806.03535).
1350 International Conference on Medical Image Computing and Computer-Assisted Intervention (MICCAI), Granada, Spain, September 2018.
1351 - Martin Weigert, Uwe Schmidt, Robert Haase, Ko Sugawara, and Gene Myers.
1352 [*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).
1353 The IEEE Winter Conference on Applications of Computer Vision (WACV), Snowmass Village, Colorado, March 2020.
1355 Note: Only available if the `stardist` package is installed.
1356 """
1358 implemented_id: ClassVar[Literal["stardist_postprocessing"]] = (
1359 "stardist_postprocessing"
1360 )
1361 if TYPE_CHECKING:
1362 id: Literal["stardist_postprocessing"] = "stardist_postprocessing"
1363 else:
1364 id: Literal["stardist_postprocessing"]
1366 kwargs: Union[StardistPostprocessingKwargs2D, StardistPostprocessingKwargs3D]
1369class CellposeFlowDynamicsKwargs(KwargsNode):
1370 """key word arguments for [CellposeFlowDynamicsDescr][]"""
1372 cellprob_threshold: float
1373 flow_threshold: float
1374 do_3D: bool
1375 min_size: int = 15
1376 """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."""
1377 output_dtype: Literal["uint16", "uint32"] = "uint16"
1380class CellposeFlowDynamicsDescr(NodeWithExplicitlySetFields):
1381 """Cellpose flow dynamics postprocessing as described in:
1382 - Carsen Stringer and Marius Pachitariu. [*Cellpose: a generalist algorithm for cellular segmentation*](https://www.nature.com/articles/s41592-020-01018-x). Nature Methods, 2021.
1384 Note: Only available if the `cellpose` package is installed.
1385 """
1387 implemented_id: ClassVar[Literal["cellpose_flow_dynamics"]] = (
1388 "cellpose_flow_dynamics"
1389 )
1390 if TYPE_CHECKING:
1391 id: Literal["cellpose_flow_dynamics"] = "cellpose_flow_dynamics"
1392 else:
1393 id: Literal["cellpose_flow_dynamics"]
1395 kwargs: CellposeFlowDynamicsKwargs
1398class CustomProcessingDescr(NodeWithExplicitlySetFields, FileDescr):
1399 """Custom (post)processing op — source file shipped inline with the model.
1401 Supports (post)processing that cannot be expressed by the built-in named
1402 operations (watershed, connected components, etc.)
1403 using a simple Python callable interface.
1405 The op is implemented in a ``.py`` file packaged alongside the model weights.
1406 Two styles are supported:
1408 *Callable class* — kwargs go to ``__init__``, tensors arrive in ``__call__``:
1410 .. code-block:: python
1412 # my_postprocess.py
1413 import numpy as np
1415 class my_postprocess:
1416 def __init__(self, threshold: float = 0.5) -> None:
1417 self.threshold = threshold
1418 def __call__(self, *arrays: np.ndarray) -> np.ndarray:
1419 # arrays = model output tensors in rdf.yaml declaration order
1420 return (arrays[0] > self.threshold).astype(np.uint8)
1422 *Factory function* — alternative closure style, identical runtime behaviour:
1424 .. code-block:: python
1426 # my_postprocess.py
1427 import numpy as np
1429 def my_postprocess(threshold: float = 0.5):
1430 def run(*arrays: np.ndarray) -> np.ndarray:
1431 return (arrays[0] > threshold).astype(np.uint8)
1432 return run
1434 Reference it in ``rdf.yaml`` with the source file included in the package:
1436 .. code-block:: yaml
1438 postprocessing:
1439 - id: custom
1440 callable: my_postprocess # class or function name in source
1441 source: my_postprocess.py # packaged alongside weights
1442 sha256: <hash> # sha256 of the source file
1443 kwargs: # forwarded to __init__ / factory
1444 threshold: 0.5
1446 **Security:** source files are SHA-256 verified before execution.
1447 Execution requires explicit opt-in in bioimageio.core and curator
1448 review before Zoo publication.
1449 """
1451 implemented_id: ClassVar[Literal["custom"]] = "custom"
1452 if TYPE_CHECKING:
1453 id: Literal["custom"] = "custom"
1454 else:
1455 id: Literal["custom"]
1457 callable: Annotated[
1458 str,
1459 Field(examples=["my_postprocess_factory", "MyPostprocessClass"]),
1460 ]
1461 """Name of the callable class or factory function defined in ``source``.
1463 At runtime: ``op = callable(**kwargs)``, then ``result = op(*output_tensors)``
1464 per image. Both a class with ``__call__`` and a factory function returning
1465 a callable satisfy this protocol."""
1467 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
1468 """Python source file (included when packaging the model)."""
1470 kwargs: Dict[str, YamlValue] = Field(
1471 default_factory=cast(Callable[[], Dict[str, YamlValue]], dict)
1472 )
1473 """Keyword arguments forwarded to the callable (``__init__`` or factory)."""
1475 @model_serializer(mode="wrap", when_used="unless-none")
1476 def _serialize(
1477 self, nxt: SerializerFunctionWrapHandler, info: SerializationInfo
1478 ) -> Dict[str, YamlValue]:
1479 return package_file_descr_serializer(self, nxt, info)
1482class FixedZeroMeanUnitVarianceKwargs(KwargsNode):
1483 """key word arguments for [FixedZeroMeanUnitVarianceDescr][]"""
1485 mean: float
1486 """The mean value to normalize with."""
1488 std: Annotated[float, Ge(1e-6)]
1489 """The standard deviation value to normalize with."""
1492class FixedZeroMeanUnitVarianceAlongAxisKwargs(KwargsNode):
1493 """key word arguments for [FixedZeroMeanUnitVarianceDescr][]"""
1495 mean: NotEmpty[List[float]]
1496 """The mean value(s) to normalize with."""
1498 std: NotEmpty[List[Annotated[float, Ge(1e-6)]]]
1499 """The standard deviation value(s) to normalize with.
1500 Size must match `mean` values."""
1502 axis: Annotated[NonBatchAxisId, Field(examples=["channel", "index"])]
1503 """The axis of the mean/std values to normalize each entry along that dimension
1504 separately."""
1506 @model_validator(mode="after")
1507 def _mean_and_std_match(self) -> Self:
1508 if len(self.mean) != len(self.std):
1509 raise ValueError(
1510 f"Size of `mean` ({len(self.mean)}) and `std` ({len(self.std)})"
1511 + " must match."
1512 )
1514 return self
1517class FixedZeroMeanUnitVarianceDescr(NodeWithExplicitlySetFields):
1518 """Subtract a given mean and divide by the standard deviation.
1520 Normalize with fixed, precomputed values for
1521 `FixedZeroMeanUnitVarianceKwargs.mean` and `FixedZeroMeanUnitVarianceKwargs.std`
1522 Use `FixedZeroMeanUnitVarianceAlongAxisKwargs` for independent scaling along given
1523 axes.
1525 Examples:
1526 1. scalar value for whole tensor
1527 - in YAML
1528 ```yaml
1529 preprocessing:
1530 - id: fixed_zero_mean_unit_variance
1531 kwargs:
1532 mean: 103.5
1533 std: 13.7
1534 ```
1535 - in Python
1536 >>> preprocessing = [FixedZeroMeanUnitVarianceDescr(
1537 ... kwargs=FixedZeroMeanUnitVarianceKwargs(mean=103.5, std=13.7)
1538 ... )]
1540 2. independently along an axis
1541 - in YAML
1542 ```yaml
1543 preprocessing:
1544 - id: fixed_zero_mean_unit_variance
1545 kwargs:
1546 axis: channel
1547 mean: [101.5, 102.5, 103.5]
1548 std: [11.7, 12.7, 13.7]
1549 ```
1550 - in Python
1551 >>> preprocessing = [FixedZeroMeanUnitVarianceDescr(
1552 ... kwargs=FixedZeroMeanUnitVarianceAlongAxisKwargs(
1553 ... axis=AxisId("channel"),
1554 ... mean=[101.5, 102.5, 103.5],
1555 ... std=[11.7, 12.7, 13.7],
1556 ... )
1557 ... )]
1558 """
1560 implemented_id: ClassVar[Literal["fixed_zero_mean_unit_variance"]] = (
1561 "fixed_zero_mean_unit_variance"
1562 )
1563 if TYPE_CHECKING:
1564 id: Literal["fixed_zero_mean_unit_variance"] = "fixed_zero_mean_unit_variance"
1565 else:
1566 id: Literal["fixed_zero_mean_unit_variance"]
1568 kwargs: Union[
1569 FixedZeroMeanUnitVarianceKwargs, FixedZeroMeanUnitVarianceAlongAxisKwargs
1570 ]
1573class ZeroMeanUnitVarianceKwargs(KwargsNode):
1574 """key word arguments for [ZeroMeanUnitVarianceDescr][]"""
1576 axes: Annotated[
1577 Optional[Sequence[AxisId]], Field(examples=[("batch", "x", "y")])
1578 ] = None
1579 """The subset of axes to normalize jointly, i.e. axes to reduce to compute mean/std.
1580 For example to normalize 'batch', 'x' and 'y' jointly in a tensor ('batch', 'channel', 'y', 'x')
1581 resulting in a tensor of equal shape normalized per channel, specify `axes=('batch', 'x', 'y')`.
1582 To normalize each sample independently leave out the 'batch' axis.
1583 Default: Scale all axes jointly."""
1585 eps: Annotated[float, Interval(gt=0, le=0.1)] = 1e-6
1586 """epsilon for numeric stability: `out = (tensor - mean) / (std + eps)`."""
1589class ZeroMeanUnitVarianceDescr(NodeWithExplicitlySetFields):
1590 """Subtract mean and divide by variance.
1592 Examples:
1593 Subtract tensor mean and variance
1594 - in YAML
1595 ```yaml
1596 preprocessing:
1597 - id: zero_mean_unit_variance
1598 ```
1599 - in Python
1600 >>> preprocessing = [ZeroMeanUnitVarianceDescr()]
1601 """
1603 implemented_id: ClassVar[Literal["zero_mean_unit_variance"]] = (
1604 "zero_mean_unit_variance"
1605 )
1606 if TYPE_CHECKING:
1607 id: Literal["zero_mean_unit_variance"] = "zero_mean_unit_variance"
1608 else:
1609 id: Literal["zero_mean_unit_variance"]
1611 kwargs: ZeroMeanUnitVarianceKwargs = Field(
1612 default_factory=ZeroMeanUnitVarianceKwargs.model_construct
1613 )
1616class ScaleRangeKwargs(KwargsNode):
1617 """key word arguments for [ScaleRangeDescr][]
1619 For `min_percentile`=0.0 (the default) and `max_percentile`=100 (the default)
1620 this processing step normalizes data to the [0, 1] intervall.
1621 For other percentiles the normalized values will partially be outside the [0, 1]
1622 intervall. Use `ScaleRange` followed by `ClipDescr` if you want to limit the
1623 normalized values to a range.
1624 """
1626 axes: Annotated[
1627 Optional[Sequence[AxisId]], Field(examples=[("batch", "x", "y")])
1628 ] = None
1629 """The subset of axes to normalize jointly, i.e. axes to reduce to compute the min/max percentile value.
1630 For example to normalize 'batch', 'x' and 'y' jointly in a tensor ('batch', 'channel', 'y', 'x')
1631 resulting in a tensor of equal shape normalized per channel, specify `axes=('batch', 'x', 'y')`.
1632 To normalize samples independently, leave out the "batch" axis.
1633 Default: Scale all axes jointly."""
1635 min_percentile: Annotated[float, Interval(ge=0, lt=100)] = 0.0
1636 """The lower percentile used to determine the value to align with zero."""
1638 max_percentile: Annotated[float, Interval(gt=1, le=100)] = 100.0
1639 """The upper percentile used to determine the value to align with one.
1640 Has to be bigger than `min_percentile`.
1641 The range is 1 to 100 instead of 0 to 100 to avoid mistakenly
1642 accepting percentiles specified in the range 0.0 to 1.0."""
1644 eps: Annotated[float, Interval(gt=0, le=0.1)] = 1e-6
1645 """Epsilon for numeric stability.
1646 `out = (tensor - v_lower) / (v_upper - v_lower + eps)`;
1647 with `v_lower,v_upper` values at the respective percentiles."""
1649 reference_tensor: Optional[TensorId] = None
1650 """ID of the unprocessed input tensor to compute the percentiles from.
1651 Default: The tensor itself.
1652 """
1654 @field_validator("max_percentile", mode="after")
1655 @classmethod
1656 def min_smaller_max(cls, value: float, info: ValidationInfo) -> float:
1657 if (min_p := info.data["min_percentile"]) >= value:
1658 raise ValueError(f"min_percentile {min_p} >= max_percentile {value}")
1660 return value
1663class ScaleRangeDescr(NodeWithExplicitlySetFields):
1664 """Scale with percentiles.
1666 Examples:
1667 1. Scale linearly to map 5th percentile to 0 and 99.8th percentile to 1.0
1668 - in YAML
1669 ```yaml
1670 preprocessing:
1671 - id: scale_range
1672 kwargs:
1673 axes: ['y', 'x']
1674 max_percentile: 99.8
1675 min_percentile: 5.0
1676 ```
1677 - in Python
1679 >>> preprocessing = [
1680 ... ScaleRangeDescr(
1681 ... kwargs=ScaleRangeKwargs(
1682 ... axes= (AxisId('y'), AxisId('x')),
1683 ... max_percentile= 99.8,
1684 ... min_percentile= 5.0,
1685 ... )
1686 ... )
1687 ... ]
1689 2. Combine the above scaling with additional clipping to clip values outside the range given by the percentiles.
1690 - in YAML
1691 ```yaml
1692 preprocessing:
1693 - id: scale_range
1694 kwargs:
1695 axes: ['y', 'x']
1696 max_percentile: 99.8
1697 min_percentile: 5.0
1698 - id: clip
1699 kwargs:
1700 min: 0.0
1701 max: 1.0
1702 ```
1703 - in Python
1705 >>> preprocessing = [
1706 ... ScaleRangeDescr(
1707 ... kwargs=ScaleRangeKwargs(
1708 ... axes= (AxisId('y'), AxisId('x')),
1709 ... max_percentile= 99.8,
1710 ... min_percentile= 5.0,
1711 ... )
1712 ... ),
1713 ... ClipDescr(
1714 ... kwargs=ClipKwargs(
1715 ... min=0.0,
1716 ... max=1.0,
1717 ... )
1718 ... ),
1719 ... ]
1721 """
1723 implemented_id: ClassVar[Literal["scale_range"]] = "scale_range"
1724 if TYPE_CHECKING:
1725 id: Literal["scale_range"] = "scale_range"
1726 else:
1727 id: Literal["scale_range"]
1728 kwargs: ScaleRangeKwargs = Field(default_factory=ScaleRangeKwargs.model_construct)
1731class ScaleMeanVarianceKwargs(KwargsNode):
1732 """key word arguments for [ScaleMeanVarianceKwargs][]"""
1734 reference_tensor: TensorId
1735 """ID of unprocessed input tensor to match."""
1737 axes: Annotated[
1738 Optional[Sequence[AxisId]], Field(examples=[("batch", "x", "y")])
1739 ] = None
1740 """The subset of axes to normalize jointly, i.e. axes to reduce to compute mean/std.
1741 For example to normalize 'batch', 'x' and 'y' jointly in a tensor ('batch', 'channel', 'y', 'x')
1742 resulting in a tensor of equal shape normalized per channel, specify `axes=('batch', 'x', 'y')`.
1743 To normalize samples independently, leave out the 'batch' axis.
1744 Default: Scale all axes jointly."""
1746 eps: Annotated[float, Interval(gt=0, le=0.1)] = 1e-6
1747 """Epsilon for numeric stability:
1748 `out = (tensor - mean) / (std + eps) * (ref_std + eps) + ref_mean.`"""
1751class ScaleMeanVarianceDescr(NodeWithExplicitlySetFields):
1752 """Scale a tensor's data distribution to match another tensor's mean/std.
1753 `out = (tensor - mean) / (std + eps) * (ref_std + eps) + ref_mean.`
1754 """
1756 implemented_id: ClassVar[Literal["scale_mean_variance"]] = "scale_mean_variance"
1757 if TYPE_CHECKING:
1758 id: Literal["scale_mean_variance"] = "scale_mean_variance"
1759 else:
1760 id: Literal["scale_mean_variance"]
1761 kwargs: ScaleMeanVarianceKwargs
1764PreprocessingDescr = Annotated[
1765 Union[
1766 BinarizeDescr,
1767 ClipDescr,
1768 EnsureDtypeDescr,
1769 FixedZeroMeanUnitVarianceDescr,
1770 ScaleLinearDescr,
1771 ScaleRangeDescr,
1772 SigmoidDescr,
1773 SoftmaxDescr,
1774 ZeroMeanUnitVarianceDescr,
1775 ],
1776 Discriminator("id"),
1777]
1778PostprocessingDescr = Annotated[
1779 Union[
1780 BinarizeDescr,
1781 CellposeFlowDynamicsDescr,
1782 ClipDescr,
1783 CustomProcessingDescr,
1784 EnsureDtypeDescr,
1785 FixedZeroMeanUnitVarianceDescr,
1786 ScaleLinearDescr,
1787 ScaleMeanVarianceDescr,
1788 ScaleRangeDescr,
1789 SigmoidDescr,
1790 SoftmaxDescr,
1791 StardistPostprocessingDescr,
1792 ZeroMeanUnitVarianceDescr,
1793 ],
1794 Discriminator("id"),
1795]
1797IO_AxisT = TypeVar("IO_AxisT", InputAxis, OutputAxis)
1800class TensorDescrBase(Node, Generic[IO_AxisT]):
1801 id: TensorId
1802 """Tensor id. No duplicates are allowed."""
1804 description: Annotated[str, MaxLen(128)] = ""
1805 """free text description"""
1807 axes: NotEmpty[Sequence[IO_AxisT]]
1808 """tensor axes"""
1810 @property
1811 def shape(self):
1812 return tuple(a.size for a in self.axes)
1814 @field_validator("axes", mode="after", check_fields=False)
1815 @classmethod
1816 def _validate_axes(cls, axes: Sequence[AnyAxis]) -> Sequence[AnyAxis]:
1817 batch_axes = [a for a in axes if a.type == "batch"]
1818 if len(batch_axes) > 1:
1819 raise ValueError(
1820 f"Only one batch axis (per tensor) allowed, but got {batch_axes}"
1821 )
1823 seen_ids: Set[AxisId] = set()
1824 duplicate_axes_ids: Set[AxisId] = set()
1825 for a in axes:
1826 (duplicate_axes_ids if a.id in seen_ids else seen_ids).add(a.id)
1828 if duplicate_axes_ids:
1829 raise ValueError(f"Duplicate axis ids: {duplicate_axes_ids}")
1831 return axes
1833 test_tensor: FAIR[Optional[FileDescr_package]] = None
1834 """An example tensor to use for testing.
1835 Using the model with the test input tensors is expected to yield the test output tensors.
1836 Each test tensor has be a an ndarray in the
1837 [numpy.lib file format](https://numpy.org/doc/stable/reference/generated/numpy.lib.format.html#module-numpy.lib.format).
1838 The file extension must be '.npy'."""
1840 sample_tensor: FAIR[Optional[FileDescr_package]] = None
1841 """A sample tensor to illustrate a possible input/output for the model,
1842 The sample image primarily serves to inform a human user about an example use case
1843 and is typically stored as .hdf5, .png or .tiff.
1844 It has to be readable by the [imageio library](https://imageio.readthedocs.io/en/stable/formats/index.html#supported-formats)
1845 (numpy's `.npy` format is not supported).
1846 The image dimensionality has to match the number of axes specified in this tensor description.
1847 """
1849 @model_validator(mode="after")
1850 def _validate_sample_tensor(self) -> Self:
1851 if self.sample_tensor is None or not get_validation_context().perform_io_checks:
1852 return self
1854 reader = get_reader(self.sample_tensor.source, sha256=self.sample_tensor.sha256)
1855 tensor: NDArray[Any] = imread( # pyright: ignore[reportUnknownVariableType]
1856 reader.read(),
1857 extension=PurePosixPath(reader.original_file_name).suffix,
1858 )
1859 n_dims = len(tensor.squeeze().shape)
1860 n_dims_min = n_dims_max = len(self.axes)
1862 for a in self.axes:
1863 if isinstance(a, BatchAxis):
1864 n_dims_min -= 1
1865 elif isinstance(a.size, int):
1866 if a.size == 1:
1867 n_dims_min -= 1
1868 elif isinstance(a.size, (ParameterizedSize, DataDependentSize)):
1869 if a.size.min == 1:
1870 n_dims_min -= 1
1871 elif isinstance(a.size, SizeReference):
1872 if a.size.offset < 2:
1873 # size reference may result in singleton axis
1874 n_dims_min -= 1
1875 else:
1876 assert_never(a.size)
1878 n_dims_min = max(0, n_dims_min)
1879 if n_dims < n_dims_min or n_dims > n_dims_max:
1880 raise ValueError(
1881 f"Expected sample tensor to have {n_dims_min} to"
1882 + f" {n_dims_max} dimensions, but found {n_dims} (shape: {tensor.shape})."
1883 )
1885 return self
1887 data: Union[TensorDataDescr, NotEmpty[Sequence[TensorDataDescr]]] = (
1888 IntervalOrRatioDataDescr()
1889 )
1890 """Description of the tensor's data values, optionally per channel.
1891 If specified per channel, the data `type` needs to match across channels."""
1893 @property
1894 def dtype(
1895 self,
1896 ) -> Literal[
1897 "float32",
1898 "float64",
1899 "uint8",
1900 "int8",
1901 "uint16",
1902 "int16",
1903 "uint32",
1904 "int32",
1905 "uint64",
1906 "int64",
1907 "bool",
1908 ]:
1909 """dtype as specified under `data.type` or `data[i].type`"""
1910 if isinstance(self.data, collections.abc.Sequence):
1911 return self.data[0].type
1912 else:
1913 return self.data.type
1915 @field_validator("data", mode="after")
1916 @classmethod
1917 def _check_data_type_across_channels(
1918 cls, value: Union[TensorDataDescr, NotEmpty[Sequence[TensorDataDescr]]]
1919 ) -> Union[TensorDataDescr, NotEmpty[Sequence[TensorDataDescr]]]:
1920 if not isinstance(value, list):
1921 return value
1923 dtypes = {t.type for t in value}
1924 if len(dtypes) > 1:
1925 raise ValueError(
1926 "Tensor data descriptions per channel need to agree in their data"
1927 + f" `type`, but found {dtypes}."
1928 )
1930 return value
1932 @model_validator(mode="after")
1933 def _check_data_matches_channelaxis(self) -> Self:
1934 if not isinstance(self.data, (list, tuple)):
1935 return self
1937 for a in self.axes:
1938 if isinstance(a, ChannelAxis):
1939 size = a.size
1940 assert isinstance(size, int)
1941 break
1942 else:
1943 return self
1945 if len(self.data) != size:
1946 raise ValueError(
1947 f"Got tensor data descriptions for {len(self.data)} channels, but"
1948 + f" '{a.id}' axis has size {size}."
1949 )
1951 return self
1953 def get_axis_sizes_for_array(self, array: NDArray[Any]) -> Dict[AxisId, int]:
1954 if len(array.shape) != len(self.axes):
1955 raise ValueError(
1956 f"Dimension mismatch: array shape {array.shape} (#{len(array.shape)})"
1957 + f" incompatible with {len(self.axes)} axes."
1958 )
1959 return {a.id: array.shape[i] for i, a in enumerate(self.axes)}
1962class ConstantPadding(Node):
1963 mode: Literal["constant"] = "constant"
1964 value: Union[int, float] = 0
1967class EdgePadding(Node):
1968 mode: Literal["edge"] = "edge"
1971class ReflectPadding(Node):
1972 mode: Literal["reflect"] = "reflect"
1975class SymmetricPadding(Node):
1976 mode: Literal["symmetric"] = "symmetric"
1979Padding = Union[ConstantPadding, EdgePadding, ReflectPadding, SymmetricPadding]
1982class ModelId(ResourceId):
1983 pass
1986class InputTensorDescr(TensorDescrBase[InputAxis]):
1987 id: TensorId = TensorId("input")
1988 """Input tensor id.
1989 No duplicates are allowed across all inputs and outputs."""
1991 output_of: Optional[ModelId] = None
1992 """If this input tensor is the output of another model, specify the model id here.
1993 This model's input id must match the output id of the referenced model.
1994 """
1996 @model_validator(mode="after")
1997 def _validate_output_of(self) -> Self:
1998 if self.output_of is None:
1999 return self
2001 try:
2002 with get_validation_context().replace(perform_io_checks=False):
2003 opened_ref_model = open_bioimageio_yaml(self.output_of)
2004 format_version = opened_ref_model.content["format_version"]
2005 assert isinstance(format_version, str)
2006 if format_version.startswith("0.4"):
2007 ref_model = _ModelDescr04.model_validate(opened_ref_model.content)
2008 else:
2009 ref_model = ModelDescr.model_validate(opened_ref_model.content)
2010 except Exception as e:
2011 raise ValueError(
2012 f"Failed to load model '{self.output_of}' referenced under output_of: {e}"
2013 )
2015 try:
2016 ref_model_outputs = {
2017 t.id if isinstance(t, OutputTensorDescr) else TensorId(t.name)
2018 for t in ref_model.outputs
2019 }
2020 except Exception as e:
2021 raise ValueError(
2022 f"Failed to read output IDs of model '{self.output_of}' referenced under output_of: {e}"
2023 )
2025 if self.id not in ref_model_outputs:
2026 raise ValueError(
2027 f"Input tensor '{self.id}' is specified as output of model '{self.output_of}', "
2028 + f"but that model's outputs are {ref_model_outputs}."
2029 )
2030 return self
2032 optional: bool = False
2033 """indicates that this tensor may be `None`"""
2035 pad: Optional[Padding] = None
2036 """Explicitly specify how to pad this input tensor.
2038 Use `axes[i].pad` to specify padding width.
2040 Note:
2041 Non-blockwise sample prediction only applies padding for axes with a `pad` specification.
2042 """
2044 preprocessing: List[PreprocessingDescr] = Field(
2045 default_factory=cast(Callable[[], List[PreprocessingDescr]], list)
2046 )
2047 """Description of how this input should be preprocessed.
2049 notes:
2050 - If preprocessing does not start with an 'ensure_dtype' entry, it is added
2051 to ensure an input tensor's data type matches the input tensor's data description.
2052 - If preprocessing does not end with an 'ensure_dtype' or 'binarize' entry, an
2053 'ensure_dtype' step is added to ensure preprocessing steps are not unintentionally
2054 changing the data type.
2055 """
2057 @model_validator(mode="after")
2058 def _validate_preprocessing_kwargs(self) -> Self:
2059 axes_ids = [a.id for a in self.axes]
2060 for p in self.preprocessing:
2061 kwargs_axes: Optional[Sequence[Any]] = p.kwargs.get("axes")
2062 if kwargs_axes is None:
2063 continue
2065 if not isinstance(kwargs_axes, collections.abc.Sequence):
2066 raise ValueError(
2067 f"Expected `preprocessing.i.kwargs.axes` to be a sequence, but got {type(kwargs_axes)}"
2068 )
2070 if any(a not in axes_ids for a in kwargs_axes):
2071 raise ValueError(
2072 "`preprocessing.i.kwargs.axes` needs to be subset of axes ids"
2073 )
2075 if isinstance(self.data, (NominalOrOrdinalDataDescr, IntervalOrRatioDataDescr)):
2076 dtype = self.data.type
2077 else:
2078 dtype = self.data[0].type
2080 # ensure `preprocessing` begins with `EnsureDtypeDescr`
2081 if not self.preprocessing or not isinstance(
2082 self.preprocessing[0], EnsureDtypeDescr
2083 ):
2084 self.preprocessing.insert(
2085 0, EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype=dtype))
2086 )
2088 # ensure `preprocessing` ends with `EnsureDtypeDescr` or `BinarizeDescr`
2089 if not isinstance(self.preprocessing[-1], (EnsureDtypeDescr, BinarizeDescr)):
2090 self.preprocessing.append(
2091 EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype=dtype))
2092 )
2094 return self
2097def convert_axes(
2098 axes: str,
2099 *,
2100 shape: Union[
2101 Sequence[int], _ParameterizedInputShape_v0_4, _ImplicitOutputShape_v0_4
2102 ],
2103 tensor_type: Literal["input", "output"],
2104 halo: Optional[Sequence[int]],
2105 size_refs: Mapping[_TensorName_v0_4, Mapping[str, int]],
2106):
2107 ret: List[AnyAxis] = []
2108 for i, a in enumerate(axes):
2109 axis_type = _AXIS_TYPE_MAP.get(a, a)
2110 if axis_type == "batch":
2111 ret.append(BatchAxis())
2112 continue
2114 scale = 1.0
2115 if isinstance(shape, _ParameterizedInputShape_v0_4):
2116 if shape.step[i] == 0:
2117 size = shape.min[i]
2118 else:
2119 size = ParameterizedSize(min=shape.min[i], step=shape.step[i])
2120 elif isinstance(shape, _ImplicitOutputShape_v0_4):
2121 ref_t = str(shape.reference_tensor)
2122 if ref_t.count(".") == 1:
2123 t_id, orig_a_id = ref_t.split(".")
2124 else:
2125 t_id = ref_t
2126 orig_a_id = a
2128 a_id = _AXIS_ID_MAP.get(orig_a_id, a)
2129 if not (orig_scale := shape.scale[i]):
2130 # old way to insert a new axis dimension
2131 size = int(2 * shape.offset[i])
2132 else:
2133 scale = 1 / orig_scale
2134 if axis_type in ("channel", "index"):
2135 # these axes no longer have a scale
2136 offset_from_scale = orig_scale * size_refs.get(
2137 _TensorName_v0_4(t_id), {}
2138 ).get(orig_a_id, 0)
2139 else:
2140 offset_from_scale = 0
2141 size = SizeReference(
2142 tensor_id=TensorId(t_id),
2143 axis_id=AxisId(a_id),
2144 offset=int(offset_from_scale + 2 * shape.offset[i]),
2145 )
2146 else:
2147 size = shape[i]
2149 if axis_type == "time":
2150 if tensor_type == "input":
2151 ret.append(TimeInputAxis(size=size, scale=scale))
2152 else:
2153 assert not isinstance(size, ParameterizedSize)
2154 if halo is None:
2155 ret.append(TimeOutputAxis(size=size, scale=scale))
2156 else:
2157 assert not isinstance(size, int)
2158 ret.append(
2159 TimeOutputAxisWithHalo(size=size, scale=scale, halo=halo[i])
2160 )
2162 elif axis_type == "index":
2163 if tensor_type == "input":
2164 ret.append(IndexInputAxis(size=size))
2165 else:
2166 if isinstance(size, ParameterizedSize):
2167 size = DataDependentSize(min=size.min)
2169 ret.append(IndexOutputAxis(size=size))
2170 elif axis_type == "channel":
2171 assert not isinstance(size, ParameterizedSize)
2172 if isinstance(size, SizeReference):
2173 warnings.warn(
2174 "Conversion of channel size from an implicit output shape may be"
2175 + " wrong"
2176 )
2177 ret.append(
2178 ChannelAxis(
2179 channel_names=[f"channel{i}" for i in range(size.offset)]
2180 )
2181 )
2182 else:
2183 ret.append(
2184 ChannelAxis(channel_names=[f"channel{i}" for i in range(size)])
2185 )
2186 elif axis_type == "space":
2187 if tensor_type == "input":
2188 ret.append(SpaceInputAxis(id=AxisId(a), size=size, scale=scale))
2189 else:
2190 assert not isinstance(size, ParameterizedSize)
2191 if halo is None or halo[i] == 0:
2192 ret.append(SpaceOutputAxis(id=AxisId(a), size=size, scale=scale))
2193 elif isinstance(size, int):
2194 raise NotImplementedError(
2195 f"output axis with halo and fixed size (here {size}) not allowed"
2196 )
2197 else:
2198 ret.append(
2199 SpaceOutputAxisWithHalo(
2200 id=AxisId(a), size=size, scale=scale, halo=halo[i]
2201 )
2202 )
2204 return ret
2207def _axes_letters_to_ids(
2208 axes: Optional[str],
2209) -> Optional[List[AxisId]]:
2210 if axes is None:
2211 return None
2213 return [AxisId(a) for a in axes]
2216def _get_complement_v04_axis(
2217 tensor_axes: Sequence[str], axes: Optional[Sequence[str]]
2218) -> Optional[AxisId]:
2219 if axes is None:
2220 return None
2222 non_complement_axes = set(axes) | {"b"}
2223 complement_axes = [a for a in tensor_axes if a not in non_complement_axes]
2224 if len(complement_axes) > 1:
2225 raise ValueError(
2226 f"Expected none or a single complement axis, but axes '{axes}' "
2227 + f"for tensor dims '{tensor_axes}' leave '{complement_axes}'."
2228 )
2230 return None if not complement_axes else AxisId(complement_axes[0])
2233def _convert_proc(
2234 p: Union[_PreprocessingDescr_v0_4, _PostprocessingDescr_v0_4],
2235 tensor_axes: Sequence[str],
2236) -> Union[PreprocessingDescr, PostprocessingDescr]:
2237 if isinstance(p, _BinarizeDescr_v0_4):
2238 return BinarizeDescr(kwargs=BinarizeKwargs(threshold=p.kwargs.threshold))
2239 elif isinstance(p, _ClipDescr_v0_4):
2240 return ClipDescr(kwargs=ClipKwargs(min=p.kwargs.min, max=p.kwargs.max))
2241 elif isinstance(p, _SigmoidDescr_v0_4):
2242 return SigmoidDescr()
2243 elif isinstance(p, _ScaleLinearDescr_v0_4):
2244 axes = _axes_letters_to_ids(p.kwargs.axes)
2245 if p.kwargs.axes is None:
2246 axis = None
2247 else:
2248 axis = _get_complement_v04_axis(tensor_axes, p.kwargs.axes)
2250 if axis is None:
2251 assert not isinstance(p.kwargs.gain, list)
2252 assert not isinstance(p.kwargs.offset, list)
2253 kwargs = ScaleLinearKwargs(gain=p.kwargs.gain, offset=p.kwargs.offset)
2254 else:
2255 kwargs = ScaleLinearAlongAxisKwargs(
2256 axis=axis, gain=p.kwargs.gain, offset=p.kwargs.offset
2257 )
2258 return ScaleLinearDescr(kwargs=kwargs)
2259 elif isinstance(p, _ScaleMeanVarianceDescr_v0_4):
2260 return ScaleMeanVarianceDescr(
2261 kwargs=ScaleMeanVarianceKwargs(
2262 axes=_axes_letters_to_ids(p.kwargs.axes),
2263 reference_tensor=TensorId(str(p.kwargs.reference_tensor)),
2264 eps=p.kwargs.eps,
2265 )
2266 )
2267 elif isinstance(p, _ZeroMeanUnitVarianceDescr_v0_4):
2268 if p.kwargs.mode == "fixed":
2269 mean = p.kwargs.mean
2270 std = p.kwargs.std
2271 assert mean is not None
2272 assert std is not None
2274 axis = _get_complement_v04_axis(tensor_axes, p.kwargs.axes)
2276 if axis is None:
2277 if isinstance(mean, list):
2278 raise ValueError("Expected single float value for mean, not <list>")
2279 if isinstance(std, list):
2280 raise ValueError("Expected single float value for std, not <list>")
2281 return FixedZeroMeanUnitVarianceDescr(
2282 kwargs=FixedZeroMeanUnitVarianceKwargs.model_construct(
2283 mean=mean,
2284 std=std,
2285 )
2286 )
2287 else:
2288 if not isinstance(mean, list):
2289 mean = [float(mean)]
2290 if not isinstance(std, list):
2291 std = [float(std)]
2293 return FixedZeroMeanUnitVarianceDescr(
2294 kwargs=FixedZeroMeanUnitVarianceAlongAxisKwargs(
2295 axis=axis, mean=mean, std=std
2296 )
2297 )
2299 else:
2300 axes = _axes_letters_to_ids(p.kwargs.axes) or []
2301 if p.kwargs.mode == "per_dataset":
2302 axes = [AxisId("batch")] + axes
2303 if not axes:
2304 axes = None
2305 return ZeroMeanUnitVarianceDescr(
2306 kwargs=ZeroMeanUnitVarianceKwargs(axes=axes, eps=p.kwargs.eps)
2307 )
2309 elif isinstance(p, _ScaleRangeDescr_v0_4):
2310 return ScaleRangeDescr(
2311 kwargs=ScaleRangeKwargs(
2312 axes=_axes_letters_to_ids(p.kwargs.axes),
2313 min_percentile=p.kwargs.min_percentile,
2314 max_percentile=p.kwargs.max_percentile,
2315 eps=p.kwargs.eps,
2316 )
2317 )
2318 else:
2319 assert_never(p)
2322class _InputTensorConv(
2323 Converter[
2324 _InputTensorDescr_v0_4,
2325 InputTensorDescr,
2326 FileSource,
2327 Optional[FileSource],
2328 Mapping[_TensorName_v0_4, Mapping[str, int]],
2329 ]
2330):
2331 def _convert(
2332 self,
2333 src: _InputTensorDescr_v0_4,
2334 tgt: "type[InputTensorDescr] | type[dict[str, Any]]",
2335 test_tensor: FileSource,
2336 sample_tensor: Optional[FileSource],
2337 size_refs: Mapping[_TensorName_v0_4, Mapping[str, int]],
2338 ) -> "InputTensorDescr | dict[str, Any]":
2339 axes: List[InputAxis] = convert_axes( # pyright: ignore[reportAssignmentType]
2340 src.axes,
2341 shape=src.shape,
2342 tensor_type="input",
2343 halo=None,
2344 size_refs=size_refs,
2345 )
2346 prep: List[PreprocessingDescr] = []
2347 for p in src.preprocessing:
2348 cp = _convert_proc(p, src.axes)
2349 assert not isinstance(
2350 cp,
2351 (
2352 CellposeFlowDynamicsDescr,
2353 CustomProcessingDescr,
2354 ScaleMeanVarianceDescr,
2355 StardistPostprocessingDescr,
2356 ),
2357 )
2358 prep.append(cp)
2360 prep.append(EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype="float32")))
2362 return tgt(
2363 axes=axes,
2364 id=TensorId(str(src.name)),
2365 test_tensor=FileDescr(source=test_tensor),
2366 sample_tensor=(
2367 None if sample_tensor is None else FileDescr(source=sample_tensor)
2368 ),
2369 data=dict(type=src.data_type), # pyright: ignore[reportArgumentType]
2370 preprocessing=prep,
2371 )
2374_input_tensor_conv = _InputTensorConv(_InputTensorDescr_v0_4, InputTensorDescr)
2377class OutputTensorDescr(TensorDescrBase[OutputAxis]):
2378 id: TensorId = TensorId("output")
2379 """Output tensor id.
2380 No duplicates are allowed across all inputs and outputs."""
2382 postprocessing: List[PostprocessingDescr] = Field(
2383 default_factory=cast(Callable[[], List[PostprocessingDescr]], list)
2384 )
2385 """Description of how this output should be postprocessed.
2387 note: `postprocessing` always ends with an 'ensure_dtype' operation.
2388 If not given this is added to cast to this tensor's `data.type`.
2389 """
2391 @model_validator(mode="after")
2392 def _validate_postprocessing_kwargs(self) -> Self:
2393 axes_ids = [a.id for a in self.axes]
2394 for p in self.postprocessing:
2395 kwargs_axes = p.kwargs.get("axes")
2396 if kwargs_axes is None:
2397 continue
2399 if not isinstance(kwargs_axes, collections.abc.Sequence):
2400 raise ValueError(
2401 f"expected `axes` sequence, but got {type(kwargs_axes)}"
2402 )
2404 kwargs_axes_seq: Sequence[Any] = cast(Sequence[Any], kwargs_axes)
2405 if any(a not in axes_ids for a in kwargs_axes_seq):
2406 raise ValueError("`kwargs.axes` needs to be subset of axes ids")
2408 if isinstance(self.data, (NominalOrOrdinalDataDescr, IntervalOrRatioDataDescr)):
2409 dtype = self.data.type
2410 else:
2411 dtype = self.data[0].type
2413 # ensure `postprocessing` ends with `EnsureDtypeDescr` or `BinarizeDescr`
2414 if not self.postprocessing or not isinstance(
2415 self.postprocessing[-1], (EnsureDtypeDescr, BinarizeDescr)
2416 ):
2417 self.postprocessing.append(
2418 EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype=dtype))
2419 )
2420 return self
2423class _OutputTensorConv(
2424 Converter[
2425 _OutputTensorDescr_v0_4,
2426 OutputTensorDescr,
2427 FileSource,
2428 Optional[FileSource],
2429 Mapping[_TensorName_v0_4, Mapping[str, int]],
2430 ]
2431):
2432 def _convert(
2433 self,
2434 src: _OutputTensorDescr_v0_4,
2435 tgt: "type[OutputTensorDescr] | type[dict[str, Any]]",
2436 test_tensor: FileSource,
2437 sample_tensor: Optional[FileSource],
2438 size_refs: Mapping[_TensorName_v0_4, Mapping[str, int]],
2439 ) -> "OutputTensorDescr | dict[str, Any]":
2440 # TODO: split convert_axes into convert_output_axes and convert_input_axes
2441 axes: List[OutputAxis] = convert_axes( # pyright: ignore[reportAssignmentType]
2442 src.axes,
2443 shape=src.shape,
2444 tensor_type="output",
2445 halo=src.halo,
2446 size_refs=size_refs,
2447 )
2448 data_descr: Dict[str, Any] = dict(type=src.data_type)
2449 if data_descr["type"] == "bool":
2450 data_descr["values"] = [False, True]
2452 return tgt(
2453 axes=axes,
2454 id=TensorId(str(src.name)),
2455 test_tensor=FileDescr(source=test_tensor),
2456 sample_tensor=(
2457 None if sample_tensor is None else FileDescr(source=sample_tensor)
2458 ),
2459 data=data_descr, # pyright: ignore[reportArgumentType]
2460 postprocessing=[_convert_proc(p, src.axes) for p in src.postprocessing],
2461 )
2464_output_tensor_conv = _OutputTensorConv(_OutputTensorDescr_v0_4, OutputTensorDescr)
2467TensorDescr = Union[InputTensorDescr, OutputTensorDescr]
2470def get_halos(
2471 tensors: Mapping[TensorId, TensorDescr],
2472 /,
2473) -> Dict[TensorId, Dict[AxisId, Tuple[int, int]]]:
2474 """Get all input and output halos from tensor descriptions.
2476 Note:
2477 - Input halos are to be padded
2478 - Output halos are to be cropped
2479 """
2480 halos: Dict[TensorId, Dict[AxisId, Tuple[int, int]]] = {}
2481 for descr in tensors.values():
2482 if isinstance(descr, InputTensorDescr):
2483 continue
2484 for axis in descr.axes:
2485 if not isinstance(axis, WithHalo):
2486 continue
2488 ref_scale = next(
2489 a
2490 for a in tensors[axis.size.tensor_id].axes
2491 if a.id == axis.size.axis_id
2492 ).scale
2494 # set output halo (to be cropped)
2495 halos.setdefault(descr.id, {})[axis.id] = (axis.halo, axis.halo)
2496 # set input halo (to be padded)
2497 pad_width = int(axis.halo / axis.scale * ref_scale)
2498 halos.setdefault(axis.size.tensor_id, {})[axis.size.axis_id] = (
2499 pad_width,
2500 pad_width,
2501 )
2503 return halos
2506def validate_tensors(
2507 tensors: Mapping[TensorId, Tuple[TensorDescr, Optional[NDArray[Any]]]],
2508 tensor_origin: Literal[
2509 "source", "test_tensor"
2510 ] = "source", # for more precise error messages
2511 *,
2512 pad_inputs: Union[bool, Literal["allow"]] = True,
2513 crop_outputs: Union[bool, Literal["allow"]] = True,
2514):
2515 """Validate all inputs (and optionally output tensors) against their tensor descriptions.
2517 Args:
2518 tensors: Mapping of tensor id to a tuple of tensor description and optional numpy array.
2519 tensor_origin: String to use in error messages to indicate the origin of the tensors being validated.
2520 pad_inputs: Wether to apply/allow padding of inputs before shape comparison
2521 crop_outputs: Wether to apply/allow cropping of outputs before shape comparison.
2522 """
2523 all_tensor_axes: Dict[TensorId, Dict[AxisId, Tuple[AnyAxis, Optional[int]]]] = {}
2525 def e_msg_location(d: TensorDescr):
2526 return f"{'inputs' if isinstance(d, InputTensorDescr) else 'outputs'}[{d.id}]"
2528 for descr, array in tensors.values():
2529 if array is None:
2530 axis_sizes = {a.id: None for a in descr.axes}
2531 else:
2532 try:
2533 axis_sizes = descr.get_axis_sizes_for_array(array)
2534 except ValueError as e:
2535 raise ValueError(f"{e_msg_location(descr)} {e}")
2537 all_tensor_axes[descr.id] = {a.id: (a, axis_sizes[a.id]) for a in descr.axes}
2539 # get halos to be padded/cropped to validate against halo-adjusted sizes
2540 io_halos = get_halos({k: v[0] for k, v in tensors.items()})
2542 for descr, array in tensors.values():
2543 if array is None:
2544 continue
2546 if descr.dtype in ("float32", "float64"):
2547 invalid_test_tensor_dtype = array.dtype.name not in (
2548 "float32",
2549 "float64",
2550 "uint8",
2551 "int8",
2552 "uint16",
2553 "int16",
2554 "uint32",
2555 "int32",
2556 "uint64",
2557 "int64",
2558 )
2559 else:
2560 invalid_test_tensor_dtype = array.dtype.name != descr.dtype
2562 if invalid_test_tensor_dtype:
2563 raise ValueError(
2564 f"{tensor_origin} data type '{array.dtype.name}' does not"
2565 + f" match described {e_msg_location(descr)}.dtype '{descr.dtype}'"
2566 )
2568 if array.min() > -1e-4 and array.max() < 1e-4:
2569 raise ValueError(
2570 "Output values are too small for reliable testing."
2571 + f" Values <-1e5 or >=1e5 must be present in {tensor_origin}"
2572 )
2574 for a in descr.axes:
2575 actual_size = all_tensor_axes[descr.id][a.id][1]
2577 if actual_size is None:
2578 continue
2580 if a.size is None:
2581 continue
2583 # add padding width to actual tensor size
2584 total_axis_halo = sum(io_halos.get(descr.id, {}).get(a.id, (0, 0)))
2585 if isinstance(descr, InputTensorDescr):
2586 # pad input halos
2587 actual_size_with_halo = actual_size + total_axis_halo
2588 if pad_inputs is True:
2589 check_sizes = {actual_size_with_halo}
2590 size_hint = " (after padding input halo)"
2591 elif pad_inputs == "allow":
2592 check_sizes = {actual_size, actual_size_with_halo}
2593 size_hint = " (with or without padding input halo)"
2594 elif pad_inputs is False:
2595 check_sizes = {actual_size}
2596 size_hint = ""
2597 else:
2598 assert_never(pad_inputs)
2600 elif isinstance(descr, OutputTensorDescr):
2601 # crop output halos
2602 actual_size_with_halo = max(0, actual_size - total_axis_halo)
2603 if crop_outputs is True:
2604 check_sizes = {actual_size_with_halo}
2605 size_hint = " (after cropping output halo)"
2606 elif crop_outputs == "allow":
2607 check_sizes = {actual_size, actual_size_with_halo}
2608 size_hint = " (with or without cropping output halo)"
2609 elif crop_outputs is False:
2610 check_sizes = {actual_size}
2611 size_hint = ""
2612 else:
2613 assert_never(crop_outputs)
2614 else:
2615 assert_never(descr)
2617 del actual_size # make sure we explicitly use unchanged or halo-adjusted size from here on
2619 if isinstance(a.size, int):
2620 if a.size not in check_sizes:
2621 raise ValueError(
2622 f"{e_msg_location(descr)}.axes[{a.id}]: {tensor_origin} axis "
2623 + f"has incompatible size {check_sizes}{size_hint}, expected {a.size}"
2624 )
2625 elif isinstance(a.size, (ParameterizedSize, DataDependentSize)):
2626 _ = try_all_raise_last(
2627 (partial(a.size.validate_size, s) for s in check_sizes),
2628 f"{e_msg_location(descr)}.axes[{a.id}]: {tensor_origin} axis ",
2629 )
2630 elif isinstance(a.size, SizeReference):
2631 ref_tensor_axes = all_tensor_axes.get(a.size.tensor_id)
2632 if ref_tensor_axes is None:
2633 raise ValueError(
2634 f"{e_msg_location(descr)}.axes[{a.id}].size.tensor_id: Unknown tensor"
2635 + f" reference '{a.size.tensor_id}', available: {list(all_tensor_axes)}"
2636 )
2638 ref_axis, ref_size = ref_tensor_axes.get(a.size.axis_id, (None, None))
2639 if ref_axis is None or ref_size is None:
2640 raise ValueError(
2641 f"{e_msg_location(descr)}.axes[{a.id}].size.axis_id: Unknown tensor axis"
2642 + f" reference '{a.size.tensor_id}.{a.size.axis_id}, available: {list(ref_tensor_axes)}"
2643 )
2645 if a.unit != ref_axis.unit:
2646 raise ValueError(
2647 f"{e_msg_location(descr)}.axes[{a.id}].size: `SizeReference` requires"
2648 + " axis and reference axis to have the same `unit`, but"
2649 + f" {a.unit}!={ref_axis.unit}"
2650 )
2652 if (
2653 expected_size := (
2654 ref_size * ref_axis.scale / a.scale + a.size.offset
2655 )
2656 ) not in check_sizes:
2657 raise ValueError(
2658 f"{e_msg_location(descr)}.{tensor_origin}: axis '{a.id}' of size"
2659 + f" {check_sizes} invalid for referenced size {ref_size};"
2660 + f" expected {expected_size}"
2661 )
2662 else:
2663 assert_never(a.size)
2666FileDescr_dependencies = Annotated[
2667 FileDescr_package,
2668 WithSuffix((".yaml", ".yml"), case_sensitive=True),
2669 Field(examples=[dict(source="environment.yaml")]),
2670]
2673class _ArchitectureCallableDescr(Node):
2674 callable: Annotated[Identifier, Field(examples=["MyNetworkClass", "get_my_model"])]
2675 """Identifier of the callable that returns a torch.nn.Module instance."""
2677 kwargs: Dict[str, YamlValue] = Field(
2678 default_factory=cast(Callable[[], Dict[str, YamlValue]], dict)
2679 )
2680 """key word arguments for the `callable`"""
2683class ArchitectureFromFileDescr(_ArchitectureCallableDescr, FileDescr):
2684 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
2685 """Architecture source file"""
2687 @model_serializer(mode="wrap", when_used="unless-none")
2688 def _serialize(self, nxt: SerializerFunctionWrapHandler, info: SerializationInfo):
2689 return package_file_descr_serializer(self, nxt, info)
2692class ArchitectureFromLibraryDescr(_ArchitectureCallableDescr):
2693 import_from: str
2694 """Where to import the callable from, i.e. `from <import_from> import <callable>`"""
2697class _ArchFileConv(
2698 Converter[
2699 _CallableFromFile_v0_4,
2700 ArchitectureFromFileDescr,
2701 Optional[Sha256],
2702 Dict[str, Any],
2703 ]
2704):
2705 def _convert(
2706 self,
2707 src: _CallableFromFile_v0_4,
2708 tgt: "type[ArchitectureFromFileDescr | dict[str, Any]]",
2709 sha256: Optional[Sha256],
2710 kwargs: Dict[str, Any],
2711 ) -> "ArchitectureFromFileDescr | dict[str, Any]":
2712 if src.startswith("http") and src.count(":") == 2:
2713 http, source, callable_ = src.split(":")
2714 source = ":".join((http, source))
2715 elif not src.startswith("http") and src.count(":") == 1:
2716 source, callable_ = src.split(":")
2717 else:
2718 source = str(src)
2719 callable_ = str(src)
2720 return tgt(
2721 callable=Identifier(callable_),
2722 source=cast(FileSource, source),
2723 sha256=sha256,
2724 kwargs=kwargs,
2725 )
2728_arch_file_conv = _ArchFileConv(_CallableFromFile_v0_4, ArchitectureFromFileDescr)
2731class _ArchLibConv(
2732 Converter[
2733 _CallableFromDepencency_v0_4, ArchitectureFromLibraryDescr, Dict[str, Any]
2734 ]
2735):
2736 def _convert(
2737 self,
2738 src: _CallableFromDepencency_v0_4,
2739 tgt: "type[ArchitectureFromLibraryDescr | dict[str, Any]]",
2740 kwargs: Dict[str, Any],
2741 ) -> "ArchitectureFromLibraryDescr | dict[str, Any]":
2742 *mods, callable_ = src.split(".")
2743 import_from = ".".join(mods)
2744 return tgt(
2745 import_from=import_from, callable=Identifier(callable_), kwargs=kwargs
2746 )
2749_arch_lib_conv = _ArchLibConv(
2750 _CallableFromDepencency_v0_4, ArchitectureFromLibraryDescr
2751)
2754class WeightsEntryDescrBase(FileDescr):
2755 type: ClassVar[WeightsFormat]
2756 weights_format_name: ClassVar[str] # human readable
2758 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
2759 """Source of the weights file."""
2761 authors: Optional[List[Author]] = None
2762 """Authors
2763 Either the person(s) that have trained this model resulting in the original weights file.
2764 (If this is the initial weights entry, i.e. it does not have a `parent`)
2765 Or the person(s) who have converted the weights to this weights format.
2766 (If this is a child weight, i.e. it has a `parent` field)
2767 """
2769 parent: Annotated[
2770 Optional[WeightsFormat], Field(examples=["pytorch_state_dict"])
2771 ] = None
2772 """The source weights these weights were converted from.
2773 For example, if a model's weights were converted from the `pytorch_state_dict` format to `torchscript`,
2774 The `pytorch_state_dict` weights entry has no `parent` and is the parent of the `torchscript` weights.
2775 All weight entries except one (the initial set of weights resulting from training the model),
2776 need to have this field."""
2778 comment: str = ""
2779 """A comment about this weights entry, for example how these weights were created."""
2781 @model_validator(mode="after")
2782 def _validate(self) -> Self:
2783 if self.type == self.parent:
2784 raise ValueError("Weights entry can't be it's own parent.")
2786 return self
2788 @model_serializer(mode="wrap", when_used="unless-none")
2789 def _serialize(self, nxt: SerializerFunctionWrapHandler, info: SerializationInfo):
2790 return package_file_descr_serializer(self, nxt, info)
2793class KerasHdf5WeightsDescr(WeightsEntryDescrBase):
2794 type: ClassVar[WeightsFormat] = "keras_hdf5"
2795 weights_format_name: ClassVar[str] = "Keras HDF5"
2796 tensorflow_version: Version
2797 """TensorFlow version used to create these weights."""
2800class KerasV3WeightsDescr(WeightsEntryDescrBase):
2801 type: ClassVar[WeightsFormat] = "keras_v3"
2802 weights_format_name: ClassVar[str] = "Keras v3"
2803 keras_version: Annotated[Version, Ge(Version(3))]
2804 """Keras version used to create these weights."""
2805 backend: Tuple[Literal["tensorflow", "jax", "torch"], Version]
2806 """Keras backend used to create these weights."""
2807 source: Annotated[
2808 FileSource,
2809 AfterValidator(wo_special_file_name),
2810 WithSuffix(".keras", case_sensitive=True),
2811 ]
2812 """Source of the .keras weights file."""
2815FileDescr_external_data = Annotated[
2816 FileDescr_package,
2817 WithSuffix(".data", case_sensitive=True),
2818 Field(examples=[dict(source="weights.onnx.data")]),
2819]
2822class OnnxWeightsDescr(WeightsEntryDescrBase):
2823 type: ClassVar[WeightsFormat] = "onnx"
2824 weights_format_name: ClassVar[str] = "ONNX"
2825 opset_version: Annotated[int, Ge(7)]
2826 """ONNX opset version"""
2828 external_data: Optional[FileDescr_external_data] = None
2829 """Source of the external ONNX data file holding the weights.
2830 (If present **source** holds the ONNX architecture without weights)."""
2832 @model_validator(mode="after")
2833 def _validate_external_data_unique_file_name(self) -> Self:
2834 if self.external_data is not None and (
2835 extract_file_name(self.source)
2836 == extract_file_name(self.external_data.source)
2837 ):
2838 raise ValueError(
2839 f"ONNX `external_data` file name '{extract_file_name(self.external_data.source)}'"
2840 + " must be different from ONNX `source` file name."
2841 )
2843 return self
2846class PytorchStateDictWeightsDescr(WeightsEntryDescrBase):
2847 type: ClassVar[WeightsFormat] = "pytorch_state_dict"
2848 weights_format_name: ClassVar[str] = "Pytorch State Dict"
2849 architecture: Union[ArchitectureFromFileDescr, ArchitectureFromLibraryDescr]
2850 pytorch_version: Version
2851 """Version of the PyTorch library used.
2852 If `architecture.depencencies` is specified it has to include pytorch and any version pinning has to be compatible.
2853 """
2854 dependencies: Optional[FileDescr_dependencies] = None
2855 """Custom depencies beyond pytorch described in a Conda environment file.
2856 Allows to specify custom dependencies, see conda docs:
2857 - [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)
2858 - [Creating an environment file manually](https://conda.io/projects/conda/en/latest/user-guide/tasks/manage-environments.html#creating-an-environment-file-manually)
2860 The conda environment file should include pytorch and any version pinning has to be compatible with
2861 **pytorch_version**.
2862 """
2863 strict: bool = True
2864 """Whether to allow missing or unexpected keys or to be strict about the architecture matching the state dict weights."""
2867class TensorflowJsWeightsDescr(WeightsEntryDescrBase):
2868 type: ClassVar[WeightsFormat] = "tensorflow_js"
2869 weights_format_name: ClassVar[str] = "Tensorflow.js"
2870 tensorflow_version: Version
2871 """Version of the TensorFlow library used."""
2873 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
2874 """The multi-file weights.
2875 All required files/folders should be a zip archive."""
2878class TensorflowSavedModelBundleWeightsDescr(WeightsEntryDescrBase):
2879 type: ClassVar[WeightsFormat] = "tensorflow_saved_model_bundle"
2880 weights_format_name: ClassVar[str] = "Tensorflow Saved Model"
2881 tensorflow_version: Version
2882 """Version of the TensorFlow library used."""
2884 dependencies: Optional[FileDescr_dependencies] = None
2885 """Custom dependencies beyond tensorflow.
2886 Should include tensorflow and any version pinning has to be compatible with **tensorflow_version**."""
2888 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
2889 """The multi-file weights.
2890 All required files/folders should be a zip archive."""
2893class TorchscriptWeightsDescr(WeightsEntryDescrBase):
2894 type: ClassVar[WeightsFormat] = "torchscript"
2895 weights_format_name: ClassVar[str] = "TorchScript"
2896 pytorch_version: Version
2897 """Version of the PyTorch library used."""
2900SpecificWeightsDescr = Union[
2901 KerasHdf5WeightsDescr,
2902 KerasV3WeightsDescr,
2903 OnnxWeightsDescr,
2904 PytorchStateDictWeightsDescr,
2905 TensorflowJsWeightsDescr,
2906 TensorflowSavedModelBundleWeightsDescr,
2907 TorchscriptWeightsDescr,
2908]
2911class WeightsDescr(Node):
2912 keras_hdf5: Optional[KerasHdf5WeightsDescr] = None
2913 keras_v3: Optional[KerasV3WeightsDescr] = None
2914 onnx: Optional[OnnxWeightsDescr] = None
2915 pytorch_state_dict: Optional[PytorchStateDictWeightsDescr] = None
2916 tensorflow_js: Optional[TensorflowJsWeightsDescr] = None
2917 tensorflow_saved_model_bundle: Optional[TensorflowSavedModelBundleWeightsDescr] = (
2918 None
2919 )
2920 torchscript: Optional[TorchscriptWeightsDescr] = None
2922 @model_validator(mode="after")
2923 def check_entries(self) -> Self:
2924 entries = {wtype for wtype, entry in self if entry is not None}
2926 if not entries:
2927 raise ValueError("Missing weights entry")
2929 entries_wo_parent = {
2930 wtype
2931 for wtype, entry in self
2932 if entry is not None and hasattr(entry, "parent") and entry.parent is None
2933 }
2934 if len(entries_wo_parent) != 1:
2935 issue_warning(
2936 "Exactly one weights entry may not specify the `parent` field (got"
2937 + " {value}). That entry is considered the original set of model weights."
2938 + " Other weight formats are created through conversion of the orignal or"
2939 + " already converted weights. They have to reference the weights format"
2940 + " they were converted from as their `parent`.",
2941 value=len(entries_wo_parent),
2942 field="weights",
2943 )
2945 for wtype, entry in self:
2946 if entry is None:
2947 continue
2949 assert hasattr(entry, "type")
2950 assert hasattr(entry, "parent")
2951 assert wtype == entry.type
2952 if (
2953 entry.parent is not None and entry.parent not in entries
2954 ): # self reference checked for `parent` field
2955 raise ValueError(
2956 f"`weights.{wtype}.parent={entry.parent} not in specified weight"
2957 + f" formats: {entries}"
2958 )
2960 return self
2962 def __getitem__(
2963 self,
2964 key: WeightsFormat,
2965 ):
2966 if key == "keras_hdf5":
2967 ret = self.keras_hdf5
2968 elif key == "keras_v3":
2969 ret = self.keras_v3
2970 elif key == "onnx":
2971 ret = self.onnx
2972 elif key == "pytorch_state_dict":
2973 ret = self.pytorch_state_dict
2974 elif key == "tensorflow_js":
2975 ret = self.tensorflow_js
2976 elif key == "tensorflow_saved_model_bundle":
2977 ret = self.tensorflow_saved_model_bundle
2978 elif key == "torchscript":
2979 ret = self.torchscript
2980 else:
2981 raise KeyError(key)
2983 if ret is None:
2984 raise KeyError(key)
2986 return ret
2988 @overload
2989 def __setitem__(
2990 self, key: Literal["keras_hdf5"], value: Optional[KerasHdf5WeightsDescr]
2991 ) -> None: ...
2992 @overload
2993 def __setitem__(
2994 self, key: Literal["keras_v3"], value: Optional[KerasV3WeightsDescr]
2995 ) -> None: ...
2996 @overload
2997 def __setitem__(
2998 self, key: Literal["onnx"], value: Optional[OnnxWeightsDescr]
2999 ) -> None: ...
3000 @overload
3001 def __setitem__(
3002 self,
3003 key: Literal["pytorch_state_dict"],
3004 value: Optional[PytorchStateDictWeightsDescr],
3005 ) -> None: ...
3006 @overload
3007 def __setitem__(
3008 self, key: Literal["tensorflow_js"], value: Optional[TensorflowJsWeightsDescr]
3009 ) -> None: ...
3010 @overload
3011 def __setitem__(
3012 self,
3013 key: Literal["tensorflow_saved_model_bundle"],
3014 value: Optional[TensorflowSavedModelBundleWeightsDescr],
3015 ) -> None: ...
3016 @overload
3017 def __setitem__(
3018 self, key: Literal["torchscript"], value: Optional[TorchscriptWeightsDescr]
3019 ) -> None: ...
3021 def __setitem__(
3022 self,
3023 key: WeightsFormat,
3024 value: Optional[SpecificWeightsDescr],
3025 ):
3026 if key == "keras_hdf5":
3027 if value is not None and not isinstance(value, KerasHdf5WeightsDescr):
3028 raise TypeError(
3029 f"Expected KerasHdf5WeightsDescr or None for key 'keras_hdf5', got {type(value)}"
3030 )
3031 self.keras_hdf5 = value
3032 elif key == "keras_v3":
3033 if value is not None and not isinstance(value, KerasV3WeightsDescr):
3034 raise TypeError(
3035 f"Expected KerasV3WeightsDescr or None for key 'keras_v3', got {type(value)}"
3036 )
3037 self.keras_v3 = value
3038 elif key == "onnx":
3039 if value is not None and not isinstance(value, OnnxWeightsDescr):
3040 raise TypeError(
3041 f"Expected OnnxWeightsDescr or None for key 'onnx', got {type(value)}"
3042 )
3043 self.onnx = value
3044 elif key == "pytorch_state_dict":
3045 if value is not None and not isinstance(
3046 value, PytorchStateDictWeightsDescr
3047 ):
3048 raise TypeError(
3049 f"Expected PytorchStateDictWeightsDescr or None for key 'pytorch_state_dict', got {type(value)}"
3050 )
3051 self.pytorch_state_dict = value
3052 elif key == "tensorflow_js":
3053 if value is not None and not isinstance(value, TensorflowJsWeightsDescr):
3054 raise TypeError(
3055 f"Expected TensorflowJsWeightsDescr or None for key 'tensorflow_js', got {type(value)}"
3056 )
3057 self.tensorflow_js = value
3058 elif key == "tensorflow_saved_model_bundle":
3059 if value is not None and not isinstance(
3060 value, TensorflowSavedModelBundleWeightsDescr
3061 ):
3062 raise TypeError(
3063 f"Expected TensorflowSavedModelBundleWeightsDescr or None for key 'tensorflow_saved_model_bundle', got {type(value)}"
3064 )
3065 self.tensorflow_saved_model_bundle = value
3066 elif key == "torchscript":
3067 if value is not None and not isinstance(value, TorchscriptWeightsDescr):
3068 raise TypeError(
3069 f"Expected TorchscriptWeightsDescr or None for key 'torchscript', got {type(value)}"
3070 )
3071 self.torchscript = value
3072 else:
3073 raise KeyError(key)
3075 @property
3076 def available_formats(self) -> Dict[WeightsFormat, SpecificWeightsDescr]:
3077 return {
3078 **({} if self.keras_hdf5 is None else {"keras_hdf5": self.keras_hdf5}),
3079 **({} if self.keras_v3 is None else {"keras_v3": self.keras_v3}),
3080 **({} if self.onnx is None else {"onnx": self.onnx}),
3081 **(
3082 {}
3083 if self.pytorch_state_dict is None
3084 else {"pytorch_state_dict": self.pytorch_state_dict}
3085 ),
3086 **(
3087 {}
3088 if self.tensorflow_js is None
3089 else {"tensorflow_js": self.tensorflow_js}
3090 ),
3091 **(
3092 {}
3093 if self.tensorflow_saved_model_bundle is None
3094 else {
3095 "tensorflow_saved_model_bundle": self.tensorflow_saved_model_bundle
3096 }
3097 ),
3098 **({} if self.torchscript is None else {"torchscript": self.torchscript}),
3099 }
3101 @property
3102 def missing_formats(self) -> Set[WeightsFormat]:
3103 return {
3104 wf for wf in get_args(WeightsFormat) if wf not in self.available_formats
3105 }
3108class LinkedModel(LinkedResourceBase):
3109 """Reference to a bioimage.io model."""
3111 id: ModelId
3112 """A valid model `id` from the bioimage.io collection."""
3115class _DataDepSize(NamedTuple):
3116 min: StrictInt
3117 max: Optional[StrictInt]
3120class _AxisSizes(NamedTuple):
3121 """the lenghts of all axes of model inputs and outputs"""
3123 inputs: Dict[Tuple[TensorId, AxisId], int]
3124 outputs: Dict[Tuple[TensorId, AxisId], Union[int, _DataDepSize]]
3127class _TensorSizes(NamedTuple):
3128 """_AxisSizes as nested dicts"""
3130 inputs: Dict[TensorId, Dict[AxisId, int]]
3131 outputs: Dict[TensorId, Dict[AxisId, Union[int, _DataDepSize]]]
3134class ReproducibilityTolerance(Node, extra="allow"):
3135 """Describes what small numerical differences -- if any -- may be tolerated
3136 in the generated output when executing in different environments.
3138 A tensor element *output* is considered mismatched to the **test_tensor** if
3139 abs(*output* - **test_tensor**) > **absolute_tolerance** + **relative_tolerance** * abs(**test_tensor**).
3140 (Internally we call [numpy.testing.assert_allclose](https://numpy.org/doc/stable/reference/generated/numpy.testing.assert_allclose.html).)
3142 Motivation:
3143 For testing we can request the respective deep learning frameworks to be as
3144 reproducible as possible by setting seeds and chosing deterministic algorithms,
3145 but differences in operating systems, available hardware and installed drivers
3146 may still lead to numerical differences.
3147 """
3149 relative_tolerance: RelativeTolerance = 1e-3
3150 """Maximum relative tolerance of reproduced test tensor."""
3152 absolute_tolerance: AbsoluteTolerance = 1e-3
3153 """Maximum absolute tolerance of reproduced test tensor."""
3155 mismatched_elements_per_million: MismatchedElementsPerMillion = 100
3156 """Maximum number of mismatched elements/pixels per million to tolerate."""
3158 output_ids: Sequence[TensorId] = ()
3159 """Limits the output tensor IDs these reproducibility details apply to."""
3161 weights_formats: Sequence[WeightsFormat] = ()
3162 """Limits the weights formats these details apply to."""
3165class BiasRisksLimitations(Node, extra="allow"):
3166 """Known biases, risks, technical limitations, and recommendations for model use."""
3168 known_biases: str = dedent("""\
3169 In general bioimage models may suffer from biases caused by:
3171 - Imaging protocol dependencies
3172 - Use of a specific cell type
3173 - Species-specific training data limitations
3175 """)
3176 """Biases in training data or model behavior."""
3178 risks: str = dedent("""\
3179 Common risks in bioimage analysis include:
3181 - Erroneously assuming generalization to unseen experimental conditions
3182 - Trusting (overconfident) model outputs without validation
3183 - Misinterpretation of results
3185 """)
3186 """Potential risks in the context of bioimage analysis."""
3188 limitations: Optional[str] = None
3189 """Technical limitations and failure modes."""
3191 recommendations: str = "Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model."
3192 """Mitigation strategies regarding `known_biases`, `risks`, and `limitations`, as well as applicable best practices.
3194 Consider:
3195 - How to use a validation dataset?
3196 - How to manually validate?
3197 - Feasibility of domain adaptation for different experimental setups?
3199 """
3201 def format_md(self) -> str:
3202 if self.limitations is None:
3203 limitations_header = ""
3204 else:
3205 limitations_header = "## Limitations\n\n"
3207 return f"""# Bias, Risks, and Limitations
3209{self.known_biases}
3211{self.risks}
3213{limitations_header}{self.limitations or ""}
3215## Recommendations
3217{self.recommendations}
3219"""
3222class TrainingDetails(Node, extra="allow"):
3223 training_preprocessing: Optional[str] = None
3224 """Detailed image preprocessing steps during model training:
3226 Mention:
3227 - *Normalization methods*
3228 - *Augmentation strategies*
3229 - *Resizing/resampling procedures*
3230 - *Artifact handling*
3232 """
3234 training_epochs: Optional[float] = None
3235 """Number of training epochs."""
3237 training_batch_size: Optional[float] = None
3238 """Batch size used in training."""
3240 initial_learning_rate: Optional[float] = None
3241 """Initial learning rate used in training."""
3243 learning_rate_schedule: Optional[str] = None
3244 """Learning rate schedule used in training."""
3246 loss_function: Optional[str] = None
3247 """Loss function used in training, e.g. nn.MSELoss."""
3249 loss_function_kwargs: Dict[str, YamlValue] = Field(
3250 default_factory=cast(Callable[[], Dict[str, YamlValue]], dict)
3251 )
3252 """key word arguments for the `loss_function`"""
3254 optimizer: Optional[str] = None
3255 """optimizer, e.g. torch.optim.Adam"""
3257 optimizer_kwargs: Dict[str, YamlValue] = Field(
3258 default_factory=cast(Callable[[], Dict[str, YamlValue]], dict)
3259 )
3260 """key word arguments for the `optimizer`"""
3262 regularization: Optional[str] = None
3263 """Regularization techniques used during training, e.g. drop-out or weight decay."""
3265 training_duration: Optional[float] = None
3266 """Total training duration in hours."""
3269class Evaluation(Node, extra="allow"):
3270 model_id: Optional[ModelId] = None
3271 """Model being evaluated."""
3273 dataset_id: DatasetId
3274 """Dataset used for evaluation."""
3276 dataset_source: HttpUrl
3277 """Source of the dataset."""
3279 dataset_role: Literal["train", "validation", "test", "independent", "unknown"]
3280 """Role of the dataset used for evaluation.
3282 - `train`: dataset was (part of) the training data
3283 - `validation`: dataset was (part of) the validation data used during training, e.g. used for model selection or hyperparameter tuning
3284 - `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
3285 - `independent`: dataset is entirely independent test data; not used during training or validation, and acquired from a different source/distribution than training data
3286 - `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.
3287 """
3289 sample_count: int
3290 """Number of evaluated samples."""
3292 evaluation_factors: List[Annotated[str, MaxLen(16)]]
3293 """(Abbreviations of) each evaluation factor.
3295 Evaluation factors are criteria along which model performance is evaluated, e.g. different image conditions
3296 like 'low SNR', 'high cell density', or different biological conditions like 'cell type A', 'cell type B'.
3297 An 'overall' factor may be included to summarize performance across all conditions.
3298 """
3300 evaluation_factors_long: List[str]
3301 """Descriptions (long form) of each evaluation factor."""
3303 metrics: List[Annotated[str, MaxLen(16)]]
3304 """(Abbreviations of) metrics used for evaluation."""
3306 metrics_long: List[str]
3307 """Description of each metric used."""
3309 @model_validator(mode="after")
3310 def _validate_list_lengths(self) -> Self:
3311 if len(self.evaluation_factors) != len(self.evaluation_factors_long):
3312 raise ValueError(
3313 "`evaluation_factors` and `evaluation_factors_long` must have the same length"
3314 )
3316 if len(self.metrics) != len(self.metrics_long):
3317 raise ValueError("`metrics` and `metrics_long` must have the same length")
3319 if len(self.results) != len(self.metrics):
3320 raise ValueError("`results` must have the same number of rows as `metrics`")
3322 for row in self.results:
3323 if len(row) != len(self.evaluation_factors):
3324 raise ValueError(
3325 "`results` must have the same number of columns (in every row) as `evaluation_factors`"
3326 )
3328 return self
3330 results: List[List[Union[str, float, int]]]
3331 """Results for each metric (rows; outer list) and each evaluation factor (columns; inner list)."""
3333 results_summary: Optional[str] = None
3334 """Interpretation of results for general audience.
3336 Consider:
3337 - Overall model performance
3338 - Comparison to existing methods
3339 - Limitations and areas for improvement
3341"""
3343 def format_md(self):
3344 results_header = ["Metric"] + self.evaluation_factors
3345 results_table_cells = [results_header, ["---"] * len(results_header)] + [
3346 [metric] + [str(r) for r in row]
3347 for metric, row in zip(self.metrics, self.results)
3348 ]
3350 results_table = "".join(
3351 "| " + " | ".join(row) + " |\n" for row in results_table_cells
3352 )
3353 factors = "".join(
3354 f"\n - {ef}: {efl}"
3355 for ef, efl in zip(self.evaluation_factors, self.evaluation_factors_long)
3356 )
3357 metrics = "".join(
3358 f"\n - {em}: {eml}" for em, eml in zip(self.metrics, self.metrics_long)
3359 )
3361 return f"""## Testing Data, Factors & Metrics
3363Evaluation of {self.model_id or "this"} model on the {self.dataset_id} dataset (dataset role: {self.dataset_role}).
3365### Testing Data
3367- **Source:** [{self.dataset_id}]({self.dataset_source})
3368- **Size:** {self.sample_count} evaluated samples
3370### Factors
3371{factors}
3373### Metrics
3374{metrics}
3376## Results
3378### Quantitative Results
3380{results_table}
3382### Summary
3384{self.results_summary or "missing"}
3386"""
3389class EnvironmentalImpact(Node, extra="allow"):
3390 """Environmental considerations for model training and deployment.
3392 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).
3393 """
3395 hardware_type: Optional[str] = None
3396 """GPU/CPU specifications"""
3398 hours_used: Optional[float] = None
3399 """Total compute hours"""
3401 cloud_provider: Optional[str] = None
3402 """If applicable"""
3404 compute_region: Optional[str] = None
3405 """Geographic location"""
3407 co2_emitted: Optional[float] = None
3408 """kg CO2 equivalent
3410 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).
3411 """
3413 def format_md(self):
3414 """Filled Markdown template section following [Hugging Face Model Card Template](https://huggingface.co/docs/hub/en/model-card-annotated)."""
3415 if self == self.__class__():
3416 return ""
3418 ret = "# Environmental Impact\n\n"
3419 if self.hardware_type is not None:
3420 ret += f"- **Hardware Type:** {self.hardware_type}\n"
3421 if self.hours_used is not None:
3422 ret += f"- **Hours used:** {self.hours_used}\n"
3423 if self.cloud_provider is not None:
3424 ret += f"- **Cloud Provider:** {self.cloud_provider}\n"
3425 if self.compute_region is not None:
3426 ret += f"- **Compute Region:** {self.compute_region}\n"
3427 if self.co2_emitted is not None:
3428 ret += f"- **Carbon Emitted:** {self.co2_emitted} kg CO2e\n"
3430 return ret + "\n"
3433class BioimageioConfig(Node, extra="allow"):
3434 reproducibility_tolerance: Sequence[ReproducibilityTolerance] = ()
3435 """Tolerances to allow when reproducing the model's test outputs
3436 from the model's test inputs.
3437 Only the first entry matching tensor id and weights format is considered.
3438 """
3440 funded_by: Optional[str] = None
3441 """Funding agency, grant number if applicable"""
3443 architecture_type: Optional[Annotated[str, MaxLen(32)]] = (
3444 None # TODO: add to differentiated tags
3445 )
3446 """Model architecture type, e.g., 3D U-Net, ResNet, transformer"""
3448 architecture_description: Optional[str] = None
3449 """Text description of model architecture."""
3451 modality: Optional[str] = None # TODO: add to differentiated tags
3452 """Input modality, e.g., fluorescence microscopy, electron microscopy"""
3454 target_structure: List[str] = Field( # TODO: add to differentiated tags
3455 default_factory=cast(Callable[[], List[str]], list)
3456 )
3457 """Biological structure(s) the model is designed to analyze, e.g., nuclei, mitochondria, cells"""
3459 task: Optional[str] = None # TODO: add to differentiated tags
3460 """Bioimage-specific task type, e.g., segmentation, classification, detection, denoising"""
3462 new_version: Optional[ModelId] = None
3463 """A new version of this model exists with a different model id."""
3465 out_of_scope_use: Optional[str] = None
3466 """Describe how the model may be misused in bioimage analysis contexts and what users should **not** do with the model."""
3468 bias_risks_limitations: BiasRisksLimitations = Field(
3469 default_factory=BiasRisksLimitations.model_construct
3470 )
3471 """Description of known bias, risks, and technical limitations for in-scope model use."""
3473 model_parameter_count: Optional[int] = None
3474 """Total number of model parameters."""
3476 training: TrainingDetails = Field(default_factory=TrainingDetails.model_construct)
3477 """Details on how the model was trained."""
3479 inference_time: Optional[str] = None
3480 """Average inference time per image/tile. Specify hardware and image size. Multiple examples can be given."""
3482 memory_requirements_inference: Optional[str] = None
3483 """GPU memory needed for inference. Multiple examples with different image size can be given."""
3485 memory_requirements_training: Optional[str] = None
3486 """GPU memory needed for training. Multiple examples with different image/batch sizes can be given."""
3488 evaluations: List[Evaluation] = Field(
3489 default_factory=cast(Callable[[], List[Evaluation]], list)
3490 )
3491 """Quantitative model evaluations.
3493 Note:
3494 At the moment we recommend to include only a single test dataset
3495 (with evaluation factors that may mark subsets of the dataset)
3496 to avoid confusion and make the presentation of results cleaner.
3497 """
3499 environmental_impact: EnvironmentalImpact = Field(
3500 default_factory=EnvironmentalImpact.model_construct
3501 )
3502 """Environmental considerations for model training and deployment"""
3505class Config(Node, extra="allow"):
3506 bioimageio: BioimageioConfig = Field(
3507 default_factory=BioimageioConfig.model_construct
3508 )
3509 stardist: YamlValue = None
3512class ModelDescr(GenericModelDescrBase):
3513 """Specification of the fields used in a bioimage.io-compliant RDF to describe AI models with pretrained weights.
3514 These fields are typically stored in a YAML file which we call a model resource description file (model RDF).
3515 """
3517 implemented_format_version: ClassVar[Literal["0.5.12"]] = "0.5.12"
3518 if TYPE_CHECKING:
3519 format_version: Literal["0.5.12"] = "0.5.12"
3520 else:
3521 format_version: Literal["0.5.12"]
3522 """Version of the bioimage.io model description specification used.
3523 When creating a new model always use the latest micro/patch version described here.
3524 The `format_version` is important for any consumer software to understand how to parse the fields.
3525 """
3527 implemented_type: ClassVar[Literal["model"]] = "model"
3528 if TYPE_CHECKING:
3529 type: Literal["model"] = "model"
3530 else:
3531 type: Literal["model"]
3532 """Specialized resource type 'model'"""
3534 id: Optional[ModelId] = None
3535 """bioimage.io-wide unique resource identifier
3536 assigned by bioimage.io; version **un**specific."""
3538 authors: FAIR[List[Author]] = Field(
3539 default_factory=cast(Callable[[], List[Author]], list)
3540 )
3541 """The authors are the creators of the model RDF and the primary points of contact."""
3543 documentation: FAIR[Optional[FileDescr_documentation]] = None
3544 """Additional model documentation.
3545 The recommended documentation source file name is `README.md`. An `.md` suffix is mandatory.
3546 The documentation should include a '#[#] Validation' (sub)section
3547 with details on how to quantitatively validate the model on unseen data."""
3549 @field_validator("documentation", mode="after")
3550 @classmethod
3551 def _validate_documentation(cls, value: Optional[FileDescr]) -> Optional[FileDescr]:
3552 if not get_validation_context().perform_io_checks or value is None:
3553 return value
3555 doc_reader = get_reader(value)
3556 doc_content = doc_reader.read().decode(encoding="utf-8")
3557 if not re.search("#.*[vV]alidation", doc_content):
3558 issue_warning(
3559 "No '# Validation' (sub)section found in {value}.",
3560 value=value,
3561 field="documentation",
3562 )
3564 return value
3566 inputs: NotEmpty[Sequence[InputTensorDescr]]
3567 """Describes the input tensors expected by this model."""
3569 @field_validator("inputs", mode="after")
3570 @classmethod
3571 def _validate_input_axes(
3572 cls, inputs: Sequence[InputTensorDescr]
3573 ) -> Sequence[InputTensorDescr]:
3574 input_size_refs = cls._get_axes_with_independent_size(inputs)
3576 for i, ipt in enumerate(inputs):
3577 valid_independent_refs: Dict[
3578 Tuple[TensorId, AxisId],
3579 Tuple[TensorDescr, AnyAxis, Union[int, ParameterizedSize]],
3580 ] = {
3581 **{
3582 (ipt.id, a.id): (ipt, a, a.size)
3583 for a in ipt.axes
3584 if not isinstance(a, BatchAxis)
3585 and isinstance(a.size, (int, ParameterizedSize))
3586 },
3587 **input_size_refs,
3588 }
3589 for a, ax in enumerate(ipt.axes):
3590 cls._validate_axis(
3591 "inputs",
3592 i=i,
3593 tensor_id=ipt.id,
3594 a=a,
3595 axis=ax,
3596 valid_independent_refs=valid_independent_refs,
3597 )
3598 return inputs
3600 @staticmethod
3601 def _validate_axis(
3602 field_name: str,
3603 i: int,
3604 tensor_id: TensorId,
3605 a: int,
3606 axis: AnyAxis,
3607 valid_independent_refs: Dict[
3608 Tuple[TensorId, AxisId],
3609 Tuple[TensorDescr, AnyAxis, Union[int, ParameterizedSize]],
3610 ],
3611 ):
3612 if isinstance(axis, BatchAxis) or isinstance(
3613 axis.size, (int, ParameterizedSize, DataDependentSize)
3614 ):
3615 return
3616 elif not isinstance(axis.size, SizeReference):
3617 assert_never(axis.size)
3619 # validate axis.size SizeReference
3620 ref = (axis.size.tensor_id, axis.size.axis_id)
3621 if ref not in valid_independent_refs:
3622 raise ValueError(
3623 "Invalid tensor axis reference at"
3624 + f" {field_name}[{i}].axes[{a}].size: {axis.size}."
3625 )
3626 if ref == (tensor_id, axis.id):
3627 raise ValueError(
3628 "Self-referencing not allowed for"
3629 + f" {field_name}[{i}].axes[{a}].size: {axis.size}"
3630 )
3631 if axis.type == "channel":
3632 if valid_independent_refs[ref][1].type != "channel":
3633 raise ValueError(
3634 "A channel axis' size may only reference another fixed size"
3635 + " channel axis."
3636 )
3637 if isinstance(axis.channel_names, str) and "{i}" in axis.channel_names:
3638 ref_size = valid_independent_refs[ref][2]
3639 assert isinstance(ref_size, int), (
3640 "channel axis ref (another channel axis) has to specify fixed"
3641 + " size"
3642 )
3643 generated_channel_names = [
3644 axis.channel_names.format(i=i) for i in range(1, ref_size + 1)
3645 ]
3646 axis.channel_names = generated_channel_names
3648 if (ax_unit := getattr(axis, "unit", None)) != (
3649 ref_unit := getattr(valid_independent_refs[ref][1], "unit", None)
3650 ):
3651 raise ValueError(
3652 "The units of an axis and its reference axis need to match, but"
3653 + f" '{ax_unit}' != '{ref_unit}'."
3654 )
3655 ref_axis = valid_independent_refs[ref][1]
3656 if isinstance(ref_axis, BatchAxis):
3657 raise ValueError(
3658 f"Invalid reference axis '{ref_axis.id}' for {tensor_id}.{axis.id}"
3659 + " (a batch axis is not allowed as reference)."
3660 )
3662 if isinstance(axis, WithHalo):
3663 min_size = axis.size.get_size(axis, ref_axis, n=0)
3664 if (min_size - 2 * axis.halo) < 1:
3665 raise ValueError(
3666 f"axis {axis.id} with minimum size {min_size} is too small for halo"
3667 + f" {axis.halo}."
3668 )
3670 ref_halo = axis.halo * axis.scale / ref_axis.scale
3671 if ref_halo != int(ref_halo):
3672 raise ValueError(
3673 f"Inferred halo for {'.'.join(ref)} is not an integer ({ref_halo} ="
3674 + f" {tensor_id}.{axis.id}.halo {axis.halo}"
3675 + f" * {tensor_id}.{axis.id}.scale {axis.scale}"
3676 + f" / {'.'.join(ref)}.scale {ref_axis.scale})."
3677 )
3679 def validate_input_tensors(
3680 self,
3681 sources: Union[
3682 Sequence[NDArray[Any]], Mapping[TensorId, Optional[NDArray[Any]]]
3683 ],
3684 *,
3685 pad_inputs: Union[bool, Literal["allow"]] = True,
3686 crop_outputs: Union[bool, Literal["allow"]] = True,
3687 ) -> Mapping[TensorId, Optional[NDArray[Any]]]:
3688 """Check if the given input tensors match the model's input tensor descriptions.
3689 This includes checks of tensor shapes and dtypes, but not of the actual values.
3690 """
3691 if not isinstance(sources, collections.abc.Mapping):
3692 sources = {descr.id: tensor for descr, tensor in zip(self.inputs, sources)}
3694 tensors = {
3695 **{descr.id: (descr, sources.get(descr.id)) for descr in self.inputs},
3696 **{ # outputs are required for halo
3697 descr.id: (descr, None) for descr in self.outputs
3698 },
3699 }
3700 validate_tensors(tensors, pad_inputs=pad_inputs, crop_outputs=crop_outputs)
3702 return sources
3704 @model_validator(mode="after")
3705 def _validate_test_tensors(self) -> Self:
3706 if not get_validation_context().perform_io_checks:
3707 return self
3709 test_inputs = {
3710 descr.id: (
3711 descr,
3712 None if descr.test_tensor is None else load_array(descr.test_tensor),
3713 )
3714 for descr in self.inputs
3715 }
3716 test_outputs = {
3717 descr.id: (
3718 descr,
3719 None if descr.test_tensor is None else load_array(descr.test_tensor),
3720 )
3721 for descr in self.outputs
3722 }
3724 validate_tensors(
3725 {**test_inputs, **test_outputs},
3726 tensor_origin="test_tensor",
3727 pad_inputs="allow",
3728 crop_outputs="allow",
3729 )
3731 for rep_tol in self.config.bioimageio.reproducibility_tolerance:
3732 if not rep_tol.absolute_tolerance:
3733 continue
3735 if rep_tol.output_ids:
3736 out_arrays = {
3737 k: v[1] for k, v in test_outputs.items() if k in rep_tol.output_ids
3738 }
3739 else:
3740 out_arrays = {k: v[1] for k, v in test_outputs.items()}
3742 for out_id, array in out_arrays.items():
3743 if array is None:
3744 continue
3746 if rep_tol.absolute_tolerance > (max_test_value := array.max()) * 0.01:
3747 raise ValueError(
3748 "config.bioimageio.reproducibility_tolerance.absolute_tolerance="
3749 + f"{rep_tol.absolute_tolerance} > 0.01*{max_test_value}"
3750 + f" (1% of the maximum value of the test tensor '{out_id}')"
3751 )
3753 return self
3755 @model_validator(mode="after")
3756 def _validate_tensor_references_in_proc_kwargs(self, info: ValidationInfo) -> Self:
3757 ipt_refs = {t.id for t in self.inputs}
3758 missing_refs = [
3759 k["reference_tensor"]
3760 for k in [p.kwargs for ipt in self.inputs for p in ipt.preprocessing]
3761 + [p.kwargs for out in self.outputs for p in out.postprocessing]
3762 if "reference_tensor" in k
3763 and k["reference_tensor"] is not None
3764 and k["reference_tensor"] not in ipt_refs
3765 ]
3767 if missing_refs:
3768 raise ValueError(
3769 f"`reference_tensor`s {missing_refs} not found. Valid input tensor"
3770 + f" references are: {ipt_refs}."
3771 )
3773 return self
3775 name: Annotated[
3776 str,
3777 RestrictCharacters(string.ascii_letters + string.digits + "_+- ()"),
3778 MinLen(5),
3779 MaxLen(128),
3780 warn(MaxLen(64), "Name longer than 64 characters.", INFO),
3781 ]
3782 """A human-readable name of this model.
3783 It should be no longer than 64 characters
3784 and may only contain letter, number, underscore, minus, parentheses and spaces.
3785 We recommend to chose a name that refers to the model's task and image modality.
3786 """
3788 outputs: NotEmpty[Sequence[OutputTensorDescr]]
3789 """Describes the output tensors."""
3791 @field_validator("outputs", mode="after")
3792 @classmethod
3793 def _validate_tensor_ids(
3794 cls, outputs: Sequence[OutputTensorDescr], info: ValidationInfo
3795 ) -> Sequence[OutputTensorDescr]:
3796 tensor_ids = [
3797 t.id for t in info.data.get("inputs", []) + info.data.get("outputs", [])
3798 ]
3799 duplicate_tensor_ids: List[str] = []
3800 seen: Set[str] = set()
3801 for t in tensor_ids:
3802 if t in seen:
3803 duplicate_tensor_ids.append(t)
3805 seen.add(t)
3807 if duplicate_tensor_ids:
3808 raise ValueError(f"Duplicate tensor ids: {duplicate_tensor_ids}")
3810 return outputs
3812 @staticmethod
3813 def _get_axes_with_parameterized_size(
3814 io: Union[Sequence[InputTensorDescr], Sequence[OutputTensorDescr]],
3815 ):
3816 return {
3817 f"{t.id}.{a.id}": (t, a, a.size)
3818 for t in io
3819 for a in t.axes
3820 if not isinstance(a, BatchAxis) and isinstance(a.size, ParameterizedSize)
3821 }
3823 @staticmethod
3824 def _get_axes_with_independent_size(
3825 io: Union[Sequence[InputTensorDescr], Sequence[OutputTensorDescr]],
3826 ):
3827 return {
3828 (t.id, a.id): (t, a, a.size)
3829 for t in io
3830 for a in t.axes
3831 if not isinstance(a, BatchAxis)
3832 and isinstance(a.size, (int, ParameterizedSize))
3833 }
3835 @field_validator("outputs", mode="after")
3836 @classmethod
3837 def _validate_output_axes(
3838 cls, outputs: List[OutputTensorDescr], info: ValidationInfo
3839 ) -> List[OutputTensorDescr]:
3840 input_size_refs = cls._get_axes_with_independent_size(
3841 info.data.get("inputs", [])
3842 )
3843 output_size_refs = cls._get_axes_with_independent_size(outputs)
3845 for i, out in enumerate(outputs):
3846 valid_independent_refs: Dict[
3847 Tuple[TensorId, AxisId],
3848 Tuple[TensorDescr, AnyAxis, Union[int, ParameterizedSize]],
3849 ] = {
3850 **{
3851 (out.id, a.id): (out, a, a.size)
3852 for a in out.axes
3853 if not isinstance(a, BatchAxis)
3854 and isinstance(a.size, (int, ParameterizedSize))
3855 },
3856 **input_size_refs,
3857 **output_size_refs,
3858 }
3859 for a, ax in enumerate(out.axes):
3860 cls._validate_axis(
3861 "outputs",
3862 i,
3863 out.id,
3864 a,
3865 ax,
3866 valid_independent_refs=valid_independent_refs,
3867 )
3869 return outputs
3871 packaged_by: List[Author] = Field(
3872 default_factory=cast(Callable[[], List[Author]], list)
3873 )
3874 """The persons that have packaged and uploaded this model.
3875 Only required if those persons differ from the `authors`."""
3877 parent: Optional[LinkedModel] = None
3878 """The model from which this model is derived, e.g. by fine-tuning the weights."""
3880 @model_validator(mode="after")
3881 def _validate_parent_is_not_self(self) -> Self:
3882 if self.parent is not None and self.parent.id == self.id:
3883 raise ValueError("A model description may not reference itself as parent.")
3885 return self
3887 run_mode: Annotated[
3888 Optional[RunMode],
3889 warn(None, "Run mode '{value}' has limited support across consumer softwares."),
3890 ] = None
3891 """Custom run mode for this model: for more complex prediction procedures like test time
3892 data augmentation that currently cannot be expressed in the specification.
3893 No standard run modes are defined yet."""
3895 timestamp: Datetime = Field(default_factory=Datetime.now)
3896 """Timestamp in [ISO 8601](#https://en.wikipedia.org/wiki/ISO_8601) format
3897 with a few restrictions listed [here](https://docs.python.org/3/library/datetime.html#datetime.datetime.fromisoformat).
3898 (In Python a datetime object is valid, too)."""
3900 training_data: Annotated[
3901 Union[None, LinkedDataset, DatasetDescr, DatasetDescr02],
3902 Field(union_mode="left_to_right"),
3903 ] = None
3904 """The dataset used to train this model"""
3906 weights: Annotated[WeightsDescr, WrapSerializer(package_weights)]
3907 """The weights for this model.
3908 Weights can be given for different formats, but should otherwise be equivalent.
3909 The available weight formats determine which consumers can use this model."""
3911 config: Config = Field(default_factory=Config.model_construct)
3913 @model_validator(mode="after")
3914 def _add_default_cover(self) -> Self:
3915 if not get_validation_context().perform_io_checks or self.covers:
3916 return self
3918 try:
3919 generated_covers = generate_covers(
3920 [
3921 (t, load_array(t.test_tensor))
3922 for t in self.inputs
3923 if t.test_tensor is not None
3924 ],
3925 [
3926 (t, load_array(t.test_tensor))
3927 for t in self.outputs
3928 if t.test_tensor is not None
3929 ],
3930 )
3931 except Exception as e:
3932 issue_warning(
3933 "Failed to generate cover image(s): {e}",
3934 value=self.covers,
3935 msg_context=dict(e=e),
3936 field="covers",
3937 )
3938 else:
3939 self.covers.extend(generated_covers)
3941 return self
3943 def get_input_test_arrays(self) -> List[NDArray[Any]]:
3944 return self._get_test_arrays(self.inputs)
3946 def get_output_test_arrays(self) -> List[NDArray[Any]]:
3947 return self._get_test_arrays(self.outputs)
3949 @staticmethod
3950 def _get_test_arrays(
3951 io_descr: Union[Sequence[InputTensorDescr], Sequence[OutputTensorDescr]],
3952 ):
3953 ts: List[FileDescr] = []
3954 for d in io_descr:
3955 if d.test_tensor is None:
3956 raise ValueError(
3957 f"Failed to get test arrays: description of '{d.id}' is missing a `test_tensor`."
3958 )
3959 ts.append(d.test_tensor)
3961 data = [load_array(t) for t in ts]
3962 assert all(isinstance(d, np.ndarray) for d in data)
3963 return data
3965 @staticmethod
3966 def get_batch_size(tensor_sizes: Mapping[TensorId, Mapping[AxisId, int]]) -> int:
3967 batch_size = 1
3968 tensor_with_batchsize: Optional[TensorId] = None
3969 for tid in tensor_sizes:
3970 for aid, s in tensor_sizes[tid].items():
3971 if aid != BATCH_AXIS_ID or s == 1 or s == batch_size:
3972 continue
3974 if batch_size != 1:
3975 assert tensor_with_batchsize is not None
3976 raise ValueError(
3977 f"batch size mismatch for tensors '{tensor_with_batchsize}' ({batch_size}) and '{tid}' ({s})"
3978 )
3980 batch_size = s
3981 tensor_with_batchsize = tid
3983 return batch_size
3985 def get_output_tensor_sizes(
3986 self, input_sizes: Mapping[TensorId, Mapping[AxisId, int]]
3987 ) -> Dict[TensorId, Dict[AxisId, Union[int, _DataDepSize]]]:
3988 """Returns the tensor output sizes for given **input_sizes**.
3989 Only if **input_sizes** has a valid input shape, the tensor output size is exact.
3990 Otherwise it might be larger than the actual (valid) output"""
3991 batch_size = self.get_batch_size(input_sizes)
3992 ns = self.get_ns(input_sizes)
3994 tensor_sizes = self.get_tensor_sizes(ns, batch_size=batch_size)
3995 return tensor_sizes.outputs
3997 def get_ns(self, input_sizes: Mapping[TensorId, Mapping[AxisId, int]]):
3998 """get parameter `n` for each parameterized axis
3999 such that the valid input size is >= the given input size"""
4000 ret: Dict[Tuple[TensorId, AxisId], ParameterizedSize_N] = {}
4001 axes = {t.id: {a.id: a for a in t.axes} for t in self.inputs}
4002 for tid in input_sizes:
4003 for aid, s in input_sizes[tid].items():
4004 size_descr = axes[tid][aid].size
4005 if isinstance(size_descr, ParameterizedSize):
4006 ret[(tid, aid)] = size_descr.get_n(s)
4007 elif size_descr is None or isinstance(size_descr, (int, SizeReference)):
4008 pass
4009 else:
4010 assert_never(size_descr)
4012 return ret
4014 def get_tensor_sizes(
4015 self, ns: Mapping[Tuple[TensorId, AxisId], ParameterizedSize_N], batch_size: int
4016 ) -> _TensorSizes:
4017 axis_sizes = self.get_axis_sizes(ns, batch_size=batch_size)
4018 return _TensorSizes(
4019 {
4020 t: {
4021 aa: axis_sizes.inputs[(tt, aa)]
4022 for tt, aa in axis_sizes.inputs
4023 if tt == t
4024 }
4025 for t in {tt for tt, _ in axis_sizes.inputs}
4026 },
4027 {
4028 t: {
4029 aa: axis_sizes.outputs[(tt, aa)]
4030 for tt, aa in axis_sizes.outputs
4031 if tt == t
4032 }
4033 for t in {tt for tt, _ in axis_sizes.outputs}
4034 },
4035 )
4037 def get_axis_sizes(
4038 self,
4039 ns: Mapping[Tuple[TensorId, AxisId], ParameterizedSize_N],
4040 batch_size: Optional[int] = None,
4041 *,
4042 max_input_shape: Optional[Mapping[Tuple[TensorId, AxisId], int]] = None,
4043 ) -> _AxisSizes:
4044 """Determine input and output block shape for scale factors **ns**
4045 of parameterized input sizes.
4047 Args:
4048 ns: Scale factor `n` for each axis (keyed by (tensor_id, axis_id))
4049 that is parameterized as `size = min + n * step`.
4050 batch_size: The desired size of the batch dimension.
4051 If given **batch_size** overwrites any batch size present in
4052 **max_input_shape**. Default 1.
4053 max_input_shape: Limits the derived block shapes.
4054 Each axis for which the input size, parameterized by `n`, is larger
4055 than **max_input_shape** is set to the minimal value `n_min` for which
4056 this is still true.
4057 Use this for small input samples or large values of **ns**.
4058 Or simply whenever you know the full input shape.
4060 Returns:
4061 Resolved axis sizes for model inputs and outputs.
4062 """
4063 max_input_shape = max_input_shape or {}
4064 if batch_size is None:
4065 for (_t_id, a_id), s in max_input_shape.items():
4066 if a_id == BATCH_AXIS_ID:
4067 batch_size = s
4068 break
4069 else:
4070 batch_size = 1
4072 all_axes = {
4073 t.id: {a.id: a for a in t.axes} for t in chain(self.inputs, self.outputs)
4074 }
4076 inputs: Dict[Tuple[TensorId, AxisId], int] = {}
4077 outputs: Dict[Tuple[TensorId, AxisId], Union[int, _DataDepSize]] = {}
4079 def get_axis_size(a: Union[InputAxis, OutputAxis]):
4080 if isinstance(a, BatchAxis):
4081 if (t_descr.id, a.id) in ns:
4082 logger.warning(
4083 "Ignoring unexpected size increment factor (n) for batch axis"
4084 + " of tensor '{}'.",
4085 t_descr.id,
4086 )
4087 return batch_size
4088 elif isinstance(a.size, int):
4089 if (t_descr.id, a.id) in ns:
4090 logger.warning(
4091 "Ignoring unexpected size increment factor (n) for fixed size"
4092 + " axis '{}' of tensor '{}'.",
4093 a.id,
4094 t_descr.id,
4095 )
4096 return a.size
4097 elif isinstance(a.size, ParameterizedSize):
4098 if (t_descr.id, a.id) not in ns:
4099 raise ValueError(
4100 "Size increment factor (n) missing for parametrized axis"
4101 + f" '{a.id}' of tensor '{t_descr.id}'."
4102 )
4103 n = ns[(t_descr.id, a.id)]
4104 s_max = max_input_shape.get((t_descr.id, a.id))
4105 if s_max is not None:
4106 n = min(n, a.size.get_n(s_max))
4108 return a.size.get_size(n)
4110 elif isinstance(a.size, SizeReference):
4111 if (t_descr.id, a.id) in ns:
4112 logger.warning(
4113 "Ignoring unexpected size increment factor (n) for axis '{}'"
4114 + " of tensor '{}' with size reference.",
4115 a.id,
4116 t_descr.id,
4117 )
4118 assert not isinstance(a, BatchAxis)
4119 ref_axis = all_axes[a.size.tensor_id][a.size.axis_id]
4120 assert not isinstance(ref_axis, BatchAxis)
4121 ref_key = (a.size.tensor_id, a.size.axis_id)
4122 ref_size = inputs.get(ref_key, outputs.get(ref_key))
4123 assert ref_size is not None, ref_key
4124 assert not isinstance(ref_size, _DataDepSize), ref_key
4125 return a.size.get_size(
4126 axis=a,
4127 ref_axis=ref_axis,
4128 ref_size=ref_size,
4129 )
4130 elif isinstance(a.size, DataDependentSize):
4131 if (t_descr.id, a.id) in ns:
4132 logger.warning(
4133 "Ignoring unexpected increment factor (n) for data dependent"
4134 + " size axis '{}' of tensor '{}'.",
4135 a.id,
4136 t_descr.id,
4137 )
4138 return _DataDepSize(a.size.min, a.size.max)
4139 else:
4140 assert_never(a.size)
4142 # first resolve all , but the `SizeReference` input sizes
4143 for t_descr in self.inputs:
4144 for a in t_descr.axes:
4145 if not isinstance(a.size, SizeReference):
4146 s = get_axis_size(a)
4147 assert not isinstance(s, _DataDepSize)
4148 inputs[t_descr.id, a.id] = s
4150 # resolve all other input axis sizes
4151 for t_descr in self.inputs:
4152 for a in t_descr.axes:
4153 if isinstance(a.size, SizeReference):
4154 s = get_axis_size(a)
4155 assert not isinstance(s, _DataDepSize)
4156 inputs[t_descr.id, a.id] = s
4158 # resolve all output axis sizes
4159 for t_descr in self.outputs:
4160 for a in t_descr.axes:
4161 assert not isinstance(a.size, ParameterizedSize)
4162 s = get_axis_size(a)
4163 outputs[t_descr.id, a.id] = s
4165 return _AxisSizes(inputs=inputs, outputs=outputs)
4167 @model_validator(mode="before")
4168 @classmethod
4169 def _convert(cls, data: Dict[str, Any]) -> Dict[str, Any]:
4170 cls.convert_from_old_format_wo_validation(data)
4171 return data
4173 @classmethod
4174 def convert_from_old_format_wo_validation(cls, data: Dict[str, Any]) -> None:
4175 """Convert metadata following an older format version to this classes' format
4176 without validating the result.
4177 """
4178 if (
4179 data.get("type") == "model"
4180 and isinstance(fv := data.get("format_version"), str)
4181 and fv.count(".") == 2
4182 ):
4183 fv_parts = fv.split(".")
4184 if any(not p.isdigit() for p in fv_parts):
4185 return
4187 fv_tuple = tuple(map(int, fv_parts))
4189 assert cls.implemented_format_version_tuple[0:2] == (0, 5)
4190 if fv_tuple[:2] in ((0, 3), (0, 4)):
4191 m04 = _ModelDescr_v0_4.load(data)
4192 if isinstance(m04, InvalidDescr):
4193 try:
4194 updated = _model_conv.convert_as_dict(
4195 m04 # pyright: ignore[reportArgumentType]
4196 )
4197 except Exception as e:
4198 logger.error(
4199 "Failed to convert from invalid model 0.4 description."
4200 + f"\nerror: {e}"
4201 + "\nProceeding with model 0.5 validation without conversion."
4202 )
4203 updated = None
4204 else:
4205 updated = _model_conv.convert_as_dict(m04)
4207 if updated is not None:
4208 data.clear()
4209 data.update(updated)
4211 elif fv_tuple[:2] == (0, 5):
4212 # bump patch version
4213 data["format_version"] = cls.implemented_format_version
4215 if fv_tuple[:2] in ((0, 3), (0, 4)) or (
4216 fv_tuple[:2] == (0, 5) and fv_tuple[2] < 11
4217 ):
4218 convert_plain_covers_and_docs_and_icon(data)
4221class _ModelConv(Converter[_ModelDescr_v0_4, ModelDescr]):
4222 def _convert(
4223 self, src: _ModelDescr_v0_4, tgt: "type[ModelDescr] | type[dict[str, Any]]"
4224 ) -> "ModelDescr | dict[str, Any]":
4225 name = "".join(
4226 c if c in string.ascii_letters + string.digits + "_+- ()" else " "
4227 for c in src.name
4228 )
4230 def conv_authors(auths: Optional[Sequence[_Author_v0_4]]):
4231 conv = (
4232 _author_conv.convert if TYPE_CHECKING else _author_conv.convert_as_dict
4233 )
4234 return None if auths is None else [conv(a) for a in auths]
4236 if TYPE_CHECKING:
4237 arch_file_conv = _arch_file_conv.convert
4238 arch_lib_conv = _arch_lib_conv.convert
4239 else:
4240 arch_file_conv = _arch_file_conv.convert_as_dict
4241 arch_lib_conv = _arch_lib_conv.convert_as_dict
4243 input_size_refs = {
4244 ipt.name: {
4245 a: s
4246 for a, s in zip(
4247 ipt.axes,
4248 (
4249 ipt.shape.min
4250 if isinstance(ipt.shape, _ParameterizedInputShape_v0_4)
4251 else ipt.shape
4252 ),
4253 )
4254 }
4255 for ipt in src.inputs
4256 if ipt.shape
4257 }
4258 output_size_refs = {
4259 **{
4260 out.name: {a: s for a, s in zip(out.axes, out.shape)}
4261 for out in src.outputs
4262 if not isinstance(out.shape, _ImplicitOutputShape_v0_4)
4263 },
4264 **input_size_refs,
4265 }
4267 return tgt(
4268 attachments=(
4269 []
4270 if src.attachments is None
4271 else [FileDescr(source=f) for f in src.attachments.files]
4272 ),
4273 authors=[_author_conv.convert_as_dict(a) for a in src.authors], # pyright: ignore[reportArgumentType]
4274 cite=[{"text": c.text, "doi": c.doi, "url": c.url} for c in src.cite], # pyright: ignore[reportArgumentType]
4275 config=src.config, # pyright: ignore[reportArgumentType]
4276 covers=[{"source": c} for c in src.covers], # pyright: ignore[reportArgumentType]
4277 description=src.description,
4278 documentation={"source": src.documentation} if src.documentation else None, # pyright: ignore[reportArgumentType]
4279 format_version="0.5.12",
4280 git_repo=src.git_repo, # pyright: ignore[reportArgumentType]
4281 icon={"source": src.icon} if src.icon else None, # pyright: ignore[reportArgumentType]
4282 id=None if src.id is None else ModelId(src.id),
4283 id_emoji=src.id_emoji,
4284 license=src.license, # type: ignore
4285 links=src.links,
4286 maintainers=[_maintainer_conv.convert_as_dict(m) for m in src.maintainers], # pyright: ignore[reportArgumentType]
4287 name=name,
4288 tags=src.tags,
4289 type=src.type,
4290 uploader=src.uploader,
4291 version=src.version,
4292 inputs=[ # pyright: ignore[reportArgumentType]
4293 _input_tensor_conv.convert_as_dict(ipt, tt, st, input_size_refs)
4294 for ipt, tt, st in zip(
4295 src.inputs,
4296 src.test_inputs,
4297 src.sample_inputs or [None] * len(src.test_inputs),
4298 )
4299 ],
4300 outputs=[ # pyright: ignore[reportArgumentType]
4301 _output_tensor_conv.convert_as_dict(out, tt, st, output_size_refs)
4302 for out, tt, st in zip(
4303 src.outputs,
4304 src.test_outputs,
4305 src.sample_outputs or [None] * len(src.test_outputs),
4306 )
4307 ],
4308 parent=(
4309 None
4310 if src.parent is None
4311 else LinkedModel(
4312 id=ModelId(
4313 str(src.parent.id)
4314 + (
4315 ""
4316 if src.parent.version_number is None
4317 else f"/{src.parent.version_number}"
4318 )
4319 )
4320 )
4321 ),
4322 training_data=(
4323 None
4324 if src.training_data is None
4325 else (
4326 LinkedDataset(
4327 id=DatasetId(
4328 str(src.training_data.id)
4329 + (
4330 ""
4331 if src.training_data.version_number is None
4332 else f"/{src.training_data.version_number}"
4333 )
4334 )
4335 )
4336 if isinstance(src.training_data, LinkedDataset02)
4337 else src.training_data
4338 )
4339 ),
4340 packaged_by=[_author_conv.convert_as_dict(a) for a in src.packaged_by], # pyright: ignore[reportArgumentType]
4341 run_mode=src.run_mode,
4342 timestamp=src.timestamp,
4343 weights=(WeightsDescr if TYPE_CHECKING else dict)(
4344 keras_hdf5=(w := src.weights.keras_hdf5)
4345 and (KerasHdf5WeightsDescr if TYPE_CHECKING else dict)(
4346 authors=conv_authors(w.authors),
4347 source=w.source,
4348 tensorflow_version=w.tensorflow_version or Version("1.15"),
4349 parent=w.parent,
4350 ),
4351 onnx=(w := src.weights.onnx)
4352 and (OnnxWeightsDescr if TYPE_CHECKING else dict)(
4353 source=w.source,
4354 authors=conv_authors(w.authors),
4355 parent=w.parent,
4356 opset_version=w.opset_version or 15,
4357 ),
4358 pytorch_state_dict=(w := src.weights.pytorch_state_dict)
4359 and (PytorchStateDictWeightsDescr if TYPE_CHECKING else dict)(
4360 source=w.source,
4361 authors=conv_authors(w.authors),
4362 parent=w.parent,
4363 architecture=(
4364 arch_file_conv(
4365 w.architecture,
4366 w.architecture_sha256,
4367 w.kwargs,
4368 )
4369 if isinstance(w.architecture, _CallableFromFile_v0_4)
4370 else arch_lib_conv(w.architecture, w.kwargs)
4371 ),
4372 pytorch_version=w.pytorch_version or Version("1.10"),
4373 dependencies=(
4374 None
4375 if w.dependencies is None
4376 else (FileDescr if TYPE_CHECKING else dict)(
4377 source=cast(
4378 FileSource,
4379 str(deps := w.dependencies)[
4380 (
4381 len("conda:")
4382 if str(deps).startswith("conda:")
4383 else 0
4384 ) :
4385 ],
4386 )
4387 )
4388 ),
4389 ),
4390 tensorflow_js=(w := src.weights.tensorflow_js)
4391 and (TensorflowJsWeightsDescr if TYPE_CHECKING else dict)(
4392 source=w.source,
4393 authors=conv_authors(w.authors),
4394 parent=w.parent,
4395 tensorflow_version=w.tensorflow_version or Version("1.15"),
4396 ),
4397 tensorflow_saved_model_bundle=(
4398 w := src.weights.tensorflow_saved_model_bundle
4399 )
4400 and (TensorflowSavedModelBundleWeightsDescr if TYPE_CHECKING else dict)(
4401 authors=conv_authors(w.authors),
4402 parent=w.parent,
4403 source=w.source,
4404 tensorflow_version=w.tensorflow_version or Version("1.15"),
4405 dependencies=(
4406 None
4407 if w.dependencies is None
4408 else (FileDescr if TYPE_CHECKING else dict)(
4409 source=cast(
4410 FileSource,
4411 (
4412 str(w.dependencies)[len("conda:") :]
4413 if str(w.dependencies).startswith("conda:")
4414 else str(w.dependencies)
4415 ),
4416 )
4417 )
4418 ),
4419 ),
4420 torchscript=(w := src.weights.torchscript)
4421 and (TorchscriptWeightsDescr if TYPE_CHECKING else dict)(
4422 source=w.source,
4423 authors=conv_authors(w.authors),
4424 parent=w.parent,
4425 pytorch_version=w.pytorch_version or Version("1.10"),
4426 ),
4427 ),
4428 )
4431_model_conv = _ModelConv(_ModelDescr_v0_4, ModelDescr)
4434# create better cover images for 3d data and non-image outputs
4435def generate_covers(
4436 inputs: Sequence[Tuple[InputTensorDescr, NDArray[Any]]],
4437 outputs: Sequence[Tuple[OutputTensorDescr, NDArray[Any]]],
4438) -> List[FileDescr]:
4439 def squeeze(
4440 data: NDArray[Any], axes: Sequence[AnyAxis]
4441 ) -> Tuple[NDArray[Any], List[AnyAxis]]:
4442 """apply numpy.ndarray.squeeze while keeping track of the axis descriptions remaining"""
4443 if data.ndim != len(axes):
4444 raise ValueError(
4445 f"tensor shape {data.shape} does not match described axes"
4446 + f" {[a.id for a in axes]}"
4447 )
4449 axes = [deepcopy(a) for a, s in zip(axes, data.shape) if s != 1]
4450 return data.squeeze(), axes
4452 def normalize(
4453 data: NDArray[Any], axis: Optional[Tuple[int, ...]], eps: float = 1e-7
4454 ) -> NDArray[np.float32]:
4455 data = data.astype("float32")
4456 data -= data.min(axis=axis, keepdims=True)
4457 data /= data.max(axis=axis, keepdims=True) + eps
4458 return data
4460 def to_2d_image(data: NDArray[Any], axes: Sequence[AnyAxis]):
4461 original_shape = data.shape
4462 original_axes = list(axes)
4463 data, axes = squeeze(data, axes)
4465 # take slice fom any batch or index axis if needed
4466 # and convert the first channel axis and take a slice from any additional channel axes
4467 slices: Tuple[slice, ...] = ()
4468 ndim = data.ndim
4469 ndim_need = 3 if any(isinstance(a, ChannelAxis) for a in axes) else 2
4470 has_c_axis = False
4471 for i, a in enumerate(axes):
4472 s = data.shape[i]
4473 assert s > 1
4474 if (
4475 isinstance(a, (BatchAxis, IndexInputAxis, IndexOutputAxis))
4476 and ndim > ndim_need
4477 ):
4478 data = data[slices + (slice(s // 2 - 1, s // 2),)]
4479 ndim -= 1
4480 elif isinstance(a, ChannelAxis):
4481 if has_c_axis:
4482 # second channel axis
4483 data = data[slices + (slice(0, 1),)]
4484 ndim -= 1
4485 else:
4486 has_c_axis = True
4487 if s == 2:
4488 # visualize two channels with cyan and magenta
4489 data = np.concatenate(
4490 [
4491 data[slices + (slice(1, 2),)],
4492 data[slices + (slice(0, 1),)],
4493 (
4494 data[slices + (slice(0, 1),)]
4495 + data[slices + (slice(1, 2),)]
4496 )
4497 / 2, # TODO: take maximum instead?
4498 ],
4499 axis=i,
4500 )
4501 elif data.shape[i] == 3:
4502 pass # visualize 3 channels as RGB
4503 else:
4504 # visualize first 3 channels as RGB
4505 data = data[slices + (slice(3),)]
4507 assert data.shape[i] == 3
4509 slices += (slice(None),)
4511 data, axes = squeeze(data, axes)
4512 assert len(axes) == ndim
4513 # take slice from z axis if needed
4514 slices = ()
4515 if ndim > ndim_need:
4516 for i, a in enumerate(axes):
4517 s = data.shape[i]
4518 if a.id == AxisId("z"):
4519 data = data[slices + (slice(s // 2 - 1, s // 2),)]
4520 data, axes = squeeze(data, axes)
4521 ndim -= 1
4522 break
4524 slices += (slice(None),)
4526 # take slice from any space or time axis
4527 slices = ()
4529 for i, a in enumerate(axes):
4530 if ndim <= ndim_need:
4531 break
4533 s = data.shape[i]
4534 assert s > 1
4535 if isinstance(
4536 a, (SpaceInputAxis, SpaceOutputAxis, TimeInputAxis, TimeOutputAxis)
4537 ):
4538 data = data[slices + (slice(s // 2 - 1, s // 2),)]
4539 ndim -= 1
4541 slices += (slice(None),)
4543 del slices
4544 data, axes = squeeze(data, axes)
4545 assert len(axes) == ndim
4547 if (has_c_axis and ndim != 3) or (not has_c_axis and ndim != 2):
4548 raise ValueError(
4549 f"Failed to construct cover image from shape {original_shape} with axes {[a.id for a in original_axes]}."
4550 )
4552 if not has_c_axis:
4553 assert ndim == 2
4554 data = np.repeat(data[:, :, None], 3, axis=2)
4555 axes.append(ChannelAxis(channel_names=list("RGB")))
4556 ndim += 1
4558 assert ndim == 3
4560 # transpose axis order such that longest axis comes first...
4561 axis_order: List[int] = list(np.argsort(list(data.shape)))
4562 axis_order.reverse()
4563 # ... and channel axis is last
4564 c = [i for i in range(3) if isinstance(axes[i], ChannelAxis)][0]
4565 axis_order.append(axis_order.pop(c))
4566 axes = [axes[ao] for ao in axis_order]
4567 data = data.transpose(axis_order)
4569 # h, w = data.shape[:2]
4570 # if h / w in (1.0 or 2.0):
4571 # pass
4572 # elif h / w < 2:
4573 # TODO: enforce 2:1 or 1:1 aspect ratio for generated cover images
4575 norm_along = (
4576 tuple(i for i, a in enumerate(axes) if a.type in ("space", "time")) or None
4577 )
4578 # normalize the data and map to 8 bit
4579 data = normalize(data, norm_along)
4580 data = (data * 255).astype("uint8")
4582 return data
4584 def create_diagonal_split_image(im0: NDArray[Any], im1: NDArray[Any]):
4585 assert im0.dtype == im1.dtype == np.uint8
4586 assert im0.shape == im1.shape
4587 assert im0.ndim == 3
4588 N, M, C = im0.shape
4589 assert C == 3
4590 out = np.ones((N, M, C), dtype="uint8")
4591 for c in range(C):
4592 outc = np.tril(im0[..., c])
4593 mask = outc == 0
4594 outc[mask] = np.triu(im1[..., c])[mask]
4595 out[..., c] = outc
4597 return out
4599 if not inputs:
4600 raise ValueError("Missing test input tensor for cover generation.")
4602 if not outputs:
4603 raise ValueError("Missing test output tensor for cover generation.")
4605 ipt_descr, ipt = inputs[0]
4606 out_descr, out = outputs[0]
4608 ipt_img = to_2d_image(ipt, ipt_descr.axes)
4609 out_img = to_2d_image(out, out_descr.axes)
4611 cover_folder = Path(mkdtemp())
4612 if ipt_img.shape == out_img.shape:
4613 covers = [cover_folder / "cover.png"]
4614 imwrite(covers[0], create_diagonal_split_image(ipt_img, out_img))
4615 else:
4616 covers = [cover_folder / "input.png", cover_folder / "output.png"]
4617 imwrite(covers[0], ipt_img)
4618 imwrite(covers[1], out_img)
4620 return [FileDescr(source=c) for c in covers]