Coverage for src/bioimageio/spec/model/v0_5.py: 71%
1692 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 19:19 +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 TypeVar,
28 Union,
29 cast,
30 overload,
31)
33import numpy as np
34from annotated_types import Ge, Gt, Interval, MaxLen, MinLen, Predicate
35from imageio.v3 import imread, imwrite # pyright: ignore[reportUnknownVariableType]
36from loguru import logger
37from numpy.typing import NDArray
38from pydantic import (
39 AfterValidator,
40 Discriminator,
41 Field,
42 RootModel,
43 SerializationInfo,
44 SerializerFunctionWrapHandler,
45 StrictInt,
46 Tag,
47 ValidationInfo,
48 WrapSerializer,
49 field_validator,
50 model_serializer,
51 model_validator,
52)
53from typing_extensions import Annotated, Self, TypeAlias, assert_never, get_args
55from .._internal.common_nodes import (
56 InvalidDescr,
57 KwargsNode,
58 Node,
59 NodeWithExplicitlySetFields,
60)
61from .._internal.constants import DTYPE_LIMITS
62from .._internal.field_warning import issue_warning, warn
63from .._internal.io import BioimageioYamlContent as BioimageioYamlContent
64from .._internal.io import FileDescr as FileDescr
65from .._internal.io import (
66 FileSource,
67 WithSuffix,
68 YamlValue,
69 extract_file_name,
70 get_reader,
71 wo_special_file_name,
72)
73from .._internal.io_basics import Sha256 as Sha256
74from .._internal.io_packaging import (
75 FileDescr_package,
76 package_file_descr_serializer,
77)
78from .._internal.io_utils import load_array, open_bioimageio_yaml
79from .._internal.node_converter import Converter
80from .._internal.type_guards import is_dict, is_sequence
81from .._internal.types import (
82 FAIR,
83 AbsoluteTolerance,
84 LowerCaseIdentifier,
85 LowerCaseIdentifierAnno,
86 MismatchedElementsPerMillion,
87 RelativeTolerance,
88 validate_identifier,
89 validate_is_not_keyword,
90)
91from .._internal.types import Datetime as Datetime
92from .._internal.types import Identifier as Identifier
93from .._internal.types import NotEmpty as NotEmpty
94from .._internal.types import SiUnit as SiUnit
95from .._internal.url import HttpUrl as HttpUrl
96from .._internal.utils import try_all_raise_last
97from .._internal.validation_context import get_validation_context
98from .._internal.validator_annotations import RestrictCharacters
99from .._internal.version_type import Version as Version
100from .._internal.warning_levels import INFO
101from ..dataset.v0_2 import DatasetDescr as DatasetDescr02
102from ..dataset.v0_2 import LinkedDataset as LinkedDataset02
103from ..dataset.v0_3 import DatasetDescr as DatasetDescr
104from ..dataset.v0_3 import DatasetId as DatasetId
105from ..dataset.v0_3 import LinkedDataset as LinkedDataset
106from ..dataset.v0_3 import Uploader as Uploader
107from ..generic._v0_3_converter import convert_plain_covers_and_docs_and_icon
108from ..generic.v0_3 import (
109 VALID_COVER_IMAGE_EXTENSIONS as VALID_COVER_IMAGE_EXTENSIONS,
110)
111from ..generic.v0_3 import Author as Author
112from ..generic.v0_3 import BadgeDescr as BadgeDescr
113from ..generic.v0_3 import CiteEntry as CiteEntry
114from ..generic.v0_3 import DeprecatedLicenseId as DeprecatedLicenseId
115from ..generic.v0_3 import Doi as Doi
116from ..generic.v0_3 import (
117 FileDescr_documentation,
118 GenericModelDescrBase,
119 LinkedResourceBase,
120 _author_conv, # pyright: ignore[reportPrivateUsage]
121 _maintainer_conv, # pyright: ignore[reportPrivateUsage]
122)
123from ..generic.v0_3 import LicenseId as LicenseId
124from ..generic.v0_3 import LinkedResource as LinkedResource
125from ..generic.v0_3 import Maintainer as Maintainer
126from ..generic.v0_3 import OrcidId as OrcidId
127from ..generic.v0_3 import RelativeFilePath as RelativeFilePath
128from ..generic.v0_3 import ResourceId as ResourceId
129from .v0_4 import Author as _Author_v0_4
130from .v0_4 import BinarizeDescr as _BinarizeDescr_v0_4
131from .v0_4 import CallableFromDepencency as CallableFromDepencency
132from .v0_4 import CallableFromDepencency as _CallableFromDepencency_v0_4
133from .v0_4 import CallableFromFile as _CallableFromFile_v0_4
134from .v0_4 import ClipDescr as _ClipDescr_v0_4
135from .v0_4 import ImplicitOutputShape as _ImplicitOutputShape_v0_4
136from .v0_4 import InputTensorDescr as _InputTensorDescr_v0_4
137from .v0_4 import KnownRunMode as KnownRunMode
138from .v0_4 import ModelDescr as _ModelDescr04
139from .v0_4 import ModelDescr as _ModelDescr_v0_4
140from .v0_4 import OutputTensorDescr as _OutputTensorDescr_v0_4
141from .v0_4 import ParameterizedInputShape as _ParameterizedInputShape_v0_4
142from .v0_4 import PostprocessingDescr as _PostprocessingDescr_v0_4
143from .v0_4 import PreprocessingDescr as _PreprocessingDescr_v0_4
144from .v0_4 import RunMode as RunMode
145from .v0_4 import ScaleLinearDescr as _ScaleLinearDescr_v0_4
146from .v0_4 import ScaleMeanVarianceDescr as _ScaleMeanVarianceDescr_v0_4
147from .v0_4 import ScaleRangeDescr as _ScaleRangeDescr_v0_4
148from .v0_4 import SigmoidDescr as _SigmoidDescr_v0_4
149from .v0_4 import TensorName as _TensorName_v0_4
150from .v0_4 import ZeroMeanUnitVarianceDescr as _ZeroMeanUnitVarianceDescr_v0_4
151from .v0_4 import package_weights
153SpaceUnit = Literal[
154 "attometer",
155 "angstrom",
156 "centimeter",
157 "decimeter",
158 "exameter",
159 "femtometer",
160 "foot",
161 "gigameter",
162 "hectometer",
163 "inch",
164 "kilometer",
165 "megameter",
166 "meter",
167 "micrometer",
168 "mile",
169 "millimeter",
170 "nanometer",
171 "parsec",
172 "petameter",
173 "picometer",
174 "terameter",
175 "yard",
176 "yoctometer",
177 "yottameter",
178 "zeptometer",
179 "zettameter",
180]
181"""Space unit compatible to the [OME-Zarr axes specification 0.5](https://ngff.openmicroscopy.org/0.5/#axes-md)"""
183TimeUnit = Literal[
184 "attosecond",
185 "centisecond",
186 "day",
187 "decisecond",
188 "exasecond",
189 "femtosecond",
190 "gigasecond",
191 "hectosecond",
192 "hour",
193 "kilosecond",
194 "megasecond",
195 "microsecond",
196 "millisecond",
197 "minute",
198 "nanosecond",
199 "petasecond",
200 "picosecond",
201 "second",
202 "terasecond",
203 "yoctosecond",
204 "yottasecond",
205 "zeptosecond",
206 "zettasecond",
207]
208"""Time unit compatible to the [OME-Zarr axes specification 0.5](https://ngff.openmicroscopy.org/0.5/#axes-md)"""
210AxisType = Literal["batch", "channel", "index", "time", "space"]
212_AXIS_TYPE_MAP: Mapping[str, AxisType] = {
213 "b": "batch",
214 "t": "time",
215 "i": "index",
216 "c": "channel",
217 "x": "space",
218 "y": "space",
219 "z": "space",
220}
222_AXIS_ID_MAP = {
223 "b": "batch",
224 "t": "time",
225 "i": "index",
226 "c": "channel",
227 "s": "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 b = str(a).lower()
249 normalized = _AXIS_ID_MAP.get(b, b)
250 if a != normalized:
251 logger.opt(depth=3).debug(
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 NotEmpty[str],
261 AfterValidator(_normalize_axis_id),
262 MaxLen(16),
263 AfterValidator(validate_identifier),
264 AfterValidator(validate_is_not_keyword),
265 ]
266 ]
269def _is_batch(a: str) -> bool:
270 return str(a) == "batch"
273def _is_not_batch(a: str) -> bool:
274 return not _is_batch(a)
277NonBatchAxisId = Annotated[AxisId, Predicate(_is_not_batch)]
279PreprocessingId = Literal[
280 "binarize",
281 "clip",
282 "ensure_dtype",
283 "fixed_zero_mean_unit_variance",
284 "scale_linear",
285 "scale_range",
286 "sigmoid",
287 "softmax",
288]
289PostprocessingId = Literal[
290 "binarize",
291 "clip",
292 "custom",
293 "ensure_dtype",
294 "fixed_zero_mean_unit_variance",
295 "scale_linear",
296 "scale_mean_variance",
297 "scale_range",
298 "sigmoid",
299 "softmax",
300 "zero_mean_unit_variance",
301]
304SAME_AS_TYPE = "<same as type>"
307ParameterizedSize_N: TypeAlias = int
308"""
309Annotates an integer to calculate a concrete axis size from a `ParameterizedSize`.
310"""
313class ParameterizedSize(Node):
314 """Describes a range of valid tensor axis sizes as `size = min + n*step`.
316 - **min** and **step** are given by the model description.
317 - All blocksize paramters n = 0,1,2,... yield a valid `size`.
318 - A greater blocksize paramter n = 0,1,2,... results in a greater **size**.
319 This allows to adjust the axis size more generically.
320 """
322 N: ClassVar[type[int]] = ParameterizedSize_N
323 """Positive integer to parameterize this axis"""
325 min: Annotated[int, Gt(0)]
326 step: Annotated[int, Gt(0)]
328 def validate_size(self, size: int, msg_prefix: str = "") -> int:
329 if size < self.min:
330 raise ValueError(
331 f"{msg_prefix}size {size} < {self.min} (minimum axis size)"
332 )
333 if (size - self.min) % self.step != 0:
334 raise ValueError(
335 f"{msg_prefix}size {size} is not parameterized by `min + n*step` ="
336 + f" `{self.min} + n*{self.step}`"
337 )
339 return size
341 def get_size(self, n: ParameterizedSize_N) -> int:
342 return self.min + self.step * n
344 def get_n(self, s: int) -> ParameterizedSize_N:
345 """return smallest n parameterizing a size greater or equal than `s`"""
346 return ceil((s - self.min) / self.step)
349class DataDependentSize(Node):
350 min: Annotated[int, Gt(0)] = 1
351 max: Annotated[int | None, Gt(1)] = None
353 @model_validator(mode="after")
354 def _validate_max_gt_min(self):
355 if self.max is not None and self.min >= self.max:
356 raise ValueError(f"expected `min` < `max`, but got {self.min}, {self.max}")
358 return self
360 def validate_size(self, size: int, msg_prefix: str = "") -> int:
361 if size < self.min:
362 raise ValueError(f"{msg_prefix}size {size} < {self.min}")
364 if self.max is not None and size > self.max:
365 raise ValueError(f"{msg_prefix}size {size} > {self.max}")
367 return size
370class SizeReference(Node):
371 """A tensor axis size (extent in pixels/frames) defined in relation to a reference axis.
373 `axis.size = reference.size * reference.scale / axis.scale + offset`
375 Note:
376 1. The axis and the referenced axis need to have the same unit (or no unit).
377 2. Batch axes may not be referenced.
378 3. Fractions are rounded down.
379 4. If the reference axis is `concatenable` the referencing axis is assumed to be
380 `concatenable` as well with the same block order.
382 Example:
383 An unisotropic input image of w*h=100*49 pixels depicts a phsical space of 200*196mm².
384 Let's assume that we want to express the image height h in relation to its width w
385 instead of only accepting input images of exactly 100*49 pixels
386 (for example to express a range of valid image shapes by parametrizing w, see `ParameterizedSize`).
388 >>> w = SpaceInputAxis(id=AxisId("w"), size=100, unit="millimeter", scale=2)
389 >>> h = SpaceInputAxis(
390 ... id=AxisId("h"),
391 ... size=SizeReference(tensor_id=TensorId("input"), axis_id=AxisId("w"), offset=-1),
392 ... unit="millimeter",
393 ... scale=4,
394 ... )
395 >>> print(h.size.get_size(h, w))
396 49
398 ⇒ h = w * w.scale / h.scale + offset = 100 * 2mm / 4mm - 1 = 49
399 """
401 tensor_id: TensorId
402 """tensor id of the reference axis"""
404 axis_id: AxisId
405 """axis id of the reference axis"""
407 offset: StrictInt = 0
409 def get_size(
410 self,
411 axis: ChannelAxis
412 | IndexInputAxis
413 | IndexOutputAxis
414 | TimeInputAxis
415 | SpaceInputAxis
416 | TimeOutputAxis
417 | TimeOutputAxisWithHalo
418 | SpaceOutputAxis
419 | SpaceOutputAxisWithHalo,
420 ref_axis: ChannelAxis
421 | IndexInputAxis
422 | IndexOutputAxis
423 | TimeInputAxis
424 | SpaceInputAxis
425 | TimeOutputAxis
426 | TimeOutputAxisWithHalo
427 | SpaceOutputAxis
428 | SpaceOutputAxisWithHalo,
429 n: ParameterizedSize_N = 0,
430 ref_size: int | None = None,
431 ):
432 """Compute the concrete size for a given axis and its reference axis.
434 Args:
435 axis: The axis this [SizeReference][] is the size of.
436 ref_axis: The reference axis to compute the size from.
437 n: If the **ref_axis** is parameterized (of type `ParameterizedSize`)
438 and no fixed **ref_size** is given,
439 **n** is used to compute the size of the parameterized **ref_axis**.
440 ref_size: Overwrite the reference size instead of deriving it from
441 **ref_axis**
442 (**ref_axis.scale** is still used; any given **n** is ignored).
443 """
444 assert axis.size == self, (
445 "Given `axis.size` is not defined by this `SizeReference`"
446 )
448 assert ref_axis.id == self.axis_id, (
449 f"Expected `ref_axis.id` to be {self.axis_id}, but got {ref_axis.id}."
450 )
452 assert axis.unit == ref_axis.unit, (
453 "`SizeReference` requires `axis` and `ref_axis` to have the same `unit`,"
454 f" but {axis.unit}!={ref_axis.unit}"
455 )
456 if ref_size is None:
457 if isinstance(ref_axis.size, (int, float)):
458 ref_size = ref_axis.size
459 elif isinstance(ref_axis.size, ParameterizedSize):
460 ref_size = ref_axis.size.get_size(n)
461 elif isinstance(ref_axis.size, DataDependentSize):
462 raise ValueError(
463 "Reference axis referenced in `SizeReference` may not be a `DataDependentSize`."
464 )
465 elif isinstance(ref_axis.size, SizeReference):
466 raise ValueError(
467 "Reference axis referenced in `SizeReference` may not be sized by a"
468 + " `SizeReference` itself."
469 )
470 else:
471 assert_never(ref_axis.size)
473 return int(ref_size * ref_axis.scale / axis.scale + self.offset)
475 @staticmethod
476 def _get_unit(
477 axis: ChannelAxis
478 | IndexInputAxis
479 | IndexOutputAxis
480 | TimeInputAxis
481 | SpaceInputAxis
482 | TimeOutputAxis
483 | TimeOutputAxisWithHalo
484 | SpaceOutputAxis
485 | SpaceOutputAxisWithHalo,
486 ):
487 return axis.unit
490class AxisBase(NodeWithExplicitlySetFields):
491 id: AxisId
492 """An axis id unique across all axes of one tensor."""
494 description: Annotated[str, MaxLen(128)] = ""
495 """A short description of this axis beyond its type and id."""
498class WithHalo(Node):
499 halo: Annotated[int, Ge(1)]
500 """The halo should be cropped from the output tensor to avoid boundary effects.
501 It is to be cropped from both sides, i.e. `size_after_crop = size - 2 * halo`.
502 To document a halo that is already cropped by the model use `size.offset` instead."""
504 size: Annotated[
505 SizeReference,
506 Field(examples=[{"tensor_id": "t", "axis_id": "a", "offset": 5}]),
507 ]
508 """reference to another axis with an optional offset (see [SizeReference][])"""
511BATCH_AXIS_ID = AxisId("batch")
512CHANNEL_AXIS_ID = AxisId("channel")
513DEFAULT_SPACE_AXIS_ID = AxisId("x")
514DEFAULT_INDEX_AXIS_ID = AxisId("index")
515DEFAULT_TIME_AXIS_ID = AxisId("time")
518class BatchAxis(AxisBase):
519 implemented_type: ClassVar[Literal["batch"]] = "batch"
520 if TYPE_CHECKING:
521 type: Literal["batch"] = "batch"
522 else:
523 type: Literal["batch"]
525 id: Annotated[AxisId, Predicate(_is_batch)] = BATCH_AXIS_ID
526 size: Literal[1] | None = None
527 """The batch size may be fixed to 1,
528 otherwise (the default) it may be chosen arbitrarily depending on available memory"""
530 @property
531 def scale(self):
532 return 1.0
534 @property
535 def concatenable(self):
536 return True
538 @property
539 def unit(self):
540 return None
543class ChannelAxis(AxisBase):
544 implemented_type: ClassVar[Literal["channel"]] = "channel"
545 if TYPE_CHECKING:
546 type: Literal["channel"] = "channel"
547 else:
548 type: Literal["channel"]
550 id: NonBatchAxisId = CHANNEL_AXIS_ID
552 channel_names: NotEmpty[list[str]]
553 """Name/label for each channel. The number of channels is given by `len(channel_names)`."""
555 @property
556 def size(self) -> int:
557 return len(self.channel_names)
559 @property
560 def concatenable(self):
561 return False
563 @property
564 def scale(self) -> float:
565 return 1.0
567 @property
568 def unit(self):
569 return None
572class _WithInputAxisSize(Node):
573 size: Annotated[
574 Annotated[int, Gt(0)] | ParameterizedSize | SizeReference,
575 Field(
576 examples=[
577 10,
578 ParameterizedSize(min=32, step=16).model_dump(mode="json"),
579 {"tensor_id": "t", "axis_id": "a", "offset": 5},
580 ]
581 ),
582 ]
583 """The size/length of this axis can be specified as
584 - fixed integer
585 - parameterized series of valid sizes ([ParameterizedSize][])
586 - reference to another axis with an optional offset ([SizeReference][])
587 """
590class IndexAxisBase(AxisBase):
591 implemented_type: ClassVar[Literal["index"]] = "index"
592 if TYPE_CHECKING:
593 type: Literal["index"] = "index"
594 else:
595 type: Literal["index"]
597 id: NonBatchAxisId = DEFAULT_INDEX_AXIS_ID
599 @property
600 def scale(self) -> float:
601 return 1.0
603 @property
604 def unit(self):
605 return None
608class IndexInputAxis(IndexAxisBase, _WithInputAxisSize):
609 concatenable: bool = False
610 """If a model has a `concatenable` input axis, it can be processed blockwise,
611 splitting a longer sample axis into blocks matching its input tensor description.
612 Output axes are concatenable if they have a [SizeReference][] to a concatenable
613 input axis.
614 """
617class IndexOutputAxis(IndexAxisBase):
618 size: Annotated[
619 Annotated[int, Gt(0)] | SizeReference | DataDependentSize,
620 Field(examples=[10, {"tensor_id": "t", "axis_id": "a", "offset": 5}]),
621 ]
622 """The size/length of this axis can be specified as
623 - fixed integer
624 - reference to another axis with an optional offset ([SizeReference][])
625 - data dependent size using [DataDependentSize][] (size is only known after model inference)
626 """
629class TimeAxisBase(AxisBase):
630 implemented_type: ClassVar[Literal["time"]] = "time"
631 if TYPE_CHECKING:
632 type: Literal["time"] = "time"
633 else:
634 type: Literal["time"]
636 id: NonBatchAxisId = DEFAULT_TIME_AXIS_ID
637 unit: TimeUnit | None = None
638 scale: Annotated[float, Gt(0)] = 1.0
641class TimeInputAxis(TimeAxisBase, _WithInputAxisSize):
642 concatenable: bool = False
643 """If a model has a `concatenable` input axis, it can be processed blockwise,
644 splitting a longer sample axis into blocks matching its input tensor description.
645 Output axes are concatenable if they have a [SizeReference][] to a concatenable
646 input axis.
647 """
650class SpaceAxisBase(AxisBase):
651 implemented_type: ClassVar[Literal["space"]] = "space"
652 if TYPE_CHECKING:
653 type: Literal["space"] = "space"
654 else:
655 type: Literal["space"]
657 id: Annotated[NonBatchAxisId, Field(examples=["x", "y", "z"])] = (
658 DEFAULT_SPACE_AXIS_ID
659 )
660 unit: SpaceUnit | None = None
661 scale: Annotated[float, Gt(0)] = 1.0
664class SpaceInputAxis(SpaceAxisBase, _WithInputAxisSize):
665 concatenable: bool = False
666 """If a model has a `concatenable` input axis, it can be processed blockwise,
667 splitting a longer sample axis into blocks matching its input tensor description.
668 Output axes are concatenable if they have a [SizeReference][] to a concatenable
669 input axis.
670 """
673INPUT_AXIS_TYPES = (
674 BatchAxis,
675 ChannelAxis,
676 IndexInputAxis,
677 TimeInputAxis,
678 SpaceInputAxis,
679)
680"""intended for isinstance comparisons in py<3.10"""
682_InputAxisUnion = Union[
683 BatchAxis, ChannelAxis, IndexInputAxis, TimeInputAxis, SpaceInputAxis
684]
685InputAxis = Annotated[_InputAxisUnion, Discriminator("type")]
688class _WithOutputAxisSize(Node):
689 size: Annotated[
690 Annotated[int, Gt(0)] | SizeReference,
691 Field(examples=[10, {"tensor_id": "t", "axis_id": "a", "offset": 5}]),
692 ]
693 """The size/length of this axis can be specified as
694 - fixed integer
695 - reference to another axis with an optional offset (see [SizeReference][])
696 """
699class TimeOutputAxis(TimeAxisBase, _WithOutputAxisSize):
700 pass
703class TimeOutputAxisWithHalo(TimeAxisBase, WithHalo):
704 pass
707def _get_halo_axis_discriminator_value(v: Any) -> Literal["with_halo", "wo_halo"]:
708 if isinstance(v, dict):
709 return "with_halo" if "halo" in v else "wo_halo"
710 else:
711 return "with_halo" if hasattr(v, "halo") else "wo_halo"
714_TimeOutputAxisUnion = Annotated[
715 Union[
716 Annotated[TimeOutputAxis, Tag("wo_halo")],
717 Annotated[TimeOutputAxisWithHalo, Tag("with_halo")],
718 ],
719 Discriminator(_get_halo_axis_discriminator_value),
720]
723class SpaceOutputAxis(SpaceAxisBase, _WithOutputAxisSize):
724 pass
727class SpaceOutputAxisWithHalo(SpaceAxisBase, WithHalo):
728 pass
731_SpaceOutputAxisUnion = Annotated[
732 Union[
733 Annotated[SpaceOutputAxis, Tag("wo_halo")],
734 Annotated[SpaceOutputAxisWithHalo, Tag("with_halo")],
735 ],
736 Discriminator(_get_halo_axis_discriminator_value),
737]
740_OutputAxisUnion = Union[
741 BatchAxis, ChannelAxis, IndexOutputAxis, _TimeOutputAxisUnion, _SpaceOutputAxisUnion
742]
743OutputAxis = Annotated[_OutputAxisUnion, Discriminator("type")]
745OUTPUT_AXIS_TYPES = (
746 BatchAxis,
747 ChannelAxis,
748 IndexOutputAxis,
749 TimeOutputAxis,
750 TimeOutputAxisWithHalo,
751 SpaceOutputAxis,
752 SpaceOutputAxisWithHalo,
753)
754"""intended for isinstance comparisons in py<3.10"""
757AnyAxis = Union[InputAxis, OutputAxis]
759ANY_AXIS_TYPES = INPUT_AXIS_TYPES + OUTPUT_AXIS_TYPES
760"""intended for isinstance comparisons in py<3.10"""
762TVs = Union[
763 NotEmpty[List[int]],
764 NotEmpty[List[float]],
765 NotEmpty[List[bool]],
766 NotEmpty[List[str]],
767]
770NominalOrOrdinalDType = Literal[
771 "float32",
772 "float64",
773 "uint8",
774 "int8",
775 "uint16",
776 "int16",
777 "uint32",
778 "int32",
779 "uint64",
780 "int64",
781 "bool",
782]
785class NominalOrOrdinalDataDescr(Node):
786 values: TVs
787 """A fixed set of nominal or an ascending sequence of ordinal values.
788 In this case `data.type` is required to be an unsigend integer type, e.g. 'uint8'.
789 String `values` are interpreted as labels for tensor values 0, ..., N.
790 Note: as YAML 1.2 does not natively support a "set" datatype,
791 nominal values should be given as a sequence (aka list/array) as well.
792 """
794 type: Annotated[
795 NominalOrOrdinalDType,
796 Field(
797 examples=[
798 "float32",
799 "uint8",
800 "uint16",
801 "int64",
802 "bool",
803 ],
804 ),
805 ] = "uint8"
807 @model_validator(mode="after")
808 def _validate_values_match_type(
809 self,
810 ) -> Self:
811 incompatible: list[Any] = []
812 for v in self.values:
813 if self.type == "bool":
814 if not isinstance(v, bool):
815 incompatible.append(v)
816 elif self.type in DTYPE_LIMITS:
817 if (
818 isinstance(v, (int, float))
819 and (
820 v < DTYPE_LIMITS[self.type].min
821 or v > DTYPE_LIMITS[self.type].max
822 )
823 or (isinstance(v, str) and "uint" not in self.type)
824 or (isinstance(v, float) and "int" in self.type)
825 ):
826 incompatible.append(v)
827 else:
828 incompatible.append(v)
830 if len(incompatible) == 5:
831 incompatible.append("...")
832 break
834 if incompatible:
835 raise ValueError(
836 f"data type '{self.type}' incompatible with values {incompatible}"
837 )
839 return self
841 unit: Literal["arbitrary unit"] | SiUnit | None = None
843 @property
844 def range(self):
845 if isinstance(self.values[0], str):
846 return 0, len(self.values) - 1
847 else:
848 return min(self.values), max(self.values)
851IntervalOrRatioDType = Literal[
852 "float32",
853 "float64",
854 "uint8",
855 "int8",
856 "uint16",
857 "int16",
858 "uint32",
859 "int32",
860 "uint64",
861 "int64",
862]
865class IntervalOrRatioDataDescr(Node):
866 type: Annotated[ # TODO: rename to dtype
867 IntervalOrRatioDType,
868 Field(
869 examples=["float32", "float64", "uint8", "uint16"],
870 ),
871 ] = "float32"
872 range: tuple[float | None, float | None] = (
873 None,
874 None,
875 )
876 """Tuple `(minimum, maximum)` specifying the allowed range of the data in this tensor.
877 `None` corresponds to min/max of what can be expressed by **type**."""
878 unit: Literal["arbitrary unit"] | SiUnit = "arbitrary unit"
879 scale: float = 1.0
880 """Scale for data on an interval (or ratio) scale."""
881 offset: float | None = None
882 """Offset for data on a ratio scale."""
884 @model_validator(mode="before")
885 def _replace_inf(cls, data: Any):
886 if is_dict(data) and "range" in data and is_sequence(data["range"]):
887 forbidden = (
888 "inf",
889 "-inf",
890 ".inf",
891 "-.inf",
892 float("inf"),
893 float("-inf"),
894 )
895 if any(v in forbidden for v in data["range"]):
896 issue_warning("replaced 'inf' value", value=data["range"])
898 data["range"] = tuple(
899 (None if v in forbidden else v) for v in data["range"]
900 )
902 return data
905TensorDataDescr = Union[NominalOrOrdinalDataDescr, IntervalOrRatioDataDescr]
908class BinarizeKwargs(KwargsNode):
909 """key word arguments for [BinarizeDescr][]"""
911 threshold: float
912 """The fixed threshold"""
915class BinarizeAlongAxisKwargs(KwargsNode):
916 """key word arguments for [BinarizeDescr][]"""
918 threshold: NotEmpty[list[float]]
919 """The fixed threshold values along `axis`"""
921 axis: Annotated[NonBatchAxisId, Field(examples=["channel"])]
922 """The `threshold` axis"""
925class BinarizeDescr(NodeWithExplicitlySetFields):
926 """Binarize the tensor with a fixed threshold.
928 Values above [BinarizeKwargs.threshold][]/[BinarizeAlongAxisKwargs.threshold][]
929 will be set to one, values below the threshold to zero.
931 Examples:
932 - in YAML
933 ```yaml
934 postprocessing:
935 - id: binarize
936 kwargs:
937 axis: 'channel'
938 threshold: [0.25, 0.5, 0.75]
939 ```
940 - in Python:
942 >>> postprocessing = [BinarizeDescr(
943 ... kwargs=BinarizeAlongAxisKwargs(
944 ... axis=AxisId('channel'),
945 ... threshold=[0.25, 0.5, 0.75],
946 ... )
947 ... )]
948 """
950 implemented_id: ClassVar[Literal["binarize"]] = "binarize"
951 if TYPE_CHECKING:
952 id: Literal["binarize"] = "binarize"
953 else:
954 id: Literal["binarize"]
955 kwargs: BinarizeKwargs | BinarizeAlongAxisKwargs
958class ClipKwargs(KwargsNode):
959 """key word arguments for [ClipDescr][]"""
961 min: float | None = None
962 """Minimum value for clipping.
964 Exclusive with [min_percentile][]
965 """
966 min_percentile: Annotated[float, Interval(ge=0, lt=100)] | None = None
967 """Minimum percentile for clipping.
969 Exclusive with [min][].
971 In range [0, 100).
972 """
974 max: float | None = None
975 """Maximum value for clipping.
977 Exclusive with `max_percentile`.
978 """
979 max_percentile: Annotated[float, Interval(gt=1, le=100)] | None = None
980 """Maximum percentile for clipping.
982 Exclusive with `max`.
984 In range (1, 100].
985 """
987 axes: Annotated[Sequence[AxisId] | None, Field(examples=[("batch", "x", "y")])] = (
988 None
989 )
990 """The subset of axes to determine percentiles jointly,
992 i.e. axes to reduce to compute min/max from `min_percentile`/`max_percentile`.
993 For example to clip 'batch', 'x' and 'y' jointly in a tensor ('batch', 'channel', 'y', 'x')
994 resulting in a tensor of equal shape with clipped values per channel, specify `axes=('batch', 'x', 'y')`.
995 To clip samples independently, leave out the 'batch' axis.
997 Only valid if `min_percentile` and/or `max_percentile` are set.
999 Default: Compute percentiles over all axes jointly."""
1001 @model_validator(mode="after")
1002 def _validate(self) -> Self:
1003 if (self.min is not None) and (self.min_percentile is not None):
1004 raise ValueError(
1005 "Only one of `min` and `min_percentile` may be set, not both."
1006 )
1007 if (self.max is not None) and (self.max_percentile is not None):
1008 raise ValueError(
1009 "Only one of `max` and `max_percentile` may be set, not both."
1010 )
1011 if (
1012 self.min is None
1013 and self.min_percentile is None
1014 and self.max is None
1015 and self.max_percentile is None
1016 ):
1017 raise ValueError(
1018 "At least one of `min`, `min_percentile`, `max`, or `max_percentile` must be set."
1019 )
1021 if (
1022 self.axes is not None
1023 and self.min_percentile is None
1024 and self.max_percentile is None
1025 ):
1026 raise ValueError(
1027 "If `axes` is set, at least one of `min_percentile` or `max_percentile` must be set."
1028 )
1030 return self
1033class ClipDescr(NodeWithExplicitlySetFields):
1034 """Set tensor values below min to min and above max to max.
1036 See `ScaleRangeDescr` for examples.
1037 """
1039 implemented_id: ClassVar[Literal["clip"]] = "clip"
1040 if TYPE_CHECKING:
1041 id: Literal["clip"] = "clip"
1042 else:
1043 id: Literal["clip"]
1045 kwargs: ClipKwargs
1048class EnsureDtypeKwargs(KwargsNode):
1049 """key word arguments for [EnsureDtypeDescr][]"""
1051 dtype: Literal[
1052 "float32",
1053 "float64",
1054 "uint8",
1055 "int8",
1056 "uint16",
1057 "int16",
1058 "uint32",
1059 "int32",
1060 "uint64",
1061 "int64",
1062 "bool",
1063 ]
1066class EnsureDtypeDescr(NodeWithExplicitlySetFields):
1067 """Cast the tensor data type to `EnsureDtypeKwargs.dtype` (if not matching).
1069 This can for example be used to ensure the inner neural network model gets a
1070 different input tensor data type than the fully described bioimage.io model does.
1072 Examples:
1073 The described bioimage.io model (incl. preprocessing) accepts any
1074 float32-compatible tensor, normalizes it with percentiles and clipping and then
1075 casts it to uint8, which is what the neural network in this example expects.
1076 - in YAML
1077 ```yaml
1078 inputs:
1079 - data:
1080 type: float32 # described bioimage.io model is compatible with any float32 input tensor
1081 preprocessing:
1082 - id: scale_range
1083 kwargs:
1084 axes: ['y', 'x']
1085 max_percentile: 99.8
1086 min_percentile: 5.0
1087 - id: clip
1088 kwargs:
1089 min: 0.0
1090 max: 1.0
1091 - id: ensure_dtype # the neural network of the model requires uint8
1092 kwargs:
1093 dtype: uint8
1094 ```
1095 - in Python:
1096 >>> preprocessing = [
1097 ... ScaleRangeDescr(
1098 ... kwargs=ScaleRangeKwargs(
1099 ... axes= (AxisId('y'), AxisId('x')),
1100 ... max_percentile= 99.8,
1101 ... min_percentile= 5.0,
1102 ... )
1103 ... ),
1104 ... ClipDescr(kwargs=ClipKwargs(min=0.0, max=1.0)),
1105 ... EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype="uint8")),
1106 ... ]
1107 """
1109 implemented_id: ClassVar[Literal["ensure_dtype"]] = "ensure_dtype"
1110 if TYPE_CHECKING:
1111 id: Literal["ensure_dtype"] = "ensure_dtype"
1112 else:
1113 id: Literal["ensure_dtype"]
1115 kwargs: EnsureDtypeKwargs
1118class ScaleLinearKwargs(KwargsNode):
1119 """Key word arguments for [ScaleLinearDescr][]"""
1121 gain: float = 1.0
1122 """multiplicative factor"""
1124 offset: float = 0.0
1125 """additive term"""
1127 @model_validator(mode="after")
1128 def _validate(self) -> Self:
1129 if self.gain == 1.0 and self.offset == 0.0:
1130 raise ValueError(
1131 "Redundant linear scaling not allowd. Set `gain` != 1.0 and/or `offset`"
1132 + " != 0.0."
1133 )
1135 return self
1138class ScaleLinearAlongAxisKwargs(KwargsNode):
1139 """Key word arguments for [ScaleLinearDescr][]"""
1141 axis: Annotated[NonBatchAxisId, Field(examples=["channel"])]
1142 """The axis of gain and offset values."""
1144 gain: float | NotEmpty[list[float]] = 1.0
1145 """multiplicative factor"""
1147 offset: float | NotEmpty[list[float]] = 0.0
1148 """additive term"""
1150 @model_validator(mode="after")
1151 def _validate(self) -> Self:
1152 if isinstance(self.gain, list):
1153 if isinstance(self.offset, list):
1154 if len(self.gain) != len(self.offset):
1155 raise ValueError(
1156 f"Size of `gain` ({len(self.gain)}) and `offset` ({len(self.offset)}) must match."
1157 )
1158 else:
1159 self.offset = [float(self.offset)] * len(self.gain)
1160 elif isinstance(self.offset, list):
1161 self.gain = [float(self.gain)] * len(self.offset)
1162 else:
1163 raise ValueError(
1164 "Do not specify an `axis` for scalar gain and offset values."
1165 )
1167 if all(g == 1.0 for g in self.gain) and all(off == 0.0 for off in self.offset):
1168 raise ValueError(
1169 "Redundant linear scaling not allowd. Set `gain` != 1.0 and/or `offset`"
1170 + " != 0.0."
1171 )
1173 return self
1176class ScaleLinearDescr(NodeWithExplicitlySetFields):
1177 """Fixed linear scaling.
1179 Examples:
1180 1. Scale with scalar gain and offset
1181 - in YAML
1182 ```yaml
1183 preprocessing:
1184 - id: scale_linear
1185 kwargs:
1186 gain: 2.0
1187 offset: 3.0
1188 ```
1189 - in Python:
1191 >>> preprocessing = [
1192 ... ScaleLinearDescr(kwargs=ScaleLinearKwargs(gain= 2.0, offset=3.0))
1193 ... ]
1195 2. Independent scaling along an axis
1196 - in YAML
1197 ```yaml
1198 preprocessing:
1199 - id: scale_linear
1200 kwargs:
1201 axis: 'channel'
1202 gain: [1.0, 2.0, 3.0]
1203 ```
1204 - in Python:
1206 >>> preprocessing = [
1207 ... ScaleLinearDescr(
1208 ... kwargs=ScaleLinearAlongAxisKwargs(
1209 ... axis=AxisId("channel"),
1210 ... gain=[1.0, 2.0, 3.0],
1211 ... )
1212 ... )
1213 ... ]
1215 """
1217 implemented_id: ClassVar[Literal["scale_linear"]] = "scale_linear"
1218 if TYPE_CHECKING:
1219 id: Literal["scale_linear"] = "scale_linear"
1220 else:
1221 id: Literal["scale_linear"]
1222 kwargs: ScaleLinearKwargs | ScaleLinearAlongAxisKwargs
1225class SigmoidDescr(NodeWithExplicitlySetFields):
1226 """The logistic sigmoid function, a.k.a. expit function.
1228 Examples:
1229 - in YAML
1230 ```yaml
1231 postprocessing:
1232 - id: sigmoid
1233 ```
1234 - in Python:
1236 >>> postprocessing = [SigmoidDescr()]
1237 """
1239 implemented_id: ClassVar[Literal["sigmoid"]] = "sigmoid"
1240 if TYPE_CHECKING:
1241 id: Literal["sigmoid"] = "sigmoid"
1242 else:
1243 id: Literal["sigmoid"]
1245 @property
1246 def kwargs(self) -> KwargsNode:
1247 """empty kwargs"""
1248 return KwargsNode()
1251class SoftmaxKwargs(KwargsNode):
1252 """key word arguments for [SoftmaxDescr][]"""
1254 axis: Annotated[NonBatchAxisId, Field(examples=["channel"])] = CHANNEL_AXIS_ID
1255 """The axis to apply the softmax function along.
1256 Note:
1257 Defaults to 'channel' axis
1258 (which may not exist, in which case
1259 a different axis id has to be specified).
1260 """
1263class SoftmaxDescr(NodeWithExplicitlySetFields):
1264 """The softmax function.
1266 Examples:
1267 - in YAML
1268 ```yaml
1269 postprocessing:
1270 - id: softmax
1271 kwargs:
1272 axis: channel
1273 ```
1274 - in Python:
1276 >>> postprocessing = [SoftmaxDescr(kwargs=SoftmaxKwargs(axis=AxisId("channel")))]
1277 """
1279 implemented_id: ClassVar[Literal["softmax"]] = "softmax"
1280 if TYPE_CHECKING:
1281 id: Literal["softmax"] = "softmax"
1282 else:
1283 id: Literal["softmax"]
1285 kwargs: SoftmaxKwargs = Field(default_factory=SoftmaxKwargs.model_construct)
1288class _StardistPostprocessingKwargsBase(KwargsNode):
1289 """key word arguments for [StardistPostprocessingDescr][]"""
1291 prob_threshold: float
1292 """The probability threshold for object candidate selection."""
1294 nms_threshold: float
1295 """The IoU threshold for non-maximum suppression."""
1297 n_rays: int
1298 """Number of radial lines (rays) cast from the center of an object to its boundary."""
1301class StardistPostprocessingKwargs2D(_StardistPostprocessingKwargsBase):
1302 grid: tuple[int, int]
1303 """Grid size of network predictions."""
1305 b: int | tuple[tuple[int, int], tuple[int, int]]
1306 """Border region in which object probability is set to zero."""
1309class StardistPostprocessingKwargs3D(_StardistPostprocessingKwargsBase):
1310 grid: tuple[int, int, int]
1311 """Grid size of network predictions."""
1313 b: int | tuple[tuple[int, int], tuple[int, int], tuple[int, int]]
1314 """Border region in which object probability is set to zero."""
1316 anisotropy: tuple[float, float, float]
1317 """Anisotropy factors for 3D star-convex polyhedra, i.e. the physical pixel size along each spatial axis."""
1319 overlap_label: int | None = None
1320 """Optional label to apply to any area of overlapping predicted objects."""
1323class StardistPostprocessingDescr(NodeWithExplicitlySetFields):
1324 """Stardist postprocessing including non-maximum suppression and converting polygon representations to instance labels
1326 as described in:
1327 - Uwe Schmidt, Martin Weigert, Coleman Broaddus, and Gene Myers.
1328 [*Cell Detection with Star-convex Polygons*](https://arxiv.org/abs/1806.03535).
1329 International Conference on Medical Image Computing and Computer-Assisted Intervention (MICCAI), Granada, Spain, September 2018.
1330 - Martin Weigert, Uwe Schmidt, Robert Haase, Ko Sugawara, and Gene Myers.
1331 [*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).
1332 The IEEE Winter Conference on Applications of Computer Vision (WACV), Snowmass Village, Colorado, March 2020.
1334 Note: Only available if the `stardist` package is installed.
1335 """
1337 implemented_id: ClassVar[Literal["stardist_postprocessing"]] = (
1338 "stardist_postprocessing"
1339 )
1340 if TYPE_CHECKING:
1341 id: Literal["stardist_postprocessing"] = "stardist_postprocessing"
1342 else:
1343 id: Literal["stardist_postprocessing"]
1345 kwargs: StardistPostprocessingKwargs2D | StardistPostprocessingKwargs3D
1348class CellposeFlowDynamicsKwargs(KwargsNode):
1349 """key word arguments for [CellposeFlowDynamicsDescr][]"""
1351 cellprob_threshold: float
1352 flow_threshold: float
1353 do_3D: bool
1354 min_size: int = 15
1355 """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."""
1356 output_dtype: Literal["uint16", "uint32"] = "uint16"
1359class CellposeFlowDynamicsDescr(NodeWithExplicitlySetFields):
1360 """Cellpose flow dynamics postprocessing as described in:
1361 - Carsen Stringer and Marius Pachitariu. [*Cellpose: a generalist algorithm for cellular segmentation*](https://www.nature.com/articles/s41592-020-01018-x). Nature Methods, 2021.
1363 Note: Only available if the `cellpose` package is installed.
1364 """
1366 implemented_id: ClassVar[Literal["cellpose_flow_dynamics"]] = (
1367 "cellpose_flow_dynamics"
1368 )
1369 if TYPE_CHECKING:
1370 id: Literal["cellpose_flow_dynamics"] = "cellpose_flow_dynamics"
1371 else:
1372 id: Literal["cellpose_flow_dynamics"]
1374 kwargs: CellposeFlowDynamicsKwargs
1377class CustomProcessingDescr(NodeWithExplicitlySetFields, FileDescr):
1378 """Custom (post)processing op — source file shipped inline with the model.
1380 Supports (post)processing that cannot be expressed by the built-in named
1381 operations (watershed, connected components, etc.)
1382 using a simple Python callable interface.
1384 The op is implemented in a ``.py`` file packaged alongside the model weights.
1385 Two styles are supported:
1387 *Callable class* — kwargs go to ``__init__``, tensors arrive in ``__call__``:
1389 .. code-block:: python
1391 # my_postprocess.py
1392 import numpy as np
1394 class my_postprocess:
1395 def __init__(self, threshold: float = 0.5) -> None:
1396 self.threshold = threshold
1397 def __call__(self, *arrays: np.ndarray) -> np.ndarray:
1398 # arrays = model output tensors in rdf.yaml declaration order
1399 return (arrays[0] > self.threshold).astype(np.uint8)
1401 *Factory function* — alternative closure style, identical runtime behaviour:
1403 .. code-block:: python
1405 # my_postprocess.py
1406 import numpy as np
1408 def my_postprocess(threshold: float = 0.5):
1409 def run(*arrays: np.ndarray) -> np.ndarray:
1410 return (arrays[0] > threshold).astype(np.uint8)
1411 return run
1413 Reference it in ``rdf.yaml`` with the source file included in the package:
1415 .. code-block:: yaml
1417 postprocessing:
1418 - id: custom
1419 callable: my_postprocess # class or function name in source
1420 source: my_postprocess.py # packaged alongside weights
1421 sha256: <hash> # sha256 of the source file
1422 kwargs: # forwarded to __init__ / factory
1423 threshold: 0.5
1425 **Security:** source files are SHA-256 verified before execution.
1426 Execution requires explicit opt-in in bioimageio.core and curator
1427 review before Zoo publication.
1428 """
1430 implemented_id: ClassVar[Literal["custom"]] = "custom"
1431 if TYPE_CHECKING:
1432 id: Literal["custom"] = "custom"
1433 else:
1434 id: Literal["custom"]
1436 callable: Annotated[
1437 str,
1438 Field(examples=["my_postprocess_factory", "MyPostprocessClass"]),
1439 ]
1440 """Name of the callable class or factory function defined in ``source``.
1442 At runtime: ``op = callable(**kwargs)``, then ``result = op(*output_tensors)``
1443 per image. Both a class with ``__call__`` and a factory function returning
1444 a callable satisfy this protocol."""
1446 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
1447 """Python source file (included when packaging the model)."""
1449 kwargs: dict[str, YamlValue] = Field(
1450 default_factory=cast(Callable[[], Dict[str, YamlValue]], dict)
1451 )
1452 """Keyword arguments forwarded to the callable (``__init__`` or factory)."""
1454 @model_serializer(mode="wrap", when_used="unless-none")
1455 def _serialize(
1456 self, nxt: SerializerFunctionWrapHandler, info: SerializationInfo
1457 ) -> dict[str, YamlValue]:
1458 return package_file_descr_serializer(self, nxt, info)
1461class FixedZeroMeanUnitVarianceKwargs(KwargsNode):
1462 """key word arguments for [FixedZeroMeanUnitVarianceDescr][]"""
1464 mean: float
1465 """The mean value to normalize with."""
1467 std: Annotated[float, Ge(1e-6)]
1468 """The standard deviation value to normalize with."""
1471class FixedZeroMeanUnitVarianceAlongAxisKwargs(KwargsNode):
1472 """key word arguments for [FixedZeroMeanUnitVarianceDescr][]"""
1474 mean: NotEmpty[list[float]]
1475 """The mean value(s) to normalize with."""
1477 std: NotEmpty[list[Annotated[float, Ge(1e-6)]]]
1478 """The standard deviation value(s) to normalize with.
1479 Size must match `mean` values."""
1481 axis: Annotated[NonBatchAxisId, Field(examples=["channel", "index"])]
1482 """The axis of the mean/std values to normalize each entry along that dimension
1483 separately."""
1485 @model_validator(mode="after")
1486 def _mean_and_std_match(self) -> Self:
1487 if len(self.mean) != len(self.std):
1488 raise ValueError(
1489 f"Size of `mean` ({len(self.mean)}) and `std` ({len(self.std)})"
1490 + " must match."
1491 )
1493 return self
1496class FixedZeroMeanUnitVarianceDescr(NodeWithExplicitlySetFields):
1497 """Subtract a given mean and divide by the standard deviation.
1499 Normalize with fixed, precomputed values for
1500 `FixedZeroMeanUnitVarianceKwargs.mean` and `FixedZeroMeanUnitVarianceKwargs.std`
1501 Use `FixedZeroMeanUnitVarianceAlongAxisKwargs` for independent scaling along given
1502 axes.
1504 Examples:
1505 1. scalar value for whole tensor
1506 - in YAML
1507 ```yaml
1508 preprocessing:
1509 - id: fixed_zero_mean_unit_variance
1510 kwargs:
1511 mean: 103.5
1512 std: 13.7
1513 ```
1514 - in Python
1515 >>> preprocessing = [FixedZeroMeanUnitVarianceDescr(
1516 ... kwargs=FixedZeroMeanUnitVarianceKwargs(mean=103.5, std=13.7)
1517 ... )]
1519 2. independently along an axis
1520 - in YAML
1521 ```yaml
1522 preprocessing:
1523 - id: fixed_zero_mean_unit_variance
1524 kwargs:
1525 axis: channel
1526 mean: [101.5, 102.5, 103.5]
1527 std: [11.7, 12.7, 13.7]
1528 ```
1529 - in Python
1530 >>> preprocessing = [FixedZeroMeanUnitVarianceDescr(
1531 ... kwargs=FixedZeroMeanUnitVarianceAlongAxisKwargs(
1532 ... axis=AxisId("channel"),
1533 ... mean=[101.5, 102.5, 103.5],
1534 ... std=[11.7, 12.7, 13.7],
1535 ... )
1536 ... )]
1537 """
1539 implemented_id: ClassVar[Literal["fixed_zero_mean_unit_variance"]] = (
1540 "fixed_zero_mean_unit_variance"
1541 )
1542 if TYPE_CHECKING:
1543 id: Literal["fixed_zero_mean_unit_variance"] = "fixed_zero_mean_unit_variance"
1544 else:
1545 id: Literal["fixed_zero_mean_unit_variance"]
1547 kwargs: FixedZeroMeanUnitVarianceKwargs | FixedZeroMeanUnitVarianceAlongAxisKwargs
1550class ZeroMeanUnitVarianceKwargs(KwargsNode):
1551 """key word arguments for [ZeroMeanUnitVarianceDescr][]"""
1553 axes: Annotated[Sequence[AxisId] | None, Field(examples=[("batch", "x", "y")])] = (
1554 None
1555 )
1556 """The subset of axes to normalize jointly, i.e. axes to reduce to compute mean/std.
1557 For example to normalize 'batch', 'x' and 'y' jointly in a tensor ('batch', 'channel', 'y', 'x')
1558 resulting in a tensor of equal shape normalized per channel, specify `axes=('batch', 'x', 'y')`.
1559 To normalize each sample independently leave out the 'batch' axis.
1560 Default: Scale all axes jointly."""
1562 eps: Annotated[float, Interval(gt=0, le=0.1)] = 1e-6
1563 """epsilon for numeric stability: `out = (tensor - mean) / (std + eps)`."""
1566class ZeroMeanUnitVarianceDescr(NodeWithExplicitlySetFields):
1567 """Subtract mean and divide by variance.
1569 Examples:
1570 Subtract tensor mean and variance
1571 - in YAML
1572 ```yaml
1573 preprocessing:
1574 - id: zero_mean_unit_variance
1575 ```
1576 - in Python
1577 >>> preprocessing = [ZeroMeanUnitVarianceDescr()]
1578 """
1580 implemented_id: ClassVar[Literal["zero_mean_unit_variance"]] = (
1581 "zero_mean_unit_variance"
1582 )
1583 if TYPE_CHECKING:
1584 id: Literal["zero_mean_unit_variance"] = "zero_mean_unit_variance"
1585 else:
1586 id: Literal["zero_mean_unit_variance"]
1588 kwargs: ZeroMeanUnitVarianceKwargs = Field(
1589 default_factory=ZeroMeanUnitVarianceKwargs.model_construct
1590 )
1593class ScaleRangeKwargs(KwargsNode):
1594 """key word arguments for [ScaleRangeDescr][]
1596 For `min_percentile`=0.0 (the default) and `max_percentile`=100 (the default)
1597 this processing step normalizes data to the [0, 1] intervall.
1598 For other percentiles the normalized values will partially be outside the [0, 1]
1599 intervall. Use `ScaleRange` followed by `ClipDescr` if you want to limit the
1600 normalized values to a range.
1601 """
1603 axes: Annotated[Sequence[AxisId] | None, Field(examples=[("batch", "x", "y")])] = (
1604 None
1605 )
1606 """The subset of axes to normalize jointly, i.e. axes to reduce to compute the min/max percentile value.
1607 For example to normalize 'batch', 'x' and 'y' jointly in a tensor ('batch', 'channel', 'y', 'x')
1608 resulting in a tensor of equal shape normalized per channel, specify `axes=('batch', 'x', 'y')`.
1609 To normalize samples independently, leave out the "batch" axis.
1610 Default: Scale all axes jointly."""
1612 min_percentile: Annotated[float, Interval(ge=0, lt=100)] = 0.0
1613 """The lower percentile used to determine the value to align with zero."""
1615 max_percentile: Annotated[float, Interval(gt=1, le=100)] = 100.0
1616 """The upper percentile used to determine the value to align with one.
1617 Has to be bigger than `min_percentile`.
1618 The range is 1 to 100 instead of 0 to 100 to avoid mistakenly
1619 accepting percentiles specified in the range 0.0 to 1.0."""
1621 eps: Annotated[float, Interval(gt=0, le=0.1)] = 1e-6
1622 """Epsilon for numeric stability.
1623 `out = (tensor - v_lower) / (v_upper - v_lower + eps)`;
1624 with `v_lower,v_upper` values at the respective percentiles."""
1626 reference_tensor: TensorId | None = None
1627 """ID of the unprocessed input tensor to compute the percentiles from.
1628 Default: The tensor itself.
1629 """
1631 @field_validator("max_percentile", mode="after")
1632 @classmethod
1633 def min_smaller_max(cls, value: float, info: ValidationInfo) -> float:
1634 if (min_p := info.data["min_percentile"]) >= value:
1635 raise ValueError(f"min_percentile {min_p} >= max_percentile {value}")
1637 return value
1640class ScaleRangeDescr(NodeWithExplicitlySetFields):
1641 """Scale with percentiles.
1643 Examples:
1644 1. Scale linearly to map 5th percentile to 0 and 99.8th percentile to 1.0
1645 - in YAML
1646 ```yaml
1647 preprocessing:
1648 - id: scale_range
1649 kwargs:
1650 axes: ['y', 'x']
1651 max_percentile: 99.8
1652 min_percentile: 5.0
1653 ```
1654 - in Python
1656 >>> preprocessing = [
1657 ... ScaleRangeDescr(
1658 ... kwargs=ScaleRangeKwargs(
1659 ... axes= (AxisId('y'), AxisId('x')),
1660 ... max_percentile= 99.8,
1661 ... min_percentile= 5.0,
1662 ... )
1663 ... )
1664 ... ]
1666 2. Combine the above scaling with additional clipping to clip values outside the range given by the percentiles.
1667 - in YAML
1668 ```yaml
1669 preprocessing:
1670 - id: scale_range
1671 kwargs:
1672 axes: ['y', 'x']
1673 max_percentile: 99.8
1674 min_percentile: 5.0
1675 - id: clip
1676 kwargs:
1677 min: 0.0
1678 max: 1.0
1679 ```
1680 - in Python
1682 >>> preprocessing = [
1683 ... ScaleRangeDescr(
1684 ... kwargs=ScaleRangeKwargs(
1685 ... axes= (AxisId('y'), AxisId('x')),
1686 ... max_percentile= 99.8,
1687 ... min_percentile= 5.0,
1688 ... )
1689 ... ),
1690 ... ClipDescr(
1691 ... kwargs=ClipKwargs(
1692 ... min=0.0,
1693 ... max=1.0,
1694 ... )
1695 ... ),
1696 ... ]
1698 """
1700 implemented_id: ClassVar[Literal["scale_range"]] = "scale_range"
1701 if TYPE_CHECKING:
1702 id: Literal["scale_range"] = "scale_range"
1703 else:
1704 id: Literal["scale_range"]
1705 kwargs: ScaleRangeKwargs = Field(default_factory=ScaleRangeKwargs.model_construct)
1708class ScaleMeanVarianceKwargs(KwargsNode):
1709 """key word arguments for [ScaleMeanVarianceKwargs][]"""
1711 reference_tensor: TensorId
1712 """ID of unprocessed input tensor to match."""
1714 axes: Annotated[Sequence[AxisId] | None, Field(examples=[("batch", "x", "y")])] = (
1715 None
1716 )
1717 """The subset of axes to normalize jointly, i.e. axes to reduce to compute mean/std.
1718 For example to normalize 'batch', 'x' and 'y' jointly in a tensor ('batch', 'channel', 'y', 'x')
1719 resulting in a tensor of equal shape normalized per channel, specify `axes=('batch', 'x', 'y')`.
1720 To normalize samples independently, leave out the 'batch' axis.
1721 Default: Scale all axes jointly."""
1723 eps: Annotated[float, Interval(gt=0, le=0.1)] = 1e-6
1724 """Epsilon for numeric stability:
1725 `out = (tensor - mean) / (std + eps) * (ref_std + eps) + ref_mean.`"""
1728class ScaleMeanVarianceDescr(NodeWithExplicitlySetFields):
1729 """Scale a tensor's data distribution to match another tensor's mean/std.
1730 `out = (tensor - mean) / (std + eps) * (ref_std + eps) + ref_mean.`
1731 """
1733 implemented_id: ClassVar[Literal["scale_mean_variance"]] = "scale_mean_variance"
1734 if TYPE_CHECKING:
1735 id: Literal["scale_mean_variance"] = "scale_mean_variance"
1736 else:
1737 id: Literal["scale_mean_variance"]
1738 kwargs: ScaleMeanVarianceKwargs
1741PreprocessingDescr = Annotated[
1742 Union[
1743 BinarizeDescr,
1744 ClipDescr,
1745 EnsureDtypeDescr,
1746 FixedZeroMeanUnitVarianceDescr,
1747 ScaleLinearDescr,
1748 ScaleRangeDescr,
1749 SigmoidDescr,
1750 SoftmaxDescr,
1751 ZeroMeanUnitVarianceDescr,
1752 ],
1753 Discriminator("id"),
1754]
1755PostprocessingDescr = Annotated[
1756 Union[
1757 BinarizeDescr,
1758 CellposeFlowDynamicsDescr,
1759 ClipDescr,
1760 CustomProcessingDescr,
1761 EnsureDtypeDescr,
1762 FixedZeroMeanUnitVarianceDescr,
1763 ScaleLinearDescr,
1764 ScaleMeanVarianceDescr,
1765 ScaleRangeDescr,
1766 SigmoidDescr,
1767 SoftmaxDescr,
1768 StardistPostprocessingDescr,
1769 ZeroMeanUnitVarianceDescr,
1770 ],
1771 Discriminator("id"),
1772]
1774IO_AxisT = TypeVar("IO_AxisT", InputAxis, OutputAxis)
1777class TensorDescrBase(Node, Generic[IO_AxisT]):
1778 id: TensorId
1779 """Tensor id. No duplicates are allowed."""
1781 description: Annotated[str, MaxLen(128)] = ""
1782 """free text description"""
1784 axes: NotEmpty[Sequence[IO_AxisT]]
1785 """tensor axes"""
1787 @property
1788 def shape(self):
1789 return tuple(a.size for a in self.axes)
1791 @field_validator("axes", mode="after", check_fields=False)
1792 @classmethod
1793 def _validate_axes(cls, axes: Sequence[AnyAxis]) -> Sequence[AnyAxis]:
1794 batch_axes = [a for a in axes if a.type == "batch"]
1795 if len(batch_axes) > 1:
1796 raise ValueError(
1797 f"Only one batch axis (per tensor) allowed, but got {batch_axes}"
1798 )
1800 seen_ids: set[AxisId] = set()
1801 duplicate_axes_ids: set[AxisId] = set()
1802 for a in axes:
1803 (duplicate_axes_ids if a.id in seen_ids else seen_ids).add(a.id)
1805 if duplicate_axes_ids:
1806 raise ValueError(f"Duplicate axis ids: {duplicate_axes_ids}")
1808 return axes
1810 test_tensor: FAIR[FileDescr_package | None] = None
1811 """An example tensor to use for testing.
1812 Using the model with the test input tensors is expected to yield the test output tensors.
1813 Each test tensor has be a an ndarray in the
1814 [numpy.lib file format](https://numpy.org/doc/stable/reference/generated/numpy.lib.format.html#module-numpy.lib.format).
1815 The file extension must be '.npy'."""
1817 sample_tensor: FAIR[FileDescr_package | None] = None
1818 """A sample tensor to illustrate a possible input/output for the model,
1819 The sample image primarily serves to inform a human user about an example use case
1820 and is typically stored as .hdf5, .png or .tiff.
1821 It has to be readable by the [imageio library](https://imageio.readthedocs.io/en/stable/formats/index.html#supported-formats)
1822 (numpy's `.npy` format is not supported).
1823 The image dimensionality has to match the number of axes specified in this tensor description.
1824 """
1826 @model_validator(mode="after")
1827 def _validate_sample_tensor(self) -> Self:
1828 if self.sample_tensor is None or not get_validation_context().perform_io_checks:
1829 return self
1831 reader = get_reader(self.sample_tensor.source, sha256=self.sample_tensor.sha256)
1832 tensor: NDArray[Any] = imread( # pyright: ignore[reportUnknownVariableType]
1833 reader.read(),
1834 extension=PurePosixPath(reader.original_file_name).suffix,
1835 )
1836 n_dims = len(tensor.squeeze().shape)
1837 n_dims_min = n_dims_max = len(self.axes)
1839 for a in self.axes:
1840 if isinstance(a, BatchAxis):
1841 n_dims_min -= 1
1842 elif isinstance(a.size, int):
1843 if a.size == 1:
1844 n_dims_min -= 1
1845 elif isinstance(a.size, (ParameterizedSize, DataDependentSize)):
1846 if a.size.min == 1:
1847 n_dims_min -= 1
1848 elif isinstance(a.size, SizeReference):
1849 if a.size.offset < 2:
1850 # size reference may result in singleton axis
1851 n_dims_min -= 1
1852 else:
1853 assert_never(a.size)
1855 n_dims_min = max(0, n_dims_min)
1856 if n_dims < n_dims_min or n_dims > n_dims_max:
1857 raise ValueError(
1858 f"Expected sample tensor to have {n_dims_min} to"
1859 + f" {n_dims_max} dimensions, but found {n_dims} (shape: {tensor.shape})."
1860 )
1862 return self
1864 data: TensorDataDescr | NotEmpty[Sequence[TensorDataDescr]] = (
1865 IntervalOrRatioDataDescr()
1866 )
1867 """Description of the tensor's data values, optionally per channel.
1868 If specified per channel, the data `type` needs to match across channels."""
1870 @property
1871 def dtype(
1872 self,
1873 ) -> Literal[
1874 "float32",
1875 "float64",
1876 "uint8",
1877 "int8",
1878 "uint16",
1879 "int16",
1880 "uint32",
1881 "int32",
1882 "uint64",
1883 "int64",
1884 "bool",
1885 ]:
1886 """dtype as specified under `data.type` or `data[i].type`"""
1887 if isinstance(self.data, collections.abc.Sequence):
1888 return self.data[0].type
1889 else:
1890 return self.data.type
1892 @field_validator("data", mode="after")
1893 @classmethod
1894 def _check_data_type_across_channels(
1895 cls, value: TensorDataDescr | NotEmpty[Sequence[TensorDataDescr]]
1896 ) -> TensorDataDescr | NotEmpty[Sequence[TensorDataDescr]]:
1897 if not isinstance(value, list):
1898 return value
1900 dtypes = {t.type for t in value}
1901 if len(dtypes) > 1:
1902 raise ValueError(
1903 "Tensor data descriptions per channel need to agree in their data"
1904 + f" `type`, but found {dtypes}."
1905 )
1907 return value
1909 @model_validator(mode="after")
1910 def _check_data_matches_channelaxis(self) -> Self:
1911 if not isinstance(self.data, (list, tuple)):
1912 return self
1914 for a in self.axes:
1915 if isinstance(a, ChannelAxis):
1916 size = a.size
1917 assert isinstance(size, int)
1918 break
1919 else:
1920 return self
1922 if len(self.data) != size:
1923 raise ValueError(
1924 f"Got tensor data descriptions for {len(self.data)} channels, but"
1925 + f" '{a.id}' axis has size {size}."
1926 )
1928 return self
1930 def get_axis_sizes_for_array(self, array: NDArray[Any]) -> dict[AxisId, int]:
1931 if len(array.shape) != len(self.axes):
1932 raise ValueError(
1933 f"Dimension mismatch: array shape {array.shape} (#{len(array.shape)})"
1934 + f" incompatible with {len(self.axes)} axes."
1935 )
1936 return {a.id: array.shape[i] for i, a in enumerate(self.axes)}
1939class ConstantPadding(Node):
1940 mode: Literal["constant"] = "constant"
1941 value: int | float = 0
1944class EdgePadding(Node):
1945 mode: Literal["edge"] = "edge"
1948class ReflectPadding(Node):
1949 mode: Literal["reflect"] = "reflect"
1952class SymmetricPadding(Node):
1953 mode: Literal["symmetric"] = "symmetric"
1956Padding = Union[ConstantPadding, EdgePadding, ReflectPadding, SymmetricPadding]
1959class ModelId(ResourceId):
1960 pass
1963class InputTensorDescr(TensorDescrBase[InputAxis]):
1964 id: TensorId = TensorId("input")
1965 """Input tensor id.
1966 No duplicates are allowed across all inputs and outputs."""
1968 output_of: ModelId | None = None
1969 """If this input tensor is the output of another model, specify the model id here.
1970 This model's input id must match the output id of the referenced model.
1971 """
1973 @model_validator(mode="after")
1974 def _validate_output_of(self) -> Self:
1975 if self.output_of is None:
1976 return self
1978 try:
1979 with get_validation_context().replace(perform_io_checks=False):
1980 opened_ref_model = open_bioimageio_yaml(self.output_of)
1981 format_version = opened_ref_model.content["format_version"]
1982 assert isinstance(format_version, str)
1983 if format_version.startswith("0.4"):
1984 ref_model = _ModelDescr04.model_validate(opened_ref_model.content)
1985 else:
1986 ref_model = ModelDescr.model_validate(opened_ref_model.content)
1987 except Exception as e:
1988 raise ValueError(
1989 f"Failed to load model '{self.output_of}' referenced under output_of: {e}"
1990 )
1992 try:
1993 ref_model_outputs = {
1994 t.id if isinstance(t, OutputTensorDescr) else TensorId(t.name)
1995 for t in ref_model.outputs
1996 }
1997 except Exception as e:
1998 raise ValueError(
1999 f"Failed to read output IDs of model '{self.output_of}' referenced under output_of: {e}"
2000 )
2002 if self.id not in ref_model_outputs:
2003 raise ValueError(
2004 f"Input tensor '{self.id}' is specified as output of model '{self.output_of}', "
2005 + f"but that model's outputs are {ref_model_outputs}."
2006 )
2007 return self
2009 optional: bool = False
2010 """indicates that this tensor may be `None`"""
2012 pad: Padding | None = None
2013 """Explicitly specify how to pad this input tensor.
2015 Use `axes[i].pad` to specify padding width.
2017 Note:
2018 Non-blockwise sample prediction only applies padding for axes with a `pad` specification.
2019 """
2021 preprocessing: list[PreprocessingDescr] = Field(
2022 default_factory=cast(Callable[[], List[PreprocessingDescr]], list)
2023 )
2024 """Description of how this input should be preprocessed.
2026 notes:
2027 - If preprocessing does not start with an 'ensure_dtype' entry, it is added
2028 to ensure an input tensor's data type matches the input tensor's data description.
2029 - If preprocessing does not end with an 'ensure_dtype' or 'binarize' entry, an
2030 'ensure_dtype' step is added to ensure preprocessing steps are not unintentionally
2031 changing the data type.
2032 """
2034 @model_validator(mode="after")
2035 def _validate_preprocessing_kwargs(self) -> Self:
2036 axes_ids = [a.id for a in self.axes]
2037 for p in self.preprocessing:
2038 kwargs_axes: Sequence[Any] | None = p.kwargs.get("axes")
2039 if kwargs_axes is None:
2040 continue
2042 if not isinstance(kwargs_axes, collections.abc.Sequence):
2043 raise ValueError(
2044 f"Expected `preprocessing.i.kwargs.axes` to be a sequence, but got {type(kwargs_axes)}"
2045 )
2047 if any(a not in axes_ids for a in kwargs_axes):
2048 raise ValueError(
2049 "`preprocessing.i.kwargs.axes` needs to be subset of axes ids"
2050 )
2052 if isinstance(self.data, (NominalOrOrdinalDataDescr, IntervalOrRatioDataDescr)):
2053 dtype = self.data.type
2054 else:
2055 dtype = self.data[0].type
2057 # ensure `preprocessing` begins with `EnsureDtypeDescr`
2058 if not self.preprocessing or not isinstance(
2059 self.preprocessing[0], EnsureDtypeDescr
2060 ):
2061 self.preprocessing.insert(
2062 0, EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype=dtype))
2063 )
2065 # ensure `preprocessing` ends with `EnsureDtypeDescr` or `BinarizeDescr`
2066 if not isinstance(self.preprocessing[-1], (EnsureDtypeDescr, BinarizeDescr)):
2067 self.preprocessing.append(
2068 EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype=dtype))
2069 )
2071 return self
2074def convert_axes(
2075 axes: str,
2076 *,
2077 shape: Sequence[int] | _ParameterizedInputShape_v0_4 | _ImplicitOutputShape_v0_4,
2078 tensor_type: Literal["input", "output"],
2079 halo: Sequence[int] | None,
2080 size_refs: Mapping[_TensorName_v0_4, Mapping[str, int]],
2081):
2082 ret: list[AnyAxis] = []
2083 for i, a in enumerate(axes):
2084 axis_type = _AXIS_TYPE_MAP.get(a, a)
2085 if axis_type == "batch":
2086 ret.append(BatchAxis())
2087 continue
2089 scale = 1.0
2090 if isinstance(shape, _ParameterizedInputShape_v0_4):
2091 if shape.step[i] == 0:
2092 size = shape.min[i]
2093 else:
2094 size = ParameterizedSize(min=shape.min[i], step=shape.step[i])
2095 elif isinstance(shape, _ImplicitOutputShape_v0_4):
2096 ref_t = str(shape.reference_tensor)
2097 if ref_t.count(".") == 1:
2098 t_id, orig_a_id = ref_t.split(".")
2099 else:
2100 t_id = ref_t
2101 orig_a_id = a
2103 a_id = _AXIS_ID_MAP.get(orig_a_id, a)
2104 if not (orig_scale := shape.scale[i]):
2105 # old way to insert a new axis dimension
2106 size = int(2 * shape.offset[i])
2107 else:
2108 scale = 1 / orig_scale
2109 if axis_type in ("channel", "index"):
2110 # these axes no longer have a scale
2111 offset_from_scale = orig_scale * size_refs.get(
2112 _TensorName_v0_4(t_id), {}
2113 ).get(orig_a_id, 0)
2114 else:
2115 offset_from_scale = 0
2116 size = SizeReference(
2117 tensor_id=TensorId(t_id),
2118 axis_id=AxisId(a_id),
2119 offset=int(offset_from_scale + 2 * shape.offset[i]),
2120 )
2121 else:
2122 size = shape[i]
2124 if axis_type == "time":
2125 if tensor_type == "input":
2126 ret.append(TimeInputAxis(size=size, scale=scale))
2127 else:
2128 assert not isinstance(size, ParameterizedSize)
2129 if halo is None:
2130 ret.append(TimeOutputAxis(size=size, scale=scale))
2131 else:
2132 assert not isinstance(size, int)
2133 ret.append(
2134 TimeOutputAxisWithHalo(size=size, scale=scale, halo=halo[i])
2135 )
2137 elif axis_type == "index":
2138 if tensor_type == "input":
2139 ret.append(IndexInputAxis(size=size))
2140 else:
2141 if isinstance(size, ParameterizedSize):
2142 size = DataDependentSize(min=size.min)
2144 ret.append(IndexOutputAxis(size=size))
2145 elif axis_type == "channel":
2146 assert not isinstance(size, ParameterizedSize)
2147 if isinstance(size, SizeReference):
2148 warnings.warn(
2149 "Conversion of channel size from an implicit output shape may be"
2150 + " wrong"
2151 )
2152 ret.append(
2153 ChannelAxis(
2154 channel_names=[f"channel{i}" for i in range(size.offset)]
2155 )
2156 )
2157 else:
2158 ret.append(
2159 ChannelAxis(channel_names=[f"channel{i}" for i in range(size)])
2160 )
2161 elif axis_type == "space":
2162 if tensor_type == "input":
2163 ret.append(SpaceInputAxis(id=AxisId(a), size=size, scale=scale))
2164 else:
2165 assert not isinstance(size, ParameterizedSize)
2166 if halo is None or halo[i] == 0:
2167 ret.append(SpaceOutputAxis(id=AxisId(a), size=size, scale=scale))
2168 elif isinstance(size, int):
2169 raise NotImplementedError(
2170 f"output axis with halo and fixed size (here {size}) not allowed"
2171 )
2172 else:
2173 ret.append(
2174 SpaceOutputAxisWithHalo(
2175 id=AxisId(a), size=size, scale=scale, halo=halo[i]
2176 )
2177 )
2179 return ret
2182def _axes_letters_to_ids(
2183 axes: str | None,
2184) -> list[AxisId] | None:
2185 if axes is None:
2186 return None
2188 return [AxisId(a) for a in axes]
2191def _get_complement_v04_axis(
2192 tensor_axes: Sequence[str], axes: Sequence[str] | None
2193) -> AxisId | None:
2194 if axes is None:
2195 return None
2197 non_complement_axes = set(axes) | {"b"}
2198 complement_axes = [a for a in tensor_axes if a not in non_complement_axes]
2199 if len(complement_axes) > 1:
2200 raise ValueError(
2201 f"Expected none or a single complement axis, but axes '{axes}' "
2202 + f"for tensor dims '{tensor_axes}' leave '{complement_axes}'."
2203 )
2205 return None if not complement_axes else AxisId(complement_axes[0])
2208def _convert_proc(
2209 p: _PreprocessingDescr_v0_4 | _PostprocessingDescr_v0_4,
2210 tensor_axes: Sequence[str],
2211) -> PreprocessingDescr | PostprocessingDescr:
2212 if isinstance(p, _BinarizeDescr_v0_4):
2213 return BinarizeDescr(kwargs=BinarizeKwargs(threshold=p.kwargs.threshold))
2214 elif isinstance(p, _ClipDescr_v0_4):
2215 return ClipDescr(kwargs=ClipKwargs(min=p.kwargs.min, max=p.kwargs.max))
2216 elif isinstance(p, _SigmoidDescr_v0_4):
2217 return SigmoidDescr()
2218 elif isinstance(p, _ScaleLinearDescr_v0_4):
2219 axes = _axes_letters_to_ids(p.kwargs.axes)
2220 if p.kwargs.axes is None:
2221 axis = None
2222 else:
2223 axis = _get_complement_v04_axis(tensor_axes, p.kwargs.axes)
2225 if axis is None:
2226 assert not isinstance(p.kwargs.gain, list)
2227 assert not isinstance(p.kwargs.offset, list)
2228 kwargs = ScaleLinearKwargs(gain=p.kwargs.gain, offset=p.kwargs.offset)
2229 else:
2230 kwargs = ScaleLinearAlongAxisKwargs(
2231 axis=axis, gain=p.kwargs.gain, offset=p.kwargs.offset
2232 )
2233 return ScaleLinearDescr(kwargs=kwargs)
2234 elif isinstance(p, _ScaleMeanVarianceDescr_v0_4):
2235 return ScaleMeanVarianceDescr(
2236 kwargs=ScaleMeanVarianceKwargs(
2237 axes=_axes_letters_to_ids(p.kwargs.axes),
2238 reference_tensor=TensorId(str(p.kwargs.reference_tensor)),
2239 eps=p.kwargs.eps,
2240 )
2241 )
2242 elif isinstance(p, _ZeroMeanUnitVarianceDescr_v0_4):
2243 if p.kwargs.mode == "fixed":
2244 mean = p.kwargs.mean
2245 std = p.kwargs.std
2246 assert mean is not None
2247 assert std is not None
2249 axis = _get_complement_v04_axis(tensor_axes, p.kwargs.axes)
2251 if axis is None:
2252 if isinstance(mean, list):
2253 raise ValueError("Expected single float value for mean, not <list>")
2254 if isinstance(std, list):
2255 raise ValueError("Expected single float value for std, not <list>")
2256 return FixedZeroMeanUnitVarianceDescr(
2257 kwargs=FixedZeroMeanUnitVarianceKwargs.model_construct(
2258 mean=mean,
2259 std=std,
2260 )
2261 )
2262 else:
2263 if not isinstance(mean, list):
2264 mean = [float(mean)]
2265 if not isinstance(std, list):
2266 std = [float(std)]
2268 return FixedZeroMeanUnitVarianceDescr(
2269 kwargs=FixedZeroMeanUnitVarianceAlongAxisKwargs(
2270 axis=axis, mean=mean, std=std
2271 )
2272 )
2274 else:
2275 axes = _axes_letters_to_ids(p.kwargs.axes) or []
2276 if p.kwargs.mode == "per_dataset":
2277 axes = [AxisId("batch")] + axes
2278 if not axes:
2279 axes = None
2280 return ZeroMeanUnitVarianceDescr(
2281 kwargs=ZeroMeanUnitVarianceKwargs(axes=axes, eps=p.kwargs.eps)
2282 )
2284 elif isinstance(p, _ScaleRangeDescr_v0_4):
2285 return ScaleRangeDescr(
2286 kwargs=ScaleRangeKwargs(
2287 axes=_axes_letters_to_ids(p.kwargs.axes),
2288 min_percentile=p.kwargs.min_percentile,
2289 max_percentile=p.kwargs.max_percentile,
2290 eps=p.kwargs.eps,
2291 )
2292 )
2293 else:
2294 assert_never(p)
2297class _InputTensorConv(
2298 Converter[
2299 _InputTensorDescr_v0_4,
2300 InputTensorDescr,
2301 FileSource,
2302 Optional[FileSource],
2303 Mapping[_TensorName_v0_4, Mapping[str, int]],
2304 ]
2305):
2306 def _convert(
2307 self,
2308 src: _InputTensorDescr_v0_4,
2309 tgt: type[InputTensorDescr | dict[str, Any]],
2310 test_tensor: FileSource,
2311 sample_tensor: FileSource | None,
2312 size_refs: Mapping[_TensorName_v0_4, Mapping[str, int]],
2313 ) -> InputTensorDescr | dict[str, Any]:
2314 axes: list[InputAxis] = convert_axes( # pyright: ignore[reportAssignmentType]
2315 src.axes,
2316 shape=src.shape,
2317 tensor_type="input",
2318 halo=None,
2319 size_refs=size_refs,
2320 )
2321 prep: list[PreprocessingDescr] = []
2322 for p in src.preprocessing:
2323 cp = _convert_proc(p, src.axes)
2324 assert not isinstance(
2325 cp,
2326 (
2327 CellposeFlowDynamicsDescr,
2328 CustomProcessingDescr,
2329 ScaleMeanVarianceDescr,
2330 StardistPostprocessingDescr,
2331 ),
2332 )
2333 prep.append(cp)
2335 prep.append(EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype="float32")))
2337 return tgt(
2338 axes=axes,
2339 id=TensorId(str(src.name)),
2340 test_tensor=FileDescr(source=test_tensor),
2341 sample_tensor=(
2342 None if sample_tensor is None else FileDescr(source=sample_tensor)
2343 ),
2344 data={"type": src.data_type}, # pyright: ignore[reportArgumentType]
2345 preprocessing=prep,
2346 )
2349_input_tensor_conv = _InputTensorConv(_InputTensorDescr_v0_4, InputTensorDescr)
2352class OutputTensorDescr(TensorDescrBase[OutputAxis]):
2353 id: TensorId = TensorId("output")
2354 """Output tensor id.
2355 No duplicates are allowed across all inputs and outputs."""
2357 postprocessing: list[PostprocessingDescr] = Field(
2358 default_factory=cast(Callable[[], List[PostprocessingDescr]], list)
2359 )
2360 """Description of how this output should be postprocessed.
2362 note: `postprocessing` always ends with an 'ensure_dtype' operation.
2363 If not given this is added to cast to this tensor's `data.type`.
2364 """
2366 @model_validator(mode="after")
2367 def _validate_postprocessing_kwargs(self) -> Self:
2368 axes_ids = [a.id for a in self.axes]
2369 for p in self.postprocessing:
2370 kwargs_axes = p.kwargs.get("axes")
2371 if kwargs_axes is None:
2372 continue
2374 if not isinstance(kwargs_axes, collections.abc.Sequence):
2375 raise ValueError(
2376 f"expected `axes` sequence, but got {type(kwargs_axes)}"
2377 )
2379 kwargs_axes_seq: Sequence[Any] = cast(Sequence[Any], kwargs_axes)
2380 if any(a not in axes_ids for a in kwargs_axes_seq):
2381 raise ValueError("`kwargs.axes` needs to be subset of axes ids")
2383 if isinstance(self.data, (NominalOrOrdinalDataDescr, IntervalOrRatioDataDescr)):
2384 dtype = self.data.type
2385 else:
2386 dtype = self.data[0].type
2388 # ensure `postprocessing` ends with `EnsureDtypeDescr` or `BinarizeDescr`
2389 if not self.postprocessing or not isinstance(
2390 self.postprocessing[-1], (EnsureDtypeDescr, BinarizeDescr)
2391 ):
2392 self.postprocessing.append(
2393 EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype=dtype))
2394 )
2395 return self
2398class _OutputTensorConv(
2399 Converter[
2400 _OutputTensorDescr_v0_4,
2401 OutputTensorDescr,
2402 FileSource,
2403 Optional[FileSource],
2404 Mapping[_TensorName_v0_4, Mapping[str, int]],
2405 ]
2406):
2407 def _convert(
2408 self,
2409 src: _OutputTensorDescr_v0_4,
2410 tgt: type[OutputTensorDescr | dict[str, Any]],
2411 test_tensor: FileSource,
2412 sample_tensor: FileSource | None,
2413 size_refs: Mapping[_TensorName_v0_4, Mapping[str, int]],
2414 ) -> OutputTensorDescr | dict[str, Any]:
2415 # TODO: split convert_axes into convert_output_axes and convert_input_axes
2416 axes: list[OutputAxis] = convert_axes( # pyright: ignore[reportAssignmentType]
2417 src.axes,
2418 shape=src.shape,
2419 tensor_type="output",
2420 halo=src.halo,
2421 size_refs=size_refs,
2422 )
2423 data_descr: dict[str, Any] = {"type": src.data_type}
2424 if data_descr["type"] == "bool":
2425 data_descr["values"] = [False, True]
2427 return tgt(
2428 axes=axes,
2429 id=TensorId(str(src.name)),
2430 test_tensor=FileDescr(source=test_tensor),
2431 sample_tensor=(
2432 None if sample_tensor is None else FileDescr(source=sample_tensor)
2433 ),
2434 data=data_descr, # pyright: ignore[reportArgumentType]
2435 postprocessing=[_convert_proc(p, src.axes) for p in src.postprocessing],
2436 )
2439_output_tensor_conv = _OutputTensorConv(_OutputTensorDescr_v0_4, OutputTensorDescr)
2442TensorDescr = Union[InputTensorDescr, OutputTensorDescr]
2445def get_halos(
2446 tensors: Mapping[TensorId, TensorDescr],
2447 /,
2448) -> dict[TensorId, dict[AxisId, tuple[int, int]]]:
2449 """Get all input and output halos from tensor descriptions.
2451 Note:
2452 - Input halos are to be padded
2453 - Output halos are to be cropped
2454 """
2455 halos: dict[TensorId, dict[AxisId, tuple[int, int]]] = {}
2456 for descr in tensors.values():
2457 if isinstance(descr, InputTensorDescr):
2458 continue
2459 for axis in descr.axes:
2460 if not isinstance(axis, WithHalo):
2461 continue
2463 ref_scale = next(
2464 a
2465 for a in tensors[axis.size.tensor_id].axes
2466 if a.id == axis.size.axis_id
2467 ).scale
2469 # set output halo (to be cropped)
2470 halos.setdefault(descr.id, {})[axis.id] = (axis.halo, axis.halo)
2471 # set input halo (to be padded)
2472 pad_width = int(axis.halo / axis.scale * ref_scale)
2473 halos.setdefault(axis.size.tensor_id, {})[axis.size.axis_id] = (
2474 pad_width,
2475 pad_width,
2476 )
2478 return halos
2481def validate_tensors(
2482 tensors: Mapping[TensorId, tuple[TensorDescr, NDArray[Any] | None]],
2483 tensor_origin: Literal[
2484 "source", "test_tensor"
2485 ] = "source", # for more precise error messages
2486 *,
2487 pad_inputs: bool | Literal["allow"] = True,
2488 crop_outputs: bool | Literal["allow"] = True,
2489):
2490 """Validate all inputs (and optionally output tensors) against their tensor descriptions.
2492 Args:
2493 tensors: Mapping of tensor id to a tuple of tensor description and optional numpy array.
2494 tensor_origin: String to use in error messages to indicate the origin of the tensors being validated.
2495 pad_inputs: Wether to apply/allow padding of inputs before shape comparison
2496 crop_outputs: Wether to apply/allow cropping of outputs before shape comparison.
2497 """
2498 all_tensor_axes: dict[TensorId, dict[AxisId, tuple[AnyAxis, int | None]]] = {}
2500 def e_msg_location(d: TensorDescr):
2501 return f"{'inputs' if isinstance(d, InputTensorDescr) else 'outputs'}[{d.id}]"
2503 for descr, array in tensors.values():
2504 if array is None:
2505 axis_sizes = {a.id: None for a in descr.axes}
2506 else:
2507 try:
2508 axis_sizes = descr.get_axis_sizes_for_array(array)
2509 except ValueError as e:
2510 raise ValueError(f"{e_msg_location(descr)} {e}")
2512 all_tensor_axes[descr.id] = {a.id: (a, axis_sizes[a.id]) for a in descr.axes}
2514 # get halos to be padded/cropped to validate against halo-adjusted sizes
2515 io_halos = get_halos({k: v[0] for k, v in tensors.items()})
2517 for descr, array in tensors.values():
2518 if array is None:
2519 continue
2521 if descr.dtype in ("float32", "float64"):
2522 invalid_test_tensor_dtype = array.dtype.name not in (
2523 "float32",
2524 "float64",
2525 "uint8",
2526 "int8",
2527 "uint16",
2528 "int16",
2529 "uint32",
2530 "int32",
2531 "uint64",
2532 "int64",
2533 )
2534 else:
2535 invalid_test_tensor_dtype = array.dtype.name != descr.dtype
2537 if invalid_test_tensor_dtype:
2538 raise ValueError(
2539 f"{tensor_origin} data type '{array.dtype.name}' does not"
2540 + f" match described {e_msg_location(descr)}.dtype '{descr.dtype}'"
2541 )
2543 if array.min() > -1e-4 and array.max() < 1e-4:
2544 raise ValueError(
2545 "Output values are too small for reliable testing."
2546 + f" Values <-1e5 or >=1e5 must be present in {tensor_origin}"
2547 )
2549 for a in descr.axes:
2550 actual_size = all_tensor_axes[descr.id][a.id][1]
2552 if actual_size is None:
2553 continue
2555 if a.size is None:
2556 continue
2558 # add padding width to actual tensor size
2559 total_axis_halo = sum(io_halos.get(descr.id, {}).get(a.id, (0, 0)))
2560 if isinstance(descr, InputTensorDescr):
2561 # pad input halos
2562 actual_size_with_halo = actual_size + total_axis_halo
2563 if pad_inputs is True:
2564 check_sizes = {actual_size_with_halo}
2565 size_hint = " (after padding input halo)"
2566 elif pad_inputs == "allow":
2567 check_sizes = {actual_size, actual_size_with_halo}
2568 size_hint = " (with or without padding input halo)"
2569 elif pad_inputs is False:
2570 check_sizes = {actual_size}
2571 size_hint = ""
2572 else:
2573 assert_never(pad_inputs)
2575 elif isinstance(descr, OutputTensorDescr):
2576 # crop output halos
2577 actual_size_with_halo = max(0, actual_size - total_axis_halo)
2578 if crop_outputs is True:
2579 check_sizes = {actual_size_with_halo}
2580 size_hint = " (after cropping output halo)"
2581 elif crop_outputs == "allow":
2582 check_sizes = {actual_size, actual_size_with_halo}
2583 size_hint = " (with or without cropping output halo)"
2584 elif crop_outputs is False:
2585 check_sizes = {actual_size}
2586 size_hint = ""
2587 else:
2588 assert_never(crop_outputs)
2589 else:
2590 assert_never(descr)
2592 del actual_size # make sure we explicitly use unchanged or halo-adjusted size from here on
2594 if isinstance(a.size, int):
2595 if a.size not in check_sizes:
2596 raise ValueError(
2597 f"{e_msg_location(descr)}.axes[{a.id}]: {tensor_origin} axis "
2598 + f"has incompatible size {check_sizes}{size_hint}, expected {a.size}"
2599 )
2600 elif isinstance(a.size, (ParameterizedSize, DataDependentSize)):
2601 _ = try_all_raise_last(
2602 (partial(a.size.validate_size, s) for s in check_sizes),
2603 f"{e_msg_location(descr)}.axes[{a.id}]: {tensor_origin} axis ",
2604 )
2605 elif isinstance(a.size, SizeReference):
2606 ref_tensor_axes = all_tensor_axes.get(a.size.tensor_id)
2607 if ref_tensor_axes is None:
2608 raise ValueError(
2609 f"{e_msg_location(descr)}.axes[{a.id}].size.tensor_id: Unknown tensor"
2610 + f" reference '{a.size.tensor_id}', available: {list(all_tensor_axes)}"
2611 )
2613 ref_axis, ref_size = ref_tensor_axes.get(a.size.axis_id, (None, None))
2614 if ref_axis is None or ref_size is None:
2615 raise ValueError(
2616 f"{e_msg_location(descr)}.axes[{a.id}].size.axis_id: Unknown tensor axis"
2617 + f" reference '{a.size.tensor_id}.{a.size.axis_id}, available: {list(ref_tensor_axes)}"
2618 )
2620 if a.unit != ref_axis.unit:
2621 raise ValueError(
2622 f"{e_msg_location(descr)}.axes[{a.id}].size: `SizeReference` requires"
2623 + " axis and reference axis to have the same `unit`, but"
2624 + f" {a.unit}!={ref_axis.unit}"
2625 )
2627 if (
2628 expected_size := (
2629 ref_size * ref_axis.scale / a.scale + a.size.offset
2630 )
2631 ) not in check_sizes:
2632 raise ValueError(
2633 f"{e_msg_location(descr)}.{tensor_origin}: axis '{a.id}' of size"
2634 + f" {check_sizes} invalid for referenced size {ref_size};"
2635 + f" expected {expected_size}"
2636 )
2637 else:
2638 assert_never(a.size)
2641FileDescr_dependencies = Annotated[
2642 FileDescr_package,
2643 WithSuffix((".yaml", ".yml"), case_sensitive=True),
2644 Field(examples=[{"source": "environment.yaml"}]),
2645]
2648class _ArchitectureCallableDescr(Node):
2649 callable: Annotated[Identifier, Field(examples=["MyNetworkClass", "get_my_model"])]
2650 """Identifier of the callable that returns a torch.nn.Module instance."""
2652 kwargs: dict[str, YamlValue] = Field(
2653 default_factory=cast(Callable[[], Dict[str, YamlValue]], dict)
2654 )
2655 """key word arguments for the `callable`"""
2658class ArchitectureFromFileDescr(_ArchitectureCallableDescr, FileDescr):
2659 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
2660 """Architecture source file"""
2662 @model_serializer(mode="wrap", when_used="unless-none")
2663 def _serialize(self, nxt: SerializerFunctionWrapHandler, info: SerializationInfo):
2664 return package_file_descr_serializer(self, nxt, info)
2667class ArchitectureFromLibraryDescr(_ArchitectureCallableDescr):
2668 import_from: str
2669 """Where to import the callable from, i.e. `from <import_from> import <callable>`"""
2672class _ArchFileConv(
2673 Converter[
2674 _CallableFromFile_v0_4,
2675 ArchitectureFromFileDescr,
2676 Optional[Sha256],
2677 Dict[str, Any],
2678 ]
2679):
2680 def _convert(
2681 self,
2682 src: _CallableFromFile_v0_4,
2683 tgt: type[ArchitectureFromFileDescr | dict[str, Any]],
2684 sha256: Sha256 | None,
2685 kwargs: dict[str, Any],
2686 ) -> ArchitectureFromFileDescr | dict[str, Any]:
2687 if src.startswith("http") and src.count(":") == 2:
2688 http, source, callable_ = src.split(":")
2689 source = f"{http}:{source}"
2690 elif not src.startswith("http") and src.count(":") == 1:
2691 source, callable_ = src.split(":")
2692 else:
2693 source = str(src)
2694 callable_ = str(src)
2695 return tgt(
2696 callable=Identifier(callable_),
2697 source=cast(FileSource, source),
2698 sha256=sha256,
2699 kwargs=kwargs,
2700 )
2703_arch_file_conv = _ArchFileConv(_CallableFromFile_v0_4, ArchitectureFromFileDescr)
2706class _ArchLibConv(
2707 Converter[
2708 _CallableFromDepencency_v0_4, ArchitectureFromLibraryDescr, Dict[str, Any]
2709 ]
2710):
2711 def _convert(
2712 self,
2713 src: _CallableFromDepencency_v0_4,
2714 tgt: type[ArchitectureFromLibraryDescr | dict[str, Any]],
2715 kwargs: dict[str, Any],
2716 ) -> ArchitectureFromLibraryDescr | dict[str, Any]:
2717 *mods, callable_ = src.split(".")
2718 import_from = ".".join(mods)
2719 return tgt(
2720 import_from=import_from, callable=Identifier(callable_), kwargs=kwargs
2721 )
2724_arch_lib_conv = _ArchLibConv(
2725 _CallableFromDepencency_v0_4, ArchitectureFromLibraryDescr
2726)
2729class WeightsEntryDescrBase(FileDescr):
2730 type: ClassVar[WeightsFormat]
2731 weights_format_name: ClassVar[str] # human readable
2733 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
2734 """Source of the weights file."""
2736 authors: list[Author] | None = None
2737 """Authors
2738 Either the person(s) that have trained this model resulting in the original weights file.
2739 (If this is the initial weights entry, i.e. it does not have a `parent`)
2740 Or the person(s) who have converted the weights to this weights format.
2741 (If this is a child weight, i.e. it has a `parent` field)
2742 """
2744 parent: Annotated[WeightsFormat | None, Field(examples=["pytorch_state_dict"])] = (
2745 None
2746 )
2747 """The source weights these weights were converted from.
2748 For example, if a model's weights were converted from the `pytorch_state_dict` format to `torchscript`,
2749 The `pytorch_state_dict` weights entry has no `parent` and is the parent of the `torchscript` weights.
2750 All weight entries except one (the initial set of weights resulting from training the model),
2751 need to have this field."""
2753 comment: str = ""
2754 """A comment about this weights entry, for example how these weights were created."""
2756 @model_validator(mode="after")
2757 def _validate(self) -> Self:
2758 if self.type == self.parent:
2759 raise ValueError("Weights entry can't be it's own parent.")
2761 return self
2763 @model_serializer(mode="wrap", when_used="unless-none")
2764 def _serialize(self, nxt: SerializerFunctionWrapHandler, info: SerializationInfo):
2765 return package_file_descr_serializer(self, nxt, info)
2768class KerasHdf5WeightsDescr(WeightsEntryDescrBase):
2769 type: ClassVar[WeightsFormat] = "keras_hdf5"
2770 weights_format_name: ClassVar[str] = "Keras HDF5"
2771 tensorflow_version: Version
2772 """TensorFlow version used to create these weights."""
2775class KerasV3WeightsDescr(WeightsEntryDescrBase):
2776 type: ClassVar[WeightsFormat] = "keras_v3"
2777 weights_format_name: ClassVar[str] = "Keras v3"
2778 keras_version: Annotated[Version, Ge(Version(3))]
2779 """Keras version used to create these weights."""
2780 backend: tuple[Literal["tensorflow", "jax", "torch"], Version]
2781 """Keras backend used to create these weights."""
2782 source: Annotated[
2783 FileSource,
2784 AfterValidator(wo_special_file_name),
2785 WithSuffix(".keras", case_sensitive=True),
2786 ]
2787 """Source of the .keras weights file."""
2790FileDescr_external_data = Annotated[
2791 FileDescr_package,
2792 WithSuffix(".data", case_sensitive=True),
2793 Field(examples=[{"source": "weights.onnx.data"}]),
2794]
2797class OnnxWeightsDescr(WeightsEntryDescrBase):
2798 type: ClassVar[WeightsFormat] = "onnx"
2799 weights_format_name: ClassVar[str] = "ONNX"
2800 opset_version: Annotated[int, Ge(7)]
2801 """ONNX opset version"""
2803 external_data: FileDescr_external_data | None = None
2804 """Source of the external ONNX data file holding the weights.
2805 (If present **source** holds the ONNX architecture without weights)."""
2807 @model_validator(mode="after")
2808 def _validate_external_data_unique_file_name(self) -> Self:
2809 if self.external_data is not None and (
2810 extract_file_name(self.source)
2811 == extract_file_name(self.external_data.source)
2812 ):
2813 raise ValueError(
2814 f"ONNX `external_data` file name '{extract_file_name(self.external_data.source)}'"
2815 + " must be different from ONNX `source` file name."
2816 )
2818 return self
2821class PytorchStateDictWeightsDescr(WeightsEntryDescrBase):
2822 type: ClassVar[WeightsFormat] = "pytorch_state_dict"
2823 weights_format_name: ClassVar[str] = "Pytorch State Dict"
2824 architecture: ArchitectureFromFileDescr | ArchitectureFromLibraryDescr
2825 pytorch_version: Version
2826 """Version of the PyTorch library used.
2827 If `architecture.depencencies` is specified it has to include pytorch and any version pinning has to be compatible.
2828 """
2829 dependencies: FileDescr_dependencies | None = None
2830 """Custom depencies beyond pytorch described in a Conda environment file.
2831 Allows to specify custom dependencies, see conda docs:
2832 - [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)
2833 - [Creating an environment file manually](https://conda.io/projects/conda/en/latest/user-guide/tasks/manage-environments.html#creating-an-environment-file-manually)
2835 The conda environment file should include pytorch and any version pinning has to be compatible with
2836 **pytorch_version**.
2837 """
2838 strict: bool = True
2839 """Whether to allow missing or unexpected keys or to be strict about the architecture matching the state dict weights."""
2842class TensorflowJsWeightsDescr(WeightsEntryDescrBase):
2843 type: ClassVar[WeightsFormat] = "tensorflow_js"
2844 weights_format_name: ClassVar[str] = "Tensorflow.js"
2845 tensorflow_version: Version
2846 """Version of the TensorFlow library used."""
2848 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
2849 """The multi-file weights.
2850 All required files/folders should be a zip archive."""
2853class TensorflowSavedModelBundleWeightsDescr(WeightsEntryDescrBase):
2854 type: ClassVar[WeightsFormat] = "tensorflow_saved_model_bundle"
2855 weights_format_name: ClassVar[str] = "Tensorflow Saved Model"
2856 tensorflow_version: Version
2857 """Version of the TensorFlow library used."""
2859 dependencies: FileDescr_dependencies | None = None
2860 """Custom dependencies beyond tensorflow.
2861 Should include tensorflow and any version pinning has to be compatible with **tensorflow_version**."""
2863 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
2864 """The multi-file weights.
2865 All required files/folders should be a zip archive."""
2868class TorchscriptWeightsDescr(WeightsEntryDescrBase):
2869 type: ClassVar[WeightsFormat] = "torchscript"
2870 weights_format_name: ClassVar[str] = "TorchScript"
2871 pytorch_version: Version
2872 """Version of the PyTorch library used."""
2875SpecificWeightsDescr = Union[
2876 KerasHdf5WeightsDescr,
2877 KerasV3WeightsDescr,
2878 OnnxWeightsDescr,
2879 PytorchStateDictWeightsDescr,
2880 TensorflowJsWeightsDescr,
2881 TensorflowSavedModelBundleWeightsDescr,
2882 TorchscriptWeightsDescr,
2883]
2886class WeightsDescr(Node):
2887 keras_hdf5: KerasHdf5WeightsDescr | None = None
2888 keras_v3: KerasV3WeightsDescr | None = None
2889 onnx: OnnxWeightsDescr | None = None
2890 pytorch_state_dict: PytorchStateDictWeightsDescr | None = None
2891 tensorflow_js: TensorflowJsWeightsDescr | None = None
2892 tensorflow_saved_model_bundle: TensorflowSavedModelBundleWeightsDescr | None = None
2893 torchscript: TorchscriptWeightsDescr | None = None
2895 @model_validator(mode="after")
2896 def check_entries(self) -> Self:
2897 entries = {wtype for wtype, entry in self if entry is not None}
2899 if not entries:
2900 raise ValueError("Missing weights entry")
2902 entries_wo_parent = {
2903 wtype
2904 for wtype, entry in self
2905 if entry is not None and hasattr(entry, "parent") and entry.parent is None
2906 }
2907 if len(entries_wo_parent) != 1:
2908 issue_warning(
2909 "Exactly one weights entry may not specify the `parent` field (got"
2910 + " {value}). That entry is considered the original set of model weights."
2911 + " Other weight formats are created through conversion of the orignal or"
2912 + " already converted weights. They have to reference the weights format"
2913 + " they were converted from as their `parent`.",
2914 value=len(entries_wo_parent),
2915 field="weights",
2916 )
2918 for wtype, entry in self:
2919 if entry is None:
2920 continue
2922 assert hasattr(entry, "type")
2923 assert hasattr(entry, "parent")
2924 assert wtype == entry.type
2925 if (
2926 entry.parent is not None and entry.parent not in entries
2927 ): # self reference checked for `parent` field
2928 raise ValueError(
2929 f"`weights.{wtype}.parent={entry.parent} not in specified weight"
2930 + f" formats: {entries}"
2931 )
2933 return self
2935 def __getitem__(
2936 self,
2937 key: WeightsFormat,
2938 ):
2939 if key == "keras_hdf5":
2940 ret = self.keras_hdf5
2941 elif key == "keras_v3":
2942 ret = self.keras_v3
2943 elif key == "onnx":
2944 ret = self.onnx
2945 elif key == "pytorch_state_dict":
2946 ret = self.pytorch_state_dict
2947 elif key == "tensorflow_js":
2948 ret = self.tensorflow_js
2949 elif key == "tensorflow_saved_model_bundle":
2950 ret = self.tensorflow_saved_model_bundle
2951 elif key == "torchscript":
2952 ret = self.torchscript
2953 else:
2954 raise KeyError(key)
2956 if ret is None:
2957 raise KeyError(key)
2959 return ret
2961 @overload
2962 def __setitem__(
2963 self, key: Literal["keras_hdf5"], value: KerasHdf5WeightsDescr | None
2964 ) -> None: ...
2965 @overload
2966 def __setitem__(
2967 self, key: Literal["keras_v3"], value: KerasV3WeightsDescr | None
2968 ) -> None: ...
2969 @overload
2970 def __setitem__(
2971 self, key: Literal["onnx"], value: OnnxWeightsDescr | None
2972 ) -> None: ...
2973 @overload
2974 def __setitem__(
2975 self,
2976 key: Literal["pytorch_state_dict"],
2977 value: PytorchStateDictWeightsDescr | None,
2978 ) -> None: ...
2979 @overload
2980 def __setitem__(
2981 self, key: Literal["tensorflow_js"], value: TensorflowJsWeightsDescr | None
2982 ) -> None: ...
2983 @overload
2984 def __setitem__(
2985 self,
2986 key: Literal["tensorflow_saved_model_bundle"],
2987 value: TensorflowSavedModelBundleWeightsDescr | None,
2988 ) -> None: ...
2989 @overload
2990 def __setitem__(
2991 self, key: Literal["torchscript"], value: TorchscriptWeightsDescr | None
2992 ) -> None: ...
2994 def __setitem__(
2995 self,
2996 key: WeightsFormat,
2997 value: SpecificWeightsDescr | None,
2998 ):
2999 if key == "keras_hdf5":
3000 if value is not None and not isinstance(value, KerasHdf5WeightsDescr):
3001 raise TypeError(
3002 f"Expected KerasHdf5WeightsDescr or None for key 'keras_hdf5', got {type(value)}"
3003 )
3004 self.keras_hdf5 = value
3005 elif key == "keras_v3":
3006 if value is not None and not isinstance(value, KerasV3WeightsDescr):
3007 raise TypeError(
3008 f"Expected KerasV3WeightsDescr or None for key 'keras_v3', got {type(value)}"
3009 )
3010 self.keras_v3 = value
3011 elif key == "onnx":
3012 if value is not None and not isinstance(value, OnnxWeightsDescr):
3013 raise TypeError(
3014 f"Expected OnnxWeightsDescr or None for key 'onnx', got {type(value)}"
3015 )
3016 self.onnx = value
3017 elif key == "pytorch_state_dict":
3018 if value is not None and not isinstance(
3019 value, PytorchStateDictWeightsDescr
3020 ):
3021 raise TypeError(
3022 f"Expected PytorchStateDictWeightsDescr or None for key 'pytorch_state_dict', got {type(value)}"
3023 )
3024 self.pytorch_state_dict = value
3025 elif key == "tensorflow_js":
3026 if value is not None and not isinstance(value, TensorflowJsWeightsDescr):
3027 raise TypeError(
3028 f"Expected TensorflowJsWeightsDescr or None for key 'tensorflow_js', got {type(value)}"
3029 )
3030 self.tensorflow_js = value
3031 elif key == "tensorflow_saved_model_bundle":
3032 if value is not None and not isinstance(
3033 value, TensorflowSavedModelBundleWeightsDescr
3034 ):
3035 raise TypeError(
3036 f"Expected TensorflowSavedModelBundleWeightsDescr or None for key 'tensorflow_saved_model_bundle', got {type(value)}"
3037 )
3038 self.tensorflow_saved_model_bundle = value
3039 elif key == "torchscript":
3040 if value is not None and not isinstance(value, TorchscriptWeightsDescr):
3041 raise TypeError(
3042 f"Expected TorchscriptWeightsDescr or None for key 'torchscript', got {type(value)}"
3043 )
3044 self.torchscript = value
3045 else:
3046 raise KeyError(key)
3048 @property
3049 def available_formats(self) -> dict[WeightsFormat, SpecificWeightsDescr]:
3050 return {
3051 **({} if self.keras_hdf5 is None else {"keras_hdf5": self.keras_hdf5}),
3052 **({} if self.keras_v3 is None else {"keras_v3": self.keras_v3}),
3053 **({} if self.onnx is None else {"onnx": self.onnx}),
3054 **(
3055 {}
3056 if self.pytorch_state_dict is None
3057 else {"pytorch_state_dict": self.pytorch_state_dict}
3058 ),
3059 **(
3060 {}
3061 if self.tensorflow_js is None
3062 else {"tensorflow_js": self.tensorflow_js}
3063 ),
3064 **(
3065 {}
3066 if self.tensorflow_saved_model_bundle is None
3067 else {
3068 "tensorflow_saved_model_bundle": self.tensorflow_saved_model_bundle
3069 }
3070 ),
3071 **({} if self.torchscript is None else {"torchscript": self.torchscript}),
3072 }
3074 @property
3075 def missing_formats(self) -> set[WeightsFormat]:
3076 return {
3077 wf for wf in get_args(WeightsFormat) if wf not in self.available_formats
3078 }
3081class LinkedModel(LinkedResourceBase):
3082 """Reference to a bioimage.io model."""
3084 id: ModelId
3085 """A valid model `id` from the bioimage.io collection."""
3088class _DataDepSize(NamedTuple):
3089 min: StrictInt
3090 max: StrictInt | None
3093class _AxisSizes(NamedTuple):
3094 """the lenghts of all axes of model inputs and outputs"""
3096 inputs: dict[tuple[TensorId, AxisId], int]
3097 outputs: dict[tuple[TensorId, AxisId], int | _DataDepSize]
3100class _TensorSizes(NamedTuple):
3101 """_AxisSizes as nested dicts"""
3103 inputs: dict[TensorId, dict[AxisId, int]]
3104 outputs: dict[TensorId, dict[AxisId, int | _DataDepSize]]
3107class ReproducibilityTolerance(Node, extra="allow"):
3108 """Describes what small numerical differences -- if any -- may be tolerated
3109 in the generated output when executing in different environments.
3111 A tensor element *output* is considered mismatched to the **test_tensor** if
3112 abs(*output* - **test_tensor**) > **absolute_tolerance** + **relative_tolerance** * abs(**test_tensor**).
3113 (Internally we call [numpy.testing.assert_allclose](https://numpy.org/doc/stable/reference/generated/numpy.testing.assert_allclose.html).)
3115 Motivation:
3116 For testing we can request the respective deep learning frameworks to be as
3117 reproducible as possible by setting seeds and chosing deterministic algorithms,
3118 but differences in operating systems, available hardware and installed drivers
3119 may still lead to numerical differences.
3120 """
3122 relative_tolerance: RelativeTolerance = 1e-3
3123 """Maximum relative tolerance of reproduced test tensor."""
3125 absolute_tolerance: AbsoluteTolerance = 1e-3
3126 """Maximum absolute tolerance of reproduced test tensor."""
3128 mismatched_elements_per_million: MismatchedElementsPerMillion = 100
3129 """Maximum number of mismatched elements/pixels per million to tolerate."""
3131 output_ids: Sequence[TensorId] = ()
3132 """Limits the output tensor IDs these reproducibility details apply to."""
3134 weights_formats: Sequence[WeightsFormat] = ()
3135 """Limits the weights formats these details apply to."""
3138class BiasRisksLimitations(Node, extra="allow"):
3139 """Known biases, risks, technical limitations, and recommendations for model use."""
3141 known_biases: str = dedent("""\
3142 In general bioimage models may suffer from biases caused by:
3144 - Imaging protocol dependencies
3145 - Use of a specific cell type
3146 - Species-specific training data limitations
3148 """)
3149 """Biases in training data or model behavior."""
3151 risks: str = dedent("""\
3152 Common risks in bioimage analysis include:
3154 - Erroneously assuming generalization to unseen experimental conditions
3155 - Trusting (overconfident) model outputs without validation
3156 - Misinterpretation of results
3158 """)
3159 """Potential risks in the context of bioimage analysis."""
3161 limitations: str | None = None
3162 """Technical limitations and failure modes."""
3164 recommendations: str = "Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model."
3165 """Mitigation strategies regarding `known_biases`, `risks`, and `limitations`, as well as applicable best practices.
3167 Consider:
3168 - How to use a validation dataset?
3169 - How to manually validate?
3170 - Feasibility of domain adaptation for different experimental setups?
3172 """
3174 def format_md(self) -> str:
3175 if self.limitations is None:
3176 limitations_header = ""
3177 else:
3178 limitations_header = "## Limitations\n\n"
3180 return f"""# Bias, Risks, and Limitations
3182{self.known_biases}
3184{self.risks}
3186{limitations_header}{self.limitations or ""}
3188## Recommendations
3190{self.recommendations}
3192"""
3195class TrainingDetails(Node, extra="allow"):
3196 training_preprocessing: str | None = None
3197 """Detailed image preprocessing steps during model training:
3199 Mention:
3200 - *Normalization methods*
3201 - *Augmentation strategies*
3202 - *Resizing/resampling procedures*
3203 - *Artifact handling*
3205 """
3207 training_epochs: float | None = None
3208 """Number of training epochs."""
3210 training_batch_size: float | None = None
3211 """Batch size used in training."""
3213 initial_learning_rate: float | None = None
3214 """Initial learning rate used in training."""
3216 learning_rate_schedule: str | None = None
3217 """Learning rate schedule used in training."""
3219 loss_function: str | None = None
3220 """Loss function used in training, e.g. nn.MSELoss."""
3222 loss_function_kwargs: dict[str, YamlValue] = Field(
3223 default_factory=cast(Callable[[], Dict[str, YamlValue]], dict)
3224 )
3225 """key word arguments for the `loss_function`"""
3227 optimizer: str | None = None
3228 """optimizer, e.g. torch.optim.Adam"""
3230 optimizer_kwargs: dict[str, YamlValue] = Field(
3231 default_factory=cast(Callable[[], Dict[str, YamlValue]], dict)
3232 )
3233 """key word arguments for the `optimizer`"""
3235 regularization: str | None = None
3236 """Regularization techniques used during training, e.g. drop-out or weight decay."""
3238 training_duration: float | None = None
3239 """Total training duration in hours."""
3242class Evaluation(Node, extra="allow"):
3243 model_id: ModelId | None = None
3244 """Model being evaluated."""
3246 dataset_id: DatasetId
3247 """Dataset used for evaluation."""
3249 dataset_source: HttpUrl
3250 """Source of the dataset."""
3252 dataset_role: Literal["train", "validation", "test", "independent", "unknown"]
3253 """Role of the dataset used for evaluation.
3255 - `train`: dataset was (part of) the training data
3256 - `validation`: dataset was (part of) the validation data used during training, e.g. used for model selection or hyperparameter tuning
3257 - `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
3258 - `independent`: dataset is entirely independent test data; not used during training or validation, and acquired from a different source/distribution than training data
3259 - `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.
3260 """
3262 sample_count: int
3263 """Number of evaluated samples."""
3265 evaluation_factors: list[Annotated[str, MaxLen(16)]]
3266 """(Abbreviations of) each evaluation factor.
3268 Evaluation factors are criteria along which model performance is evaluated, e.g. different image conditions
3269 like 'low SNR', 'high cell density', or different biological conditions like 'cell type A', 'cell type B'.
3270 An 'overall' factor may be included to summarize performance across all conditions.
3271 """
3273 evaluation_factors_long: list[str]
3274 """Descriptions (long form) of each evaluation factor."""
3276 metrics: list[Annotated[str, MaxLen(16)]]
3277 """(Abbreviations of) metrics used for evaluation."""
3279 metrics_long: list[str]
3280 """Description of each metric used."""
3282 @model_validator(mode="after")
3283 def _validate_list_lengths(self) -> Self:
3284 if len(self.evaluation_factors) != len(self.evaluation_factors_long):
3285 raise ValueError(
3286 "`evaluation_factors` and `evaluation_factors_long` must have the same length"
3287 )
3289 if len(self.metrics) != len(self.metrics_long):
3290 raise ValueError("`metrics` and `metrics_long` must have the same length")
3292 if len(self.results) != len(self.metrics):
3293 raise ValueError("`results` must have the same number of rows as `metrics`")
3295 for row in self.results:
3296 if len(row) != len(self.evaluation_factors):
3297 raise ValueError(
3298 "`results` must have the same number of columns (in every row) as `evaluation_factors`"
3299 )
3301 return self
3303 results: list[list[str | float | int]]
3304 """Results for each metric (rows; outer list) and each evaluation factor (columns; inner list)."""
3306 results_summary: str | None = None
3307 """Interpretation of results for general audience.
3309 Consider:
3310 - Overall model performance
3311 - Comparison to existing methods
3312 - Limitations and areas for improvement
3314"""
3316 def format_md(self):
3317 results_header = ["Metric"] + self.evaluation_factors
3318 results_table_cells = [results_header, ["---"] * len(results_header)] + [
3319 [metric] + [str(r) for r in row]
3320 for metric, row in zip(self.metrics, self.results)
3321 ]
3323 results_table = "".join(
3324 "| " + " | ".join(row) + " |\n" for row in results_table_cells
3325 )
3326 factors = "".join(
3327 f"\n - {ef}: {efl}"
3328 for ef, efl in zip(self.evaluation_factors, self.evaluation_factors_long)
3329 )
3330 metrics = "".join(
3331 f"\n - {em}: {eml}" for em, eml in zip(self.metrics, self.metrics_long)
3332 )
3334 return f"""## Testing Data, Factors & Metrics
3336Evaluation of {self.model_id or "this"} model on the {self.dataset_id} dataset (dataset role: {self.dataset_role}).
3338### Testing Data
3340- **Source:** [{self.dataset_id}]({self.dataset_source})
3341- **Size:** {self.sample_count} evaluated samples
3343### Factors
3344{factors}
3346### Metrics
3347{metrics}
3349## Results
3351### Quantitative Results
3353{results_table}
3355### Summary
3357{self.results_summary or "missing"}
3359"""
3362class EnvironmentalImpact(Node, extra="allow"):
3363 """Environmental considerations for model training and deployment.
3365 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).
3366 """
3368 hardware_type: str | None = None
3369 """GPU/CPU specifications"""
3371 hours_used: float | None = None
3372 """Total compute hours"""
3374 cloud_provider: str | None = None
3375 """If applicable"""
3377 compute_region: str | None = None
3378 """Geographic location"""
3380 co2_emitted: float | None = None
3381 """kg CO2 equivalent
3383 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).
3384 """
3386 def format_md(self):
3387 """Filled Markdown template section following [Hugging Face Model Card Template](https://huggingface.co/docs/hub/en/model-card-annotated)."""
3388 if self == self.__class__():
3389 return ""
3391 ret = "# Environmental Impact\n\n"
3392 if self.hardware_type is not None:
3393 ret += f"- **Hardware Type:** {self.hardware_type}\n"
3394 if self.hours_used is not None:
3395 ret += f"- **Hours used:** {self.hours_used}\n"
3396 if self.cloud_provider is not None:
3397 ret += f"- **Cloud Provider:** {self.cloud_provider}\n"
3398 if self.compute_region is not None:
3399 ret += f"- **Compute Region:** {self.compute_region}\n"
3400 if self.co2_emitted is not None:
3401 ret += f"- **Carbon Emitted:** {self.co2_emitted} kg CO2e\n"
3403 return ret + "\n"
3406class BioimageioConfig(Node, extra="allow"):
3407 reproducibility_tolerance: Sequence[ReproducibilityTolerance] = ()
3408 """Tolerances to allow when reproducing the model's test outputs
3409 from the model's test inputs.
3410 Only the first entry matching tensor id and weights format is considered.
3411 """
3413 funded_by: str | None = None
3414 """Funding agency, grant number if applicable"""
3416 architecture_type: Annotated[str, MaxLen(32)] | None = (
3417 None # TODO: add to differentiated tags
3418 )
3419 """Model architecture type, e.g., 3D U-Net, ResNet, transformer"""
3421 architecture_description: str | None = None
3422 """Text description of model architecture."""
3424 modality: str | None = None # TODO: add to differentiated tags
3425 """Input modality, e.g., fluorescence microscopy, electron microscopy"""
3427 target_structure: list[str] = Field( # TODO: add to differentiated tags
3428 default_factory=cast(Callable[[], List[str]], list)
3429 )
3430 """Biological structure(s) the model is designed to analyze, e.g., nuclei, mitochondria, cells"""
3432 task: str | None = None # TODO: add to differentiated tags
3433 """Bioimage-specific task type, e.g., segmentation, classification, detection, denoising"""
3435 new_version: ModelId | None = None
3436 """A new version of this model exists with a different model id."""
3438 out_of_scope_use: str | None = None
3439 """Describe how the model may be misused in bioimage analysis contexts and what users should **not** do with the model."""
3441 bias_risks_limitations: BiasRisksLimitations = Field(
3442 default_factory=BiasRisksLimitations.model_construct
3443 )
3444 """Description of known bias, risks, and technical limitations for in-scope model use."""
3446 model_parameter_count: int | None = None
3447 """Total number of model parameters."""
3449 training: TrainingDetails = Field(default_factory=TrainingDetails.model_construct)
3450 """Details on how the model was trained."""
3452 inference_time: str | None = None
3453 """Average inference time per image/tile. Specify hardware and image size. Multiple examples can be given."""
3455 memory_requirements_inference: str | None = None
3456 """GPU memory needed for inference. Multiple examples with different image size can be given."""
3458 memory_requirements_training: str | None = None
3459 """GPU memory needed for training. Multiple examples with different image/batch sizes can be given."""
3461 evaluations: list[Evaluation] = Field(
3462 default_factory=cast(Callable[[], List[Evaluation]], list)
3463 )
3464 """Quantitative model evaluations.
3466 Note:
3467 At the moment we recommend to include only a single test dataset
3468 (with evaluation factors that may mark subsets of the dataset)
3469 to avoid confusion and make the presentation of results cleaner.
3470 """
3472 environmental_impact: EnvironmentalImpact = Field(
3473 default_factory=EnvironmentalImpact.model_construct
3474 )
3475 """Environmental considerations for model training and deployment"""
3478class Config(Node, extra="allow"):
3479 bioimageio: BioimageioConfig = Field(
3480 default_factory=BioimageioConfig.model_construct
3481 )
3482 stardist: YamlValue = None
3485class ModelDescr(GenericModelDescrBase):
3486 """Specification of the fields used in a bioimage.io-compliant RDF to describe AI models with pretrained weights.
3487 These fields are typically stored in a YAML file which we call a model resource description file (model RDF).
3488 """
3490 implemented_format_version: ClassVar[Literal["0.5.14"]] = "0.5.14"
3491 if TYPE_CHECKING:
3492 format_version: Literal["0.5.14"] = "0.5.14"
3493 else:
3494 format_version: Literal["0.5.14"]
3495 """Version of the bioimage.io model description specification used.
3496 When creating a new model always use the latest micro/patch version described here.
3497 The `format_version` is important for any consumer software to understand how to parse the fields.
3498 """
3500 implemented_type: ClassVar[Literal["model"]] = "model"
3501 if TYPE_CHECKING:
3502 type: Literal["model"] = "model"
3503 else:
3504 type: Literal["model"]
3505 """Specialized resource type 'model'"""
3507 id: ModelId | None = None
3508 """bioimage.io-wide unique resource identifier
3509 assigned by bioimage.io; version **un**specific."""
3511 authors: FAIR[list[Author]] = Field(
3512 default_factory=cast(Callable[[], List[Author]], list)
3513 )
3514 """The authors are the creators of the model RDF and the primary points of contact."""
3516 documentation: FAIR[FileDescr_documentation | None] = None
3517 """Additional model documentation.
3518 The recommended documentation source file name is `README.md`. An `.md` suffix is mandatory.
3519 The documentation should include a '#[#] Validation' (sub)section
3520 with details on how to quantitatively validate the model on unseen data."""
3522 @field_validator("documentation", mode="after")
3523 @classmethod
3524 def _validate_documentation(cls, value: FileDescr | None) -> FileDescr | None:
3525 if not get_validation_context().perform_io_checks or value is None:
3526 return value
3528 doc_reader = get_reader(value)
3529 doc_content = doc_reader.read().decode(encoding="utf-8")
3530 if not re.search("#.*[vV]alidation", doc_content):
3531 issue_warning(
3532 "No '# Validation' (sub)section found in {value}.",
3533 value=value,
3534 field="documentation",
3535 )
3537 return value
3539 inputs: NotEmpty[Sequence[InputTensorDescr]]
3540 """Describes the input tensors expected by this model."""
3542 @field_validator("inputs", mode="after")
3543 @classmethod
3544 def _validate_input_axes(
3545 cls, inputs: Sequence[InputTensorDescr]
3546 ) -> Sequence[InputTensorDescr]:
3547 input_size_refs = cls._get_axes_with_independent_size(inputs)
3549 for i, ipt in enumerate(inputs):
3550 valid_independent_refs: dict[
3551 tuple[TensorId, AxisId],
3552 tuple[TensorDescr, AnyAxis, int | ParameterizedSize],
3553 ] = {
3554 **{
3555 (ipt.id, a.id): (ipt, a, a.size)
3556 for a in ipt.axes
3557 if not isinstance(a, BatchAxis)
3558 and isinstance(a.size, (int, ParameterizedSize))
3559 },
3560 **input_size_refs,
3561 }
3562 for a, ax in enumerate(ipt.axes):
3563 cls._validate_axis(
3564 "inputs",
3565 i=i,
3566 tensor_id=ipt.id,
3567 a=a,
3568 axis=ax,
3569 valid_independent_refs=valid_independent_refs,
3570 )
3571 return inputs
3573 @staticmethod
3574 def _validate_axis(
3575 field_name: str,
3576 i: int,
3577 tensor_id: TensorId,
3578 a: int,
3579 axis: AnyAxis,
3580 valid_independent_refs: dict[
3581 tuple[TensorId, AxisId],
3582 tuple[TensorDescr, AnyAxis, int | ParameterizedSize],
3583 ],
3584 ):
3585 if isinstance(axis, BatchAxis) or isinstance(
3586 axis.size, (int, ParameterizedSize, DataDependentSize)
3587 ):
3588 return
3589 elif not isinstance(axis.size, SizeReference):
3590 assert_never(axis.size)
3592 # validate axis.size SizeReference
3593 ref = (axis.size.tensor_id, axis.size.axis_id)
3594 if ref not in valid_independent_refs:
3595 raise ValueError(
3596 "Invalid tensor axis reference at"
3597 + f" {field_name}[{i}].axes[{a}].size: {axis.size}."
3598 )
3599 if ref == (tensor_id, axis.id):
3600 raise ValueError(
3601 "Self-referencing not allowed for"
3602 + f" {field_name}[{i}].axes[{a}].size: {axis.size}"
3603 )
3604 if axis.type == "channel":
3605 if valid_independent_refs[ref][1].type != "channel":
3606 raise ValueError(
3607 "A channel axis' size may only reference another fixed size"
3608 + " channel axis."
3609 )
3610 if isinstance(axis.channel_names, str) and "{i}" in axis.channel_names:
3611 ref_size = valid_independent_refs[ref][2]
3612 assert isinstance(ref_size, int), (
3613 "channel axis ref (another channel axis) has to specify fixed"
3614 + " size"
3615 )
3616 generated_channel_names = [
3617 axis.channel_names.format(i=i) for i in range(1, ref_size + 1)
3618 ]
3619 axis.channel_names = generated_channel_names
3621 if (ax_unit := getattr(axis, "unit", None)) != (
3622 ref_unit := getattr(valid_independent_refs[ref][1], "unit", None)
3623 ):
3624 raise ValueError(
3625 "The units of an axis and its reference axis need to match, but"
3626 + f" '{ax_unit}' != '{ref_unit}'."
3627 )
3628 ref_axis = valid_independent_refs[ref][1]
3629 if isinstance(ref_axis, BatchAxis):
3630 raise ValueError(
3631 f"Invalid reference axis '{ref_axis.id}' for {tensor_id}.{axis.id}"
3632 + " (a batch axis is not allowed as reference)."
3633 )
3635 if isinstance(axis, WithHalo):
3636 min_size = axis.size.get_size(axis, ref_axis, n=0)
3637 if (min_size - 2 * axis.halo) < 1:
3638 raise ValueError(
3639 f"axis {axis.id} with minimum size {min_size} is too small for halo"
3640 + f" {axis.halo}."
3641 )
3643 ref_halo = axis.halo * axis.scale / ref_axis.scale
3644 if ref_halo != int(ref_halo):
3645 raise ValueError(
3646 f"Inferred halo for {'.'.join(ref)} is not an integer ({ref_halo} ="
3647 + f" {tensor_id}.{axis.id}.halo {axis.halo}"
3648 + f" * {tensor_id}.{axis.id}.scale {axis.scale}"
3649 + f" / {'.'.join(ref)}.scale {ref_axis.scale})."
3650 )
3652 def validate_input_tensors(
3653 self,
3654 sources: Sequence[NDArray[Any]] | Mapping[TensorId, NDArray[Any] | None],
3655 *,
3656 pad_inputs: bool | Literal["allow"] = True,
3657 crop_outputs: bool | Literal["allow"] = True,
3658 ) -> Mapping[TensorId, NDArray[Any] | None]:
3659 """Check if the given input tensors match the model's input tensor descriptions.
3660 This includes checks of tensor shapes and dtypes, but not of the actual values.
3661 """
3662 if not isinstance(sources, collections.abc.Mapping):
3663 sources = {descr.id: tensor for descr, tensor in zip(self.inputs, sources)}
3665 tensors = {
3666 **{descr.id: (descr, sources.get(descr.id)) for descr in self.inputs},
3667 **{ # outputs are required for halo
3668 descr.id: (descr, None) for descr in self.outputs
3669 },
3670 }
3671 validate_tensors(tensors, pad_inputs=pad_inputs, crop_outputs=crop_outputs)
3673 return sources
3675 @model_validator(mode="after")
3676 def _validate_test_tensors(self) -> Self:
3677 if not get_validation_context().perform_io_checks:
3678 return self
3680 test_inputs = {
3681 descr.id: (
3682 descr,
3683 None if descr.test_tensor is None else load_array(descr.test_tensor),
3684 )
3685 for descr in self.inputs
3686 }
3687 test_outputs = {
3688 descr.id: (
3689 descr,
3690 None if descr.test_tensor is None else load_array(descr.test_tensor),
3691 )
3692 for descr in self.outputs
3693 }
3695 validate_tensors(
3696 {**test_inputs, **test_outputs},
3697 tensor_origin="test_tensor",
3698 pad_inputs="allow",
3699 crop_outputs="allow",
3700 )
3702 for rep_tol in self.config.bioimageio.reproducibility_tolerance:
3703 if not rep_tol.absolute_tolerance:
3704 continue
3706 if rep_tol.output_ids:
3707 out_arrays = {
3708 k: v[1] for k, v in test_outputs.items() if k in rep_tol.output_ids
3709 }
3710 else:
3711 out_arrays = {k: v[1] for k, v in test_outputs.items()}
3713 for out_id, array in out_arrays.items():
3714 if array is None:
3715 continue
3717 if rep_tol.absolute_tolerance > (max_test_value := array.max()) * 0.01:
3718 raise ValueError(
3719 "config.bioimageio.reproducibility_tolerance.absolute_tolerance="
3720 + f"{rep_tol.absolute_tolerance} > 0.01*{max_test_value}"
3721 + f" (1% of the maximum value of the test tensor '{out_id}')"
3722 )
3724 return self
3726 @model_validator(mode="after")
3727 def _validate_tensor_references_in_proc_kwargs(self, info: ValidationInfo) -> Self:
3728 ipt_refs = {t.id for t in self.inputs}
3729 missing_refs = [
3730 k["reference_tensor"]
3731 for k in [p.kwargs for ipt in self.inputs for p in ipt.preprocessing]
3732 + [p.kwargs for out in self.outputs for p in out.postprocessing]
3733 if "reference_tensor" in k
3734 and k["reference_tensor"] is not None
3735 and k["reference_tensor"] not in ipt_refs
3736 ]
3738 if missing_refs:
3739 raise ValueError(
3740 f"`reference_tensor`s {missing_refs} not found. Valid input tensor"
3741 + f" references are: {ipt_refs}."
3742 )
3744 return self
3746 name: Annotated[
3747 str,
3748 RestrictCharacters(string.ascii_letters + string.digits + "_+- ()"),
3749 MinLen(5),
3750 MaxLen(128),
3751 warn(MaxLen(64), "Name longer than 64 characters.", INFO),
3752 ]
3753 """A human-readable name of this model.
3754 It should be no longer than 64 characters
3755 and may only contain letter, number, underscore, minus, parentheses and spaces.
3756 We recommend to chose a name that refers to the model's task and image modality.
3757 """
3759 outputs: NotEmpty[Sequence[OutputTensorDescr]]
3760 """Describes the output tensors."""
3762 @field_validator("outputs", mode="after")
3763 @classmethod
3764 def _validate_tensor_ids(
3765 cls, outputs: Sequence[OutputTensorDescr], info: ValidationInfo
3766 ) -> Sequence[OutputTensorDescr]:
3767 tensor_ids = [
3768 t.id for t in info.data.get("inputs", []) + info.data.get("outputs", [])
3769 ]
3770 duplicate_tensor_ids: list[str] = []
3771 seen: set[str] = set()
3772 for t in tensor_ids:
3773 if t in seen:
3774 duplicate_tensor_ids.append(t)
3776 seen.add(t)
3778 if duplicate_tensor_ids:
3779 raise ValueError(f"Duplicate tensor ids: {duplicate_tensor_ids}")
3781 return outputs
3783 @staticmethod
3784 def _get_axes_with_parameterized_size(
3785 io: Sequence[InputTensorDescr] | Sequence[OutputTensorDescr],
3786 ):
3787 return {
3788 f"{t.id}.{a.id}": (t, a, a.size)
3789 for t in io
3790 for a in t.axes
3791 if not isinstance(a, BatchAxis) and isinstance(a.size, ParameterizedSize)
3792 }
3794 @staticmethod
3795 def _get_axes_with_independent_size(
3796 io: Sequence[InputTensorDescr] | Sequence[OutputTensorDescr],
3797 ):
3798 return {
3799 (t.id, a.id): (t, a, a.size)
3800 for t in io
3801 for a in t.axes
3802 if not isinstance(a, BatchAxis)
3803 and isinstance(a.size, (int, ParameterizedSize))
3804 }
3806 @field_validator("outputs", mode="after")
3807 @classmethod
3808 def _validate_output_axes(
3809 cls, outputs: list[OutputTensorDescr], info: ValidationInfo
3810 ) -> list[OutputTensorDescr]:
3811 input_size_refs = cls._get_axes_with_independent_size(
3812 info.data.get("inputs", [])
3813 )
3814 output_size_refs = cls._get_axes_with_independent_size(outputs)
3816 for i, out in enumerate(outputs):
3817 valid_independent_refs: dict[
3818 tuple[TensorId, AxisId],
3819 tuple[TensorDescr, AnyAxis, int | ParameterizedSize],
3820 ] = {
3821 **{
3822 (out.id, a.id): (out, a, a.size)
3823 for a in out.axes
3824 if not isinstance(a, BatchAxis)
3825 and isinstance(a.size, (int, ParameterizedSize))
3826 },
3827 **input_size_refs,
3828 **output_size_refs,
3829 }
3830 for a, ax in enumerate(out.axes):
3831 cls._validate_axis(
3832 "outputs",
3833 i,
3834 out.id,
3835 a,
3836 ax,
3837 valid_independent_refs=valid_independent_refs,
3838 )
3840 return outputs
3842 packaged_by: list[Author] = Field(
3843 default_factory=cast(Callable[[], List[Author]], list)
3844 )
3845 """The persons that have packaged and uploaded this model.
3846 Only required if those persons differ from the `authors`."""
3848 parent: LinkedModel | None = None
3849 """The model from which this model is derived, e.g. by fine-tuning the weights."""
3851 @model_validator(mode="after")
3852 def _validate_parent_is_not_self(self) -> Self:
3853 if self.parent is not None and self.parent.id == self.id:
3854 raise ValueError("A model description may not reference itself as parent.")
3856 return self
3858 run_mode: Annotated[
3859 RunMode | None,
3860 warn(None, "Run mode '{value}' has limited support across consumer softwares."),
3861 ] = None
3862 """Custom run mode for this model: for more complex prediction procedures like test time
3863 data augmentation that currently cannot be expressed in the specification.
3864 No standard run modes are defined yet."""
3866 timestamp: Datetime = Field(default_factory=Datetime.now)
3867 """Timestamp in [ISO 8601](#https://en.wikipedia.org/wiki/ISO_8601) format
3868 with a few restrictions listed [here](https://docs.python.org/3/library/datetime.html#datetime.datetime.fromisoformat).
3869 (In Python a datetime object is valid, too)."""
3871 training_data: Annotated[
3872 None | LinkedDataset | DatasetDescr | DatasetDescr02,
3873 Field(union_mode="left_to_right"),
3874 ] = None
3875 """The dataset used to train this model"""
3877 weights: Annotated[WeightsDescr, WrapSerializer(package_weights)]
3878 """The weights for this model.
3879 Weights can be given for different formats, but should otherwise be equivalent.
3880 The available weight formats determine which consumers can use this model."""
3882 config: Config = Field(default_factory=Config.model_construct)
3884 @model_validator(mode="after")
3885 def _add_default_cover(self) -> Self:
3886 if not get_validation_context().perform_io_checks or self.covers:
3887 return self
3889 try:
3890 generated_covers = generate_covers(
3891 [
3892 (t, load_array(t.test_tensor))
3893 for t in self.inputs
3894 if t.test_tensor is not None
3895 ],
3896 [
3897 (t, load_array(t.test_tensor))
3898 for t in self.outputs
3899 if t.test_tensor is not None
3900 ],
3901 )
3902 except Exception as e:
3903 issue_warning(
3904 "Failed to generate cover image(s): {e}",
3905 value=self.covers,
3906 msg_context={"e": e},
3907 field="covers",
3908 )
3909 else:
3910 self.covers.extend(generated_covers)
3912 return self
3914 def get_input_test_arrays(self) -> list[NDArray[Any]]:
3915 return self._get_test_arrays(self.inputs)
3917 def get_output_test_arrays(self) -> list[NDArray[Any]]:
3918 return self._get_test_arrays(self.outputs)
3920 @staticmethod
3921 def _get_test_arrays(
3922 io_descr: Sequence[InputTensorDescr] | Sequence[OutputTensorDescr],
3923 ):
3924 ts: list[FileDescr] = []
3925 for d in io_descr:
3926 if d.test_tensor is None:
3927 raise ValueError(
3928 f"Failed to get test arrays: description of '{d.id}' is missing a `test_tensor`."
3929 )
3930 ts.append(d.test_tensor)
3932 data = [load_array(t) for t in ts]
3933 assert all(isinstance(d, np.ndarray) for d in data)
3934 return data
3936 @staticmethod
3937 def get_batch_size(tensor_sizes: Mapping[TensorId, Mapping[AxisId, int]]) -> int:
3938 batch_size = 1
3939 tensor_with_batchsize: TensorId | None = None
3940 for tid in tensor_sizes:
3941 for aid, s in tensor_sizes[tid].items():
3942 if aid != BATCH_AXIS_ID or s == 1 or s == batch_size:
3943 continue
3945 if batch_size != 1:
3946 assert tensor_with_batchsize is not None
3947 raise ValueError(
3948 f"batch size mismatch for tensors '{tensor_with_batchsize}' ({batch_size}) and '{tid}' ({s})"
3949 )
3951 batch_size = s
3952 tensor_with_batchsize = tid
3954 return batch_size
3956 def get_output_tensor_sizes(
3957 self, input_sizes: Mapping[TensorId, Mapping[AxisId, int]]
3958 ) -> dict[TensorId, dict[AxisId, int | _DataDepSize]]:
3959 """Returns the tensor output sizes for given **input_sizes**.
3960 Only if **input_sizes** has a valid input shape, the tensor output size is exact.
3961 Otherwise it might be larger than the actual (valid) output"""
3962 batch_size = self.get_batch_size(input_sizes)
3963 ns = self.get_ns(input_sizes)
3965 tensor_sizes = self.get_tensor_sizes(ns, batch_size=batch_size)
3966 return tensor_sizes.outputs
3968 def get_ns(self, input_sizes: Mapping[TensorId, Mapping[AxisId, int]]):
3969 """get parameter `n` for each parameterized axis
3970 such that the valid input size is >= the given input size"""
3971 ret: dict[tuple[TensorId, AxisId], ParameterizedSize_N] = {}
3972 axes = {t.id: {a.id: a for a in t.axes} for t in self.inputs}
3973 for tid in input_sizes:
3974 for aid, s in input_sizes[tid].items():
3975 size_descr = axes[tid][aid].size
3976 if isinstance(size_descr, ParameterizedSize):
3977 ret[(tid, aid)] = size_descr.get_n(s)
3978 elif size_descr is None or isinstance(size_descr, (int, SizeReference)):
3979 pass
3980 else:
3981 assert_never(size_descr)
3983 return ret
3985 def get_tensor_sizes(
3986 self,
3987 ns: Mapping[tuple[TensorId, AxisId], ParameterizedSize_N],
3988 batch_size: int,
3989 max_input_shape: Mapping[TensorId, Mapping[AxisId, int]] | None = None,
3990 ) -> _TensorSizes:
3991 max_axis_sizes: dict[tuple[TensorId, AxisId], int] = {}
3992 for m, this_max_axis_sizes in (max_input_shape or {}).items():
3993 for a, s in this_max_axis_sizes.items():
3994 max_axis_sizes[(m, a)] = s
3996 axis_sizes = self.get_axis_sizes(
3997 ns, batch_size=batch_size, max_input_shape=max_axis_sizes
3998 )
3999 return _TensorSizes(
4000 {
4001 t: {
4002 aa: axis_sizes.inputs[(tt, aa)]
4003 for tt, aa in axis_sizes.inputs
4004 if tt == t
4005 }
4006 for t in {tt for tt, _ in axis_sizes.inputs}
4007 },
4008 {
4009 t: {
4010 aa: axis_sizes.outputs[(tt, aa)]
4011 for tt, aa in axis_sizes.outputs
4012 if tt == t
4013 }
4014 for t in {tt for tt, _ in axis_sizes.outputs}
4015 },
4016 )
4018 def get_axis_sizes(
4019 self,
4020 ns: Mapping[tuple[TensorId, AxisId], ParameterizedSize_N],
4021 batch_size: int | None = None,
4022 *,
4023 max_input_shape: Mapping[tuple[TensorId, AxisId], int] | None = None,
4024 ) -> _AxisSizes:
4025 """Determine input and output block shape for scale factors **ns**
4026 of parameterized input sizes.
4028 Args:
4029 ns: Scale factor `n` for each axis (keyed by (tensor_id, axis_id))
4030 that is parameterized as `size = min + n * step`.
4031 batch_size: The desired size of the batch dimension.
4032 If given **batch_size** overwrites any batch size present in
4033 **max_input_shape**. Default 1.
4034 max_input_shape: Limits the derived block shapes.
4035 Each axis for which the input size, parameterized by `n`, is larger
4036 than **max_input_shape** is set to the minimal value `n_min` for which
4037 this is still true.
4038 Use this for small input samples or large values of **ns**.
4039 Or simply whenever you know the full input shape.
4041 Returns:
4042 Resolved axis sizes for model inputs and outputs.
4043 """
4044 max_input_shape = max_input_shape or {}
4045 if batch_size is None:
4046 for (_t_id, a_id), s in max_input_shape.items():
4047 if a_id == BATCH_AXIS_ID:
4048 batch_size = s
4049 break
4050 else:
4051 batch_size = 1
4053 all_axes = {
4054 t.id: {a.id: a for a in t.axes} for t in chain(self.inputs, self.outputs)
4055 }
4057 inputs: dict[tuple[TensorId, AxisId], int] = {}
4058 outputs: dict[tuple[TensorId, AxisId], int | _DataDepSize] = {}
4060 def get_axis_size(a: InputAxis | OutputAxis):
4061 if isinstance(a, BatchAxis):
4062 if (t_descr.id, a.id) in ns:
4063 logger.warning(
4064 "Ignoring unexpected size increment factor (n) for batch axis"
4065 + " of tensor '{}'.",
4066 t_descr.id,
4067 )
4068 return batch_size
4069 elif isinstance(a.size, int):
4070 if (t_descr.id, a.id) in ns:
4071 logger.warning(
4072 "Ignoring unexpected size increment factor (n) for fixed size"
4073 + " axis '{}' of tensor '{}'.",
4074 a.id,
4075 t_descr.id,
4076 )
4077 return a.size
4078 elif isinstance(a.size, ParameterizedSize):
4079 if (t_descr.id, a.id) not in ns:
4080 raise ValueError(
4081 "Size increment factor (n) missing for parametrized axis"
4082 + f" '{a.id}' of tensor '{t_descr.id}'."
4083 )
4084 n = ns[(t_descr.id, a.id)]
4085 s_max = max_input_shape.get((t_descr.id, a.id))
4086 if s_max is not None:
4087 n = min(n, a.size.get_n(s_max))
4089 return a.size.get_size(n)
4091 elif isinstance(a.size, SizeReference):
4092 if (t_descr.id, a.id) in ns:
4093 logger.warning(
4094 "Ignoring unexpected size increment factor (n) for axis '{}'"
4095 + " of tensor '{}' with size reference.",
4096 a.id,
4097 t_descr.id,
4098 )
4099 assert not isinstance(a, BatchAxis)
4100 ref_axis = all_axes[a.size.tensor_id][a.size.axis_id]
4101 assert not isinstance(ref_axis, BatchAxis)
4102 ref_key = (a.size.tensor_id, a.size.axis_id)
4103 ref_size = inputs.get(ref_key, outputs.get(ref_key))
4104 assert ref_size is not None, ref_key
4105 assert not isinstance(ref_size, _DataDepSize), ref_key
4106 return a.size.get_size(
4107 axis=a,
4108 ref_axis=ref_axis,
4109 ref_size=ref_size,
4110 )
4111 elif isinstance(a.size, DataDependentSize):
4112 if (t_descr.id, a.id) in ns:
4113 logger.warning(
4114 "Ignoring unexpected increment factor (n) for data dependent"
4115 + " size axis '{}' of tensor '{}'.",
4116 a.id,
4117 t_descr.id,
4118 )
4119 return _DataDepSize(a.size.min, a.size.max)
4120 else:
4121 assert_never(a.size)
4123 # first resolve all , but the `SizeReference` input sizes
4124 for t_descr in self.inputs:
4125 for a in t_descr.axes:
4126 if not isinstance(a.size, SizeReference):
4127 s = get_axis_size(a)
4128 assert not isinstance(s, _DataDepSize)
4129 inputs[t_descr.id, a.id] = s
4131 # resolve all other input axis sizes
4132 for t_descr in self.inputs:
4133 for a in t_descr.axes:
4134 if isinstance(a.size, SizeReference):
4135 s = get_axis_size(a)
4136 assert not isinstance(s, _DataDepSize)
4137 inputs[t_descr.id, a.id] = s
4139 # resolve all output axis sizes
4140 for t_descr in self.outputs:
4141 for a in t_descr.axes:
4142 assert not isinstance(a.size, ParameterizedSize)
4143 s = get_axis_size(a)
4144 outputs[t_descr.id, a.id] = s
4146 return _AxisSizes(inputs=inputs, outputs=outputs)
4148 @model_validator(mode="before")
4149 @classmethod
4150 def _convert(cls, data: dict[str, Any]) -> dict[str, Any]:
4151 cls.convert_from_old_format_wo_validation(data)
4152 return data
4154 @classmethod
4155 def convert_from_old_format_wo_validation(cls, data: dict[str, Any]) -> None:
4156 """Convert metadata following an older format version to this classes' format
4157 without validating the result.
4158 """
4159 if (
4160 data.get("type") == "model"
4161 and isinstance(fv := data.get("format_version"), str)
4162 and fv.count(".") == 2
4163 ):
4164 fv_parts = fv.split(".")
4165 if any(not p.isdigit() for p in fv_parts):
4166 return
4168 fv_tuple = tuple(map(int, fv_parts))
4170 assert cls.implemented_format_version_tuple[0:2] == (0, 5)
4171 if fv_tuple[:2] in ((0, 3), (0, 4)):
4172 m04 = _ModelDescr_v0_4.load(data)
4173 if isinstance(m04, InvalidDescr):
4174 try:
4175 updated = _model_conv.convert_as_dict(
4176 m04 # pyright: ignore[reportArgumentType]
4177 )
4178 except Exception as e:
4179 logger.error(
4180 "Failed to convert from invalid model 0.4 description."
4181 + f"\nerror: {e}"
4182 + "\nProceeding with model 0.5 validation without conversion."
4183 )
4184 updated = None
4185 else:
4186 updated = _model_conv.convert_as_dict(m04)
4188 if updated is not None:
4189 data.clear()
4190 data.update(updated)
4192 elif fv_tuple[:2] == (0, 5):
4193 # bump patch version
4194 data["format_version"] = cls.implemented_format_version
4196 if fv_tuple[:2] in ((0, 3), (0, 4)) or (
4197 fv_tuple[:2] == (0, 5) and fv_tuple[2] < 11
4198 ):
4199 convert_plain_covers_and_docs_and_icon(data)
4202class _ModelConv(Converter[_ModelDescr_v0_4, ModelDescr]):
4203 def _convert(
4204 self, src: _ModelDescr_v0_4, tgt: type[ModelDescr | dict[str, Any]]
4205 ) -> ModelDescr | dict[str, Any]:
4206 name = "".join(
4207 c if c in string.ascii_letters + string.digits + "_+- ()" else " "
4208 for c in src.name
4209 )
4211 def conv_authors(auths: Sequence[_Author_v0_4] | None):
4212 conv = (
4213 _author_conv.convert if TYPE_CHECKING else _author_conv.convert_as_dict
4214 )
4215 return None if auths is None else [conv(a) for a in auths]
4217 if TYPE_CHECKING:
4218 arch_file_conv = _arch_file_conv.convert
4219 arch_lib_conv = _arch_lib_conv.convert
4220 else:
4221 arch_file_conv = _arch_file_conv.convert_as_dict
4222 arch_lib_conv = _arch_lib_conv.convert_as_dict
4224 input_size_refs = {
4225 ipt.name: {
4226 a: s
4227 for a, s in zip(
4228 ipt.axes,
4229 (
4230 ipt.shape.min
4231 if isinstance(ipt.shape, _ParameterizedInputShape_v0_4)
4232 else ipt.shape
4233 ),
4234 )
4235 }
4236 for ipt in src.inputs
4237 if ipt.shape
4238 }
4239 output_size_refs = {
4240 **{
4241 out.name: {a: s for a, s in zip(out.axes, out.shape)}
4242 for out in src.outputs
4243 if not isinstance(out.shape, _ImplicitOutputShape_v0_4)
4244 },
4245 **input_size_refs,
4246 }
4248 return tgt(
4249 attachments=(
4250 []
4251 if src.attachments is None
4252 else [FileDescr(source=f) for f in src.attachments.files]
4253 ),
4254 authors=[_author_conv.convert_as_dict(a) for a in src.authors], # pyright: ignore[reportArgumentType]
4255 cite=[{"text": c.text, "doi": c.doi, "url": c.url} for c in src.cite], # pyright: ignore[reportArgumentType]
4256 config=src.config, # pyright: ignore[reportArgumentType]
4257 covers=[{"source": c} for c in src.covers], # pyright: ignore[reportArgumentType]
4258 description=src.description,
4259 documentation={"source": src.documentation} if src.documentation else None, # pyright: ignore[reportArgumentType]
4260 format_version="0.5.14",
4261 git_repo=src.git_repo, # pyright: ignore[reportArgumentType]
4262 icon={"source": src.icon} if src.icon else None, # pyright: ignore[reportArgumentType]
4263 id=None if src.id is None else ModelId(src.id),
4264 id_emoji=src.id_emoji,
4265 license=src.license, # type: ignore
4266 links=src.links,
4267 maintainers=[_maintainer_conv.convert_as_dict(m) for m in src.maintainers], # pyright: ignore[reportArgumentType]
4268 name=name,
4269 tags=src.tags,
4270 type=src.type,
4271 uploader=src.uploader,
4272 version=src.version,
4273 inputs=[ # pyright: ignore[reportArgumentType]
4274 _input_tensor_conv.convert_as_dict(ipt, tt, st, input_size_refs)
4275 for ipt, tt, st in zip(
4276 src.inputs,
4277 src.test_inputs,
4278 src.sample_inputs or [None] * len(src.test_inputs),
4279 )
4280 ],
4281 outputs=[ # pyright: ignore[reportArgumentType]
4282 _output_tensor_conv.convert_as_dict(out, tt, st, output_size_refs)
4283 for out, tt, st in zip(
4284 src.outputs,
4285 src.test_outputs,
4286 src.sample_outputs or [None] * len(src.test_outputs),
4287 )
4288 ],
4289 parent=(
4290 None
4291 if src.parent is None
4292 else LinkedModel(
4293 id=ModelId(
4294 str(src.parent.id)
4295 + (
4296 ""
4297 if src.parent.version_number is None
4298 else f"/{src.parent.version_number}"
4299 )
4300 )
4301 )
4302 ),
4303 training_data=(
4304 None
4305 if src.training_data is None
4306 else (
4307 LinkedDataset(
4308 id=DatasetId(
4309 str(src.training_data.id)
4310 + (
4311 ""
4312 if src.training_data.version_number is None
4313 else f"/{src.training_data.version_number}"
4314 )
4315 )
4316 )
4317 if isinstance(src.training_data, LinkedDataset02)
4318 else src.training_data
4319 )
4320 ),
4321 packaged_by=[_author_conv.convert_as_dict(a) for a in src.packaged_by], # pyright: ignore[reportArgumentType]
4322 run_mode=src.run_mode,
4323 timestamp=src.timestamp,
4324 weights=(WeightsDescr if TYPE_CHECKING else dict)(
4325 keras_hdf5=(w := src.weights.keras_hdf5)
4326 and (KerasHdf5WeightsDescr if TYPE_CHECKING else dict)(
4327 authors=conv_authors(w.authors),
4328 source=w.source,
4329 tensorflow_version=w.tensorflow_version or Version("1.15"),
4330 parent=w.parent,
4331 ),
4332 onnx=(w := src.weights.onnx)
4333 and (OnnxWeightsDescr if TYPE_CHECKING else dict)(
4334 source=w.source,
4335 authors=conv_authors(w.authors),
4336 parent=w.parent,
4337 opset_version=w.opset_version or 15,
4338 ),
4339 pytorch_state_dict=(w := src.weights.pytorch_state_dict)
4340 and (PytorchStateDictWeightsDescr if TYPE_CHECKING else dict)(
4341 source=w.source,
4342 authors=conv_authors(w.authors),
4343 parent=w.parent,
4344 architecture=(
4345 arch_file_conv(
4346 w.architecture,
4347 w.architecture_sha256,
4348 w.kwargs,
4349 )
4350 if isinstance(w.architecture, _CallableFromFile_v0_4)
4351 else arch_lib_conv(w.architecture, w.kwargs)
4352 ),
4353 pytorch_version=w.pytorch_version or Version("1.10"),
4354 dependencies=(
4355 None
4356 if w.dependencies is None
4357 else (FileDescr if TYPE_CHECKING else dict)(
4358 source=cast(
4359 FileSource,
4360 str(deps := w.dependencies)[
4361 (
4362 len("conda:")
4363 if str(deps).startswith("conda:")
4364 else 0
4365 ) :
4366 ],
4367 )
4368 )
4369 ),
4370 ),
4371 tensorflow_js=(w := src.weights.tensorflow_js)
4372 and (TensorflowJsWeightsDescr if TYPE_CHECKING else dict)(
4373 source=w.source,
4374 authors=conv_authors(w.authors),
4375 parent=w.parent,
4376 tensorflow_version=w.tensorflow_version or Version("1.15"),
4377 ),
4378 tensorflow_saved_model_bundle=(
4379 w := src.weights.tensorflow_saved_model_bundle
4380 )
4381 and (TensorflowSavedModelBundleWeightsDescr if TYPE_CHECKING else dict)(
4382 authors=conv_authors(w.authors),
4383 parent=w.parent,
4384 source=w.source,
4385 tensorflow_version=w.tensorflow_version or Version("1.15"),
4386 dependencies=(
4387 None
4388 if w.dependencies is None
4389 else (FileDescr if TYPE_CHECKING else dict)(
4390 source=cast(
4391 FileSource,
4392 (
4393 str(w.dependencies)[len("conda:") :]
4394 if str(w.dependencies).startswith("conda:")
4395 else str(w.dependencies)
4396 ),
4397 )
4398 )
4399 ),
4400 ),
4401 torchscript=(w := src.weights.torchscript)
4402 and (TorchscriptWeightsDescr if TYPE_CHECKING else dict)(
4403 source=w.source,
4404 authors=conv_authors(w.authors),
4405 parent=w.parent,
4406 pytorch_version=w.pytorch_version or Version("1.10"),
4407 ),
4408 ),
4409 )
4412_model_conv = _ModelConv(_ModelDescr_v0_4, ModelDescr)
4415# create better cover images for 3d data and non-image outputs
4416def generate_covers(
4417 inputs: Sequence[tuple[InputTensorDescr, NDArray[Any]]],
4418 outputs: Sequence[tuple[OutputTensorDescr, NDArray[Any]]],
4419) -> list[FileDescr]:
4420 def squeeze(
4421 data: NDArray[Any], axes: Sequence[AnyAxis]
4422 ) -> tuple[NDArray[Any], list[AnyAxis]]:
4423 """apply numpy.ndarray.squeeze while keeping track of the axis descriptions remaining"""
4424 if data.ndim != len(axes):
4425 raise ValueError(
4426 f"tensor shape {data.shape} does not match described axes"
4427 + f" {[a.id for a in axes]}"
4428 )
4430 axes = [deepcopy(a) for a, s in zip(axes, data.shape) if s != 1]
4431 return data.squeeze(), axes
4433 def normalize(
4434 data: NDArray[Any], axis: tuple[int, ...] | None, eps: float = 1e-7
4435 ) -> NDArray[np.float32]:
4436 data = data.astype("float32")
4437 data -= data.min(axis=axis, keepdims=True)
4438 data /= data.max(axis=axis, keepdims=True) + eps
4439 return data
4441 def to_2d_image(data: NDArray[Any], axes: Sequence[AnyAxis]):
4442 original_shape = data.shape
4443 original_axes = list(axes)
4444 data, axes = squeeze(data, axes)
4446 # take slice fom any batch or index axis if needed
4447 # and convert the first channel axis and take a slice from any additional channel axes
4448 slices: tuple[slice, ...] = ()
4449 ndim = data.ndim
4450 ndim_need = 3 if any(isinstance(a, ChannelAxis) for a in axes) else 2
4451 has_c_axis = False
4452 for i, a in enumerate(axes):
4453 s = data.shape[i]
4454 assert s > 1
4455 if (
4456 isinstance(a, (BatchAxis, IndexInputAxis, IndexOutputAxis))
4457 and ndim > ndim_need
4458 ):
4459 data = data[slices + (slice(s // 2 - 1, s // 2),)]
4460 ndim -= 1
4461 elif isinstance(a, ChannelAxis):
4462 if has_c_axis:
4463 # second channel axis
4464 data = data[slices + (slice(0, 1),)]
4465 ndim -= 1
4466 else:
4467 has_c_axis = True
4468 if s == 2:
4469 # visualize two channels with cyan and magenta
4470 data = np.concatenate(
4471 [
4472 data[slices + (slice(1, 2),)],
4473 data[slices + (slice(0, 1),)],
4474 (
4475 data[slices + (slice(0, 1),)]
4476 + data[slices + (slice(1, 2),)]
4477 )
4478 / 2, # TODO: take maximum instead?
4479 ],
4480 axis=i,
4481 )
4482 elif data.shape[i] == 3:
4483 pass # visualize 3 channels as RGB
4484 else:
4485 # visualize first 3 channels as RGB
4486 data = data[slices + (slice(3),)]
4488 assert data.shape[i] == 3
4490 slices += (slice(None),)
4492 data, axes = squeeze(data, axes)
4493 assert len(axes) == ndim
4494 # take slice from z axis if needed
4495 slices = ()
4496 if ndim > ndim_need:
4497 for i, a in enumerate(axes):
4498 s = data.shape[i]
4499 if a.id == AxisId("z"):
4500 data = data[slices + (slice(s // 2 - 1, s // 2),)]
4501 data, axes = squeeze(data, axes)
4502 ndim -= 1
4503 break
4505 slices += (slice(None),)
4507 # take slice from any space or time axis
4508 slices = ()
4510 for i, a in enumerate(axes):
4511 if ndim <= ndim_need:
4512 break
4514 s = data.shape[i]
4515 assert s > 1
4516 if isinstance(
4517 a, (SpaceInputAxis, SpaceOutputAxis, TimeInputAxis, TimeOutputAxis)
4518 ):
4519 data = data[slices + (slice(s // 2 - 1, s // 2),)]
4520 ndim -= 1
4522 slices += (slice(None),)
4524 del slices
4525 data, axes = squeeze(data, axes)
4526 assert len(axes) == ndim
4528 if (has_c_axis and ndim != 3) or (not has_c_axis and ndim != 2):
4529 raise ValueError(
4530 f"Failed to construct cover image from shape {original_shape} with axes {[a.id for a in original_axes]}."
4531 )
4533 if not has_c_axis:
4534 assert ndim == 2
4535 data = np.repeat(data[:, :, None], 3, axis=2)
4536 axes.append(ChannelAxis(channel_names=list("RGB")))
4537 ndim += 1
4539 assert ndim == 3
4541 # transpose axis order such that longest axis comes first...
4542 axis_order: list[int] = [int(i) for i in np.argsort(list(data.shape))]
4543 axis_order.reverse()
4544 # ... and channel axis is last
4545 c = next(i for i in range(3) if isinstance(axes[i], ChannelAxis))
4546 axis_order.append(axis_order.pop(c))
4547 axes = [axes[ao] for ao in axis_order]
4548 data = data.transpose(axis_order)
4550 # h, w = data.shape[:2]
4551 # if h / w in (1.0 or 2.0):
4552 # pass
4553 # elif h / w < 2:
4554 # TODO: enforce 2:1 or 1:1 aspect ratio for generated cover images
4556 norm_along = (
4557 tuple(i for i, a in enumerate(axes) if a.type in ("space", "time")) or None
4558 )
4559 # normalize the data and map to 8 bit
4560 data = normalize(data, norm_along)
4561 data = (data * 255).astype("uint8")
4563 return data
4565 def create_diagonal_split_image(im0: NDArray[Any], im1: NDArray[Any]):
4566 assert im0.dtype == im1.dtype == np.uint8
4567 assert im0.shape == im1.shape
4568 assert im0.ndim == 3
4569 N, M, C = im0.shape
4570 assert C == 3
4571 out = np.ones((N, M, C), dtype="uint8")
4572 for c in range(C):
4573 outc = np.tril(im0[..., c])
4574 mask = outc == 0
4575 outc[mask] = np.triu(im1[..., c])[mask]
4576 out[..., c] = outc
4578 return out
4580 if not inputs:
4581 raise ValueError("Missing test input tensor for cover generation.")
4583 if not outputs:
4584 raise ValueError("Missing test output tensor for cover generation.")
4586 ipt_descr, ipt = inputs[0]
4587 out_descr, out = outputs[0]
4589 ipt_img = to_2d_image(ipt, ipt_descr.axes)
4590 out_img = to_2d_image(out, out_descr.axes)
4592 cover_folder = Path(mkdtemp())
4593 if ipt_img.shape == out_img.shape:
4594 covers = [cover_folder / "cover.png"]
4595 imwrite(covers[0], create_diagonal_split_image(ipt_img, out_img))
4596 else:
4597 covers = [cover_folder / "input.png", cover_folder / "output.png"]
4598 imwrite(covers[0], ipt_img)
4599 imwrite(covers[1], out_img)
4601 return [FileDescr(source=c) for c in covers]