Coverage for src/bioimageio/spec/model/v0_5.py: 71%
1792 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 10:20 +0000
1from __future__ import annotations
3import collections.abc
4import re
5import string
6import warnings
7from abc import ABC
8from copy import deepcopy
9from functools import partial
10from itertools import chain
11from math import ceil
12from pathlib import Path, PurePosixPath
13from tempfile import mkdtemp
14from textwrap import dedent
15from typing import (
16 TYPE_CHECKING,
17 Any,
18 Callable,
19 ClassVar,
20 Dict,
21 Generic,
22 List,
23 Literal,
24 Mapping,
25 NamedTuple,
26 Optional,
27 Sequence,
28 TypeVar,
29 Union,
30 cast,
31 overload,
32)
34import numpy as np
35from annotated_types import Ge, Gt, Interval, MaxLen, MinLen, Predicate
36from imageio.v3 import imread, imwrite # pyright: ignore[reportUnknownVariableType]
37from loguru import logger
38from numpy.typing import NDArray
39from pydantic import (
40 AfterValidator,
41 Discriminator,
42 Field,
43 RootModel,
44 SerializationInfo,
45 SerializerFunctionWrapHandler,
46 StrictInt,
47 Tag,
48 ValidationInfo,
49 WrapSerializer,
50 field_validator,
51 model_serializer,
52 model_validator,
53)
54from typing_extensions import Annotated, Self, TypeAlias, assert_never, get_args
56from .._internal.common_nodes import (
57 InvalidDescr,
58 KwargsNode,
59 Node,
60 NodeWithExplicitlySetFields,
61)
62from .._internal.constants import DTYPE_LIMITS
63from .._internal.field_warning import issue_warning, warn
64from .._internal.io import BioimageioYamlContent as BioimageioYamlContent
65from .._internal.io import FileDescr as FileDescr
66from .._internal.io import (
67 FileSource,
68 WithSuffix,
69 YamlValue,
70 extract_file_name,
71 get_reader,
72 wo_special_file_name,
73)
74from .._internal.io_basics import Sha256 as Sha256
75from .._internal.io_packaging import (
76 FileDescr_package,
77 package_file_descr_serializer,
78)
79from .._internal.io_utils import load_array, open_bioimageio_yaml
80from .._internal.node_converter import Converter
81from .._internal.type_guards import is_dict, is_sequence
82from .._internal.types import (
83 FAIR,
84 LowerCaseIdentifier,
85 LowerCaseIdentifierAnno,
86 validate_identifier,
87 validate_is_not_keyword,
88)
89from .._internal.types import (
90 AbsoluteTolerance as AbsoluteTolerance,
91)
92from .._internal.types import Datetime as Datetime
93from .._internal.types import Identifier as Identifier
94from .._internal.types import (
95 MismatchedElementsPerMillion as MismatchedElementsPerMillion,
96)
97from .._internal.types import NotEmpty as NotEmpty
98from .._internal.types import (
99 RelativeTolerance as RelativeTolerance,
100)
101from .._internal.types import SiUnit as SiUnit
102from .._internal.url import HttpUrl as HttpUrl
103from .._internal.utils import try_all_raise_last
104from .._internal.validation_context import get_validation_context
105from .._internal.validator_annotations import RestrictCharacters
106from .._internal.version_type import Version as Version
107from .._internal.warning_levels import INFO
108from ..dataset.v0_2 import DatasetDescr as DatasetDescr02
109from ..dataset.v0_2 import LinkedDataset as LinkedDataset02
110from ..dataset.v0_3 import DatasetDescr as DatasetDescr
111from ..dataset.v0_3 import DatasetId as DatasetId
112from ..dataset.v0_3 import LinkedDataset as LinkedDataset
113from ..dataset.v0_3 import Uploader as Uploader
114from ..generic._v0_3_converter import convert_plain_covers_and_docs_and_icon
115from ..generic.v0_3 import (
116 VALID_COVER_IMAGE_EXTENSIONS as VALID_COVER_IMAGE_EXTENSIONS,
117)
118from ..generic.v0_3 import Author as Author
119from ..generic.v0_3 import BadgeDescr as BadgeDescr
120from ..generic.v0_3 import CiteEntry as CiteEntry
121from ..generic.v0_3 import DeprecatedLicenseId as DeprecatedLicenseId
122from ..generic.v0_3 import Doi as Doi
123from ..generic.v0_3 import (
124 FileDescr_documentation,
125 GenericModelDescrBase,
126 LinkedResourceBase,
127 _author_conv, # pyright: ignore[reportPrivateUsage]
128 _maintainer_conv, # pyright: ignore[reportPrivateUsage]
129)
130from ..generic.v0_3 import LicenseId as LicenseId
131from ..generic.v0_3 import LinkedResource as LinkedResource
132from ..generic.v0_3 import Maintainer as Maintainer
133from ..generic.v0_3 import OrcidId as OrcidId
134from ..generic.v0_3 import RelativeFilePath as RelativeFilePath
135from ..generic.v0_3 import ResourceId as ResourceId
136from .v0_4 import Author as _Author_v0_4
137from .v0_4 import BinarizeDescr as _BinarizeDescr_v0_4
138from .v0_4 import CallableFromDepencency as CallableFromDepencency
139from .v0_4 import CallableFromDepencency as _CallableFromDepencency_v0_4
140from .v0_4 import CallableFromFile as _CallableFromFile_v0_4
141from .v0_4 import ClipDescr as _ClipDescr_v0_4
142from .v0_4 import ImplicitOutputShape as _ImplicitOutputShape_v0_4
143from .v0_4 import InputTensorDescr as _InputTensorDescr_v0_4
144from .v0_4 import KnownRunMode as KnownRunMode
145from .v0_4 import ModelDescr as _ModelDescr04
146from .v0_4 import ModelDescr as _ModelDescr_v0_4
147from .v0_4 import OutputTensorDescr as _OutputTensorDescr_v0_4
148from .v0_4 import ParameterizedInputShape as _ParameterizedInputShape_v0_4
149from .v0_4 import PostprocessingDescr as _PostprocessingDescr_v0_4
150from .v0_4 import PreprocessingDescr as _PreprocessingDescr_v0_4
151from .v0_4 import RunMode as RunMode
152from .v0_4 import ScaleLinearDescr as _ScaleLinearDescr_v0_4
153from .v0_4 import ScaleMeanVarianceDescr as _ScaleMeanVarianceDescr_v0_4
154from .v0_4 import ScaleRangeDescr as _ScaleRangeDescr_v0_4
155from .v0_4 import SigmoidDescr as _SigmoidDescr_v0_4
156from .v0_4 import TensorName as _TensorName_v0_4
157from .v0_4 import ZeroMeanUnitVarianceDescr as _ZeroMeanUnitVarianceDescr_v0_4
158from .v0_4 import package_weights
160SpaceUnit = Literal[
161 "attometer",
162 "angstrom",
163 "centimeter",
164 "decimeter",
165 "exameter",
166 "femtometer",
167 "foot",
168 "gigameter",
169 "hectometer",
170 "inch",
171 "kilometer",
172 "megameter",
173 "meter",
174 "micrometer",
175 "mile",
176 "millimeter",
177 "nanometer",
178 "parsec",
179 "petameter",
180 "picometer",
181 "terameter",
182 "yard",
183 "yoctometer",
184 "yottameter",
185 "zeptometer",
186 "zettameter",
187]
188"""Space unit compatible to the [OME-Zarr axes specification 0.5](https://ngff.openmicroscopy.org/0.5/#axes-md)"""
190TimeUnit = Literal[
191 "attosecond",
192 "centisecond",
193 "day",
194 "decisecond",
195 "exasecond",
196 "femtosecond",
197 "gigasecond",
198 "hectosecond",
199 "hour",
200 "kilosecond",
201 "megasecond",
202 "microsecond",
203 "millisecond",
204 "minute",
205 "nanosecond",
206 "petasecond",
207 "picosecond",
208 "second",
209 "terasecond",
210 "yoctosecond",
211 "yottasecond",
212 "zeptosecond",
213 "zettasecond",
214]
215"""Time unit compatible to the [OME-Zarr axes specification 0.5](https://ngff.openmicroscopy.org/0.5/#axes-md)"""
217AxisType = Literal["batch", "channel", "index", "time", "space"]
219_AXIS_TYPE_MAP: Mapping[str, AxisType] = {
220 "b": "batch",
221 "t": "time",
222 "i": "index",
223 "c": "channel",
224 "x": "space",
225 "y": "space",
226 "z": "space",
227}
229_AXIS_ID_MAP = {
230 "b": "batch",
231 "t": "time",
232 "i": "index",
233 "c": "channel",
234 "s": "channel",
235}
237WeightsFormat = Literal[
238 "keras_hdf5",
239 "keras_v3",
240 "onnx",
241 "pytorch_state_dict",
242 "tensorflow_js",
243 "tensorflow_saved_model_bundle",
244 "torchscript",
245]
248class TensorId(LowerCaseIdentifier):
249 root_model: ClassVar[type[RootModel[Any]]] = RootModel[
250 Annotated[LowerCaseIdentifierAnno, MaxLen(32)]
251 ]
254def _normalize_axis_id(a: str):
255 b = str(a).lower()
256 normalized = _AXIS_ID_MAP.get(b, b)
257 if a != normalized:
258 logger.opt(depth=3).debug(
259 "Normalized axis id from '{}' to '{}'.", a, normalized
260 )
261 return normalized
264class AxisId(LowerCaseIdentifier):
265 root_model: ClassVar[type[RootModel[Any]]] = RootModel[
266 Annotated[
267 NotEmpty[str],
268 AfterValidator(_normalize_axis_id),
269 MaxLen(16),
270 AfterValidator(validate_identifier),
271 AfterValidator(validate_is_not_keyword),
272 ]
273 ]
276def _is_batch(a: str) -> bool:
277 return str(a) == "batch"
280def _is_not_batch(a: str) -> bool:
281 return not _is_batch(a)
284NonBatchAxisId = Annotated[AxisId, Predicate(_is_not_batch)]
286PreprocessingId = Literal[
287 "binarize",
288 "clip",
289 "ensure_dtype",
290 "fixed_zero_mean_unit_variance",
291 "scale_linear",
292 "scale_range",
293 "sigmoid",
294 "softmax",
295]
296PostprocessingId = Literal[
297 "binarize",
298 "clip",
299 "custom",
300 "ensure_dtype",
301 "fixed_zero_mean_unit_variance",
302 "scale_linear",
303 "scale_mean_variance",
304 "scale_range",
305 "sigmoid",
306 "softmax",
307 "zero_mean_unit_variance",
308]
311SAME_AS_TYPE = "<same as type>"
314ParameterizedSize_N: TypeAlias = int
315"""
316Annotates an integer to calculate a concrete axis size from a `ParameterizedSize`.
317"""
320class ParameterizedSize(Node):
321 """Describes a range of valid tensor axis sizes as `size = min + n*step`.
323 - **min** and **step** are given by the model description.
324 - All blocksize paramters n = 0,1,2,... yield a valid `size`.
325 - A greater blocksize paramter n = 0,1,2,... results in a greater **size**.
326 This allows to adjust the axis size more generically.
327 """
329 N: ClassVar[type[int]] = ParameterizedSize_N
330 """Positive integer to parameterize this axis"""
332 min: Annotated[int, Gt(0)]
333 step: Annotated[int, Gt(0)]
335 def validate_size(self, size: int, msg_prefix: str = "") -> int:
336 if size < self.min:
337 raise ValueError(
338 f"{msg_prefix}size {size} < {self.min} (minimum axis size)"
339 )
340 if (size - self.min) % self.step != 0:
341 raise ValueError(
342 f"{msg_prefix}size {size} is not parameterized by `min + n*step` ="
343 + f" `{self.min} + n*{self.step}`"
344 )
346 return size
348 def get_size(self, n: ParameterizedSize_N) -> int:
349 return self.min + self.step * n
351 def get_n(self, s: int) -> ParameterizedSize_N:
352 """return smallest n parameterizing a size greater or equal than `s`"""
353 return ceil((s - self.min) / self.step)
356class DataDependentSize(Node):
357 min: Annotated[int, Gt(0)] = 1
358 max: Annotated[int | None, Gt(1)] = None
360 @model_validator(mode="after")
361 def _validate_max_gt_min(self):
362 if self.max is not None and self.min >= self.max:
363 raise ValueError(f"expected `min` < `max`, but got {self.min}, {self.max}")
365 return self
367 def validate_size(self, size: int, msg_prefix: str = "") -> int:
368 if size < self.min:
369 raise ValueError(f"{msg_prefix}size {size} < {self.min}")
371 if self.max is not None and size > self.max:
372 raise ValueError(f"{msg_prefix}size {size} > {self.max}")
374 return size
377class SizeReference(Node):
378 """A tensor axis size (extent in pixels/frames) defined in relation to a reference axis.
380 `axis.size = reference.size * reference.scale / axis.scale + offset`
382 Note:
383 1. The axis and the referenced axis need to have the same unit (or no unit).
384 2. Batch axes may not be referenced.
385 3. Fractions are rounded down.
386 4. If the reference axis is `concatenable` the referencing axis is assumed to be
387 `concatenable` as well with the same block order.
389 Example:
390 An unisotropic input image of w*h=100*49 pixels depicts a phsical space of 200*196mm².
391 Let's assume that we want to express the image height h in relation to its width w
392 instead of only accepting input images of exactly 100*49 pixels
393 (for example to express a range of valid image shapes by parametrizing w, see `ParameterizedSize`).
395 >>> w = SpaceInputAxis(id=AxisId("w"), size=100, unit="millimeter", scale=2)
396 >>> h = SpaceInputAxis(
397 ... id=AxisId("h"),
398 ... size=SizeReference(tensor_id=TensorId("input"), axis_id=AxisId("w"), offset=-1),
399 ... unit="millimeter",
400 ... scale=4,
401 ... )
402 >>> print(h.size.get_size(h, w))
403 49
405 ⇒ h = w * w.scale / h.scale + offset = 100 * 2mm / 4mm - 1 = 49
406 """
408 tensor_id: TensorId
409 """tensor id of the reference axis"""
411 axis_id: AxisId
412 """axis id of the reference axis"""
414 offset: StrictInt = 0
416 def get_size(
417 self,
418 axis: ChannelAxis
419 | IndexInputAxis
420 | IndexOutputAxis
421 | TimeInputAxis
422 | SpaceInputAxis
423 | TimeOutputAxis
424 | TimeOutputAxisWithHalo
425 | SpaceOutputAxis
426 | SpaceOutputAxisWithHalo,
427 ref_axis: ChannelAxis
428 | IndexInputAxis
429 | IndexOutputAxis
430 | TimeInputAxis
431 | SpaceInputAxis
432 | TimeOutputAxis
433 | TimeOutputAxisWithHalo
434 | SpaceOutputAxis
435 | SpaceOutputAxisWithHalo,
436 n: ParameterizedSize_N = 0,
437 ref_size: int | None = None,
438 ):
439 """Compute the concrete size for a given axis and its reference axis.
441 Args:
442 axis: The axis this [SizeReference][] is the size of.
443 ref_axis: The reference axis to compute the size from.
444 n: If the **ref_axis** is parameterized (of type `ParameterizedSize`)
445 and no fixed **ref_size** is given,
446 **n** is used to compute the size of the parameterized **ref_axis**.
447 ref_size: Overwrite the reference size instead of deriving it from
448 **ref_axis**
449 (**ref_axis.scale** is still used; any given **n** is ignored).
450 """
451 assert axis.size == self, (
452 "Given `axis.size` is not defined by this `SizeReference`"
453 )
455 assert ref_axis.id == self.axis_id, (
456 f"Expected `ref_axis.id` to be {self.axis_id}, but got {ref_axis.id}."
457 )
459 assert axis.unit == ref_axis.unit, (
460 "`SizeReference` requires `axis` and `ref_axis` to have the same `unit`,"
461 f" but {axis.unit}!={ref_axis.unit}"
462 )
463 if ref_size is None:
464 if isinstance(ref_axis.size, (int, float)):
465 ref_size = ref_axis.size
466 elif isinstance(ref_axis.size, ParameterizedSize):
467 ref_size = ref_axis.size.get_size(n)
468 elif isinstance(ref_axis.size, DataDependentSize):
469 raise ValueError(
470 "Reference axis referenced in `SizeReference` may not be a `DataDependentSize`."
471 )
472 elif isinstance(ref_axis.size, SizeReference):
473 raise ValueError(
474 "Reference axis referenced in `SizeReference` may not be sized by a"
475 + " `SizeReference` itself."
476 )
477 else:
478 assert_never(ref_axis.size)
480 return int(ref_size * ref_axis.scale / axis.scale + self.offset)
482 @staticmethod
483 def _get_unit(
484 axis: ChannelAxis
485 | IndexInputAxis
486 | IndexOutputAxis
487 | TimeInputAxis
488 | SpaceInputAxis
489 | TimeOutputAxis
490 | TimeOutputAxisWithHalo
491 | SpaceOutputAxis
492 | SpaceOutputAxisWithHalo,
493 ):
494 return axis.unit
497class AxisBase(NodeWithExplicitlySetFields):
498 id: AxisId
499 """An axis id unique across all axes of one tensor."""
501 description: Annotated[str, MaxLen(128)] = ""
502 """A short description of this axis beyond its type and id."""
505class WithHalo(Node):
506 halo: Annotated[int, Ge(1)]
507 """The halo should be cropped from the output tensor to avoid boundary effects.
508 It is to be cropped from both sides, i.e. `size_after_crop = size - 2 * halo`.
509 To document a halo that is already cropped by the model use `size.offset` instead."""
511 size: Annotated[
512 SizeReference,
513 Field(examples=[{"tensor_id": "t", "axis_id": "a", "offset": 5}]),
514 ]
515 """reference to another axis with an optional offset (see [SizeReference][])"""
518BATCH_AXIS_ID = AxisId("batch")
519CHANNEL_AXIS_ID = AxisId("channel")
520DEFAULT_SPACE_AXIS_ID = AxisId("x")
521DEFAULT_INDEX_AXIS_ID = AxisId("index")
522DEFAULT_TIME_AXIS_ID = AxisId("time")
525class BatchAxis(AxisBase):
526 implemented_type: ClassVar[Literal["batch"]] = "batch"
527 if TYPE_CHECKING:
528 type: Literal["batch"] = "batch"
529 else:
530 type: Literal["batch"]
532 id: Annotated[AxisId, Predicate(_is_batch)] = BATCH_AXIS_ID
533 size: Literal[1] | None = None
534 """The batch size may be fixed to 1,
535 otherwise (the default) it may be chosen arbitrarily depending on available memory"""
537 @property
538 def scale(self):
539 return 1.0
541 @property
542 def concatenable(self):
543 return True
545 @property
546 def unit(self):
547 return None
550class ChannelAxis(AxisBase):
551 implemented_type: ClassVar[Literal["channel"]] = "channel"
552 if TYPE_CHECKING:
553 type: Literal["channel"] = "channel"
554 else:
555 type: Literal["channel"]
557 id: NonBatchAxisId = CHANNEL_AXIS_ID
559 channel_names: NotEmpty[list[str]]
560 """Name/label for each channel. The number of channels is given by `len(channel_names)`."""
562 @property
563 def size(self) -> int:
564 return len(self.channel_names)
566 @property
567 def concatenable(self):
568 return False
570 @property
571 def scale(self) -> float:
572 return 1.0
574 @property
575 def unit(self):
576 return None
579class _WithInputAxisSize(Node):
580 size: Annotated[
581 Annotated[int, Gt(0)] | ParameterizedSize | SizeReference,
582 Field(
583 examples=[
584 10,
585 ParameterizedSize(min=32, step=16).model_dump(mode="json"),
586 {"tensor_id": "t", "axis_id": "a", "offset": 5},
587 ]
588 ),
589 ]
590 """The size/length of this axis can be specified as
591 - fixed integer
592 - parameterized series of valid sizes ([ParameterizedSize][])
593 - reference to another axis with an optional offset ([SizeReference][])
594 """
597class IndexAxisBase(AxisBase):
598 implemented_type: ClassVar[Literal["index"]] = "index"
599 if TYPE_CHECKING:
600 type: Literal["index"] = "index"
601 else:
602 type: Literal["index"]
604 id: NonBatchAxisId = DEFAULT_INDEX_AXIS_ID
606 @property
607 def scale(self) -> float:
608 return 1.0
610 @property
611 def unit(self):
612 return None
615class IndexInputAxis(IndexAxisBase, _WithInputAxisSize):
616 concatenable: bool = False
617 """If a model has a `concatenable` input axis, it can be processed blockwise,
618 splitting a longer sample axis into blocks matching its input tensor description.
619 Output axes are concatenable if they have a [SizeReference][] to a concatenable
620 input axis.
621 """
624class IndexOutputAxis(IndexAxisBase):
625 size: Annotated[
626 Annotated[int, Gt(0)] | SizeReference | DataDependentSize,
627 Field(examples=[10, {"tensor_id": "t", "axis_id": "a", "offset": 5}]),
628 ]
629 """The size/length of this axis can be specified as
630 - fixed integer
631 - reference to another axis with an optional offset ([SizeReference][])
632 - data dependent size using [DataDependentSize][] (size is only known after model inference)
633 """
636class TimeAxisBase(AxisBase):
637 implemented_type: ClassVar[Literal["time"]] = "time"
638 if TYPE_CHECKING:
639 type: Literal["time"] = "time"
640 else:
641 type: Literal["time"]
643 id: NonBatchAxisId = DEFAULT_TIME_AXIS_ID
644 unit: TimeUnit | None = None
645 scale: Annotated[float, Gt(0)] = 1.0
648class TimeInputAxis(TimeAxisBase, _WithInputAxisSize):
649 concatenable: bool = False
650 """If a model has a `concatenable` input axis, it can be processed blockwise,
651 splitting a longer sample axis into blocks matching its input tensor description.
652 Output axes are concatenable if they have a [SizeReference][] to a concatenable
653 input axis.
654 """
657class SpaceAxisBase(AxisBase):
658 implemented_type: ClassVar[Literal["space"]] = "space"
659 if TYPE_CHECKING:
660 type: Literal["space"] = "space"
661 else:
662 type: Literal["space"]
664 id: Annotated[NonBatchAxisId, Field(examples=["x", "y", "z"])] = (
665 DEFAULT_SPACE_AXIS_ID
666 )
667 unit: SpaceUnit | None = None
668 scale: Annotated[float, Gt(0)] = 1.0
671class SpaceInputAxis(SpaceAxisBase, _WithInputAxisSize):
672 concatenable: bool = False
673 """If a model has a `concatenable` input axis, it can be processed blockwise,
674 splitting a longer sample axis into blocks matching its input tensor description.
675 Output axes are concatenable if they have a [SizeReference][] to a concatenable
676 input axis.
677 """
680INPUT_AXIS_TYPES = (
681 BatchAxis,
682 ChannelAxis,
683 IndexInputAxis,
684 TimeInputAxis,
685 SpaceInputAxis,
686)
687"""intended for isinstance comparisons in py<3.10"""
689_InputAxisUnion = Union[
690 BatchAxis, ChannelAxis, IndexInputAxis, TimeInputAxis, SpaceInputAxis
691]
692InputAxis = Annotated[_InputAxisUnion, Discriminator("type")]
695class _WithOutputAxisSize(Node):
696 size: Annotated[
697 Annotated[int, Gt(0)] | SizeReference,
698 Field(examples=[10, {"tensor_id": "t", "axis_id": "a", "offset": 5}]),
699 ]
700 """The size/length of this axis can be specified as
701 - fixed integer
702 - reference to another axis with an optional offset (see [SizeReference][])
703 """
706class TimeOutputAxis(TimeAxisBase, _WithOutputAxisSize):
707 pass
710class TimeOutputAxisWithHalo(TimeAxisBase, WithHalo):
711 pass
714def _get_halo_axis_discriminator_value(v: Any) -> Literal["with_halo", "wo_halo"]:
715 if isinstance(v, dict):
716 return "with_halo" if "halo" in v else "wo_halo"
717 else:
718 return "with_halo" if hasattr(v, "halo") else "wo_halo"
721_TimeOutputAxisUnion = Annotated[
722 Union[
723 Annotated[TimeOutputAxis, Tag("wo_halo")],
724 Annotated[TimeOutputAxisWithHalo, Tag("with_halo")],
725 ],
726 Discriminator(_get_halo_axis_discriminator_value),
727]
730class SpaceOutputAxis(SpaceAxisBase, _WithOutputAxisSize):
731 pass
734class SpaceOutputAxisWithHalo(SpaceAxisBase, WithHalo):
735 pass
738_SpaceOutputAxisUnion = Annotated[
739 Union[
740 Annotated[SpaceOutputAxis, Tag("wo_halo")],
741 Annotated[SpaceOutputAxisWithHalo, Tag("with_halo")],
742 ],
743 Discriminator(_get_halo_axis_discriminator_value),
744]
747_OutputAxisUnion = Union[
748 BatchAxis, ChannelAxis, IndexOutputAxis, _TimeOutputAxisUnion, _SpaceOutputAxisUnion
749]
750OutputAxis = Annotated[_OutputAxisUnion, Discriminator("type")]
752OUTPUT_AXIS_TYPES = (
753 BatchAxis,
754 ChannelAxis,
755 IndexOutputAxis,
756 TimeOutputAxis,
757 TimeOutputAxisWithHalo,
758 SpaceOutputAxis,
759 SpaceOutputAxisWithHalo,
760)
761"""intended for isinstance comparisons in py<3.10"""
764AnyAxis = Union[InputAxis, OutputAxis]
766ANY_AXIS_TYPES = INPUT_AXIS_TYPES + OUTPUT_AXIS_TYPES
767"""intended for isinstance comparisons in py<3.10"""
769TVs = Union[
770 NotEmpty[List[int]],
771 NotEmpty[List[float]],
772 NotEmpty[List[bool]],
773 NotEmpty[List[str]],
774]
777NominalOrOrdinalDType = Literal[
778 "float32",
779 "float64",
780 "uint8",
781 "int8",
782 "uint16",
783 "int16",
784 "uint32",
785 "int32",
786 "uint64",
787 "int64",
788 "bool",
789]
792class NominalOrOrdinalDataDescr(Node):
793 values: TVs
794 """A fixed set of nominal or an ascending sequence of ordinal values.
795 In this case `data.type` is required to be an unsigend integer type, e.g. 'uint8'.
796 String `values` are interpreted as labels for tensor values 0, ..., N.
797 Note: as YAML 1.2 does not natively support a "set" datatype,
798 nominal values should be given as a sequence (aka list/array) as well.
799 """
801 type: Annotated[
802 NominalOrOrdinalDType,
803 Field(
804 examples=[
805 "float32",
806 "uint8",
807 "uint16",
808 "int64",
809 "bool",
810 ],
811 ),
812 ] = "uint8"
814 @model_validator(mode="after")
815 def _validate_values_match_type(
816 self,
817 ) -> Self:
818 incompatible: list[Any] = []
819 for v in self.values:
820 if self.type == "bool":
821 if not isinstance(v, bool):
822 incompatible.append(v)
823 elif self.type in DTYPE_LIMITS:
824 if (
825 isinstance(v, (int, float))
826 and (
827 v < DTYPE_LIMITS[self.type].min
828 or v > DTYPE_LIMITS[self.type].max
829 )
830 or (isinstance(v, str) and "uint" not in self.type)
831 or (isinstance(v, float) and "int" in self.type)
832 ):
833 incompatible.append(v)
834 else:
835 incompatible.append(v)
837 if len(incompatible) == 5:
838 incompatible.append("...")
839 break
841 if incompatible:
842 raise ValueError(
843 f"data type '{self.type}' incompatible with values {incompatible}"
844 )
846 return self
848 unit: Literal["arbitrary unit"] | SiUnit | None = None
850 @property
851 def range(self):
852 if isinstance(self.values[0], str):
853 return 0, len(self.values) - 1
854 else:
855 return min(self.values), max(self.values)
858IntervalOrRatioDType = Literal[
859 "float32",
860 "float64",
861 "uint8",
862 "int8",
863 "uint16",
864 "int16",
865 "uint32",
866 "int32",
867 "uint64",
868 "int64",
869]
872class IntervalOrRatioDataDescr(Node):
873 type: Annotated[ # TODO: rename to dtype
874 IntervalOrRatioDType,
875 Field(
876 examples=["float32", "float64", "uint8", "uint16"],
877 ),
878 ] = "float32"
879 range: tuple[float | None, float | None] = (
880 None,
881 None,
882 )
883 """Tuple `(minimum, maximum)` specifying the allowed range of the data in this tensor.
884 `None` corresponds to min/max of what can be expressed by **type**."""
885 unit: Literal["arbitrary unit"] | SiUnit = "arbitrary unit"
886 scale: float = 1.0
887 """Scale for data on an interval (or ratio) scale."""
888 offset: float | None = None
889 """Offset for data on a ratio scale."""
891 @model_validator(mode="before")
892 def _replace_inf(cls, data: Any):
893 if is_dict(data) and "range" in data and is_sequence(data["range"]):
894 forbidden = (
895 "inf",
896 "-inf",
897 ".inf",
898 "-.inf",
899 float("inf"),
900 float("-inf"),
901 )
902 if any(v in forbidden for v in data["range"]):
903 issue_warning("replaced 'inf' value", value=data["range"])
905 data["range"] = tuple(
906 (None if v in forbidden else v) for v in data["range"]
907 )
909 return data
912TensorDataDescr = Union[NominalOrOrdinalDataDescr, IntervalOrRatioDataDescr]
915class BinarizeKwargs(KwargsNode):
916 """key word arguments for [BinarizeDescr][]"""
918 threshold: float
919 """The fixed threshold"""
922class BinarizeAlongAxisKwargs(KwargsNode):
923 """key word arguments for [BinarizeDescr][]"""
925 threshold: NotEmpty[list[float]]
926 """The fixed threshold values along `axis`"""
928 axis: Annotated[NonBatchAxisId, Field(examples=["channel"])]
929 """The `threshold` axis"""
932class BinarizeDescr(NodeWithExplicitlySetFields):
933 """Binarize the tensor with a fixed threshold.
935 Values above [BinarizeKwargs.threshold][]/[BinarizeAlongAxisKwargs.threshold][]
936 will be set to one, values below the threshold to zero.
938 Examples:
939 - in YAML
940 ```yaml
941 postprocessing:
942 - id: binarize
943 kwargs:
944 axis: 'channel'
945 threshold: [0.25, 0.5, 0.75]
946 ```
947 - in Python:
949 >>> postprocessing = [BinarizeDescr(
950 ... kwargs=BinarizeAlongAxisKwargs(
951 ... axis=AxisId('channel'),
952 ... threshold=[0.25, 0.5, 0.75],
953 ... )
954 ... )]
955 """
957 implemented_id: ClassVar[Literal["binarize"]] = "binarize"
958 if TYPE_CHECKING:
959 id: Literal["binarize"] = "binarize"
960 else:
961 id: Literal["binarize"]
962 kwargs: BinarizeKwargs | BinarizeAlongAxisKwargs
965class ClipKwargs(KwargsNode):
966 """key word arguments for [ClipDescr][]"""
968 min: float | None = None
969 """Minimum value for clipping.
971 Exclusive with [min_percentile][]
972 """
973 min_percentile: Annotated[float, Interval(ge=0, lt=100)] | None = None
974 """Minimum percentile for clipping.
976 Exclusive with [min][].
978 In range [0, 100).
979 """
981 max: float | None = None
982 """Maximum value for clipping.
984 Exclusive with `max_percentile`.
985 """
986 max_percentile: Annotated[float, Interval(gt=1, le=100)] | None = None
987 """Maximum percentile for clipping.
989 Exclusive with `max`.
991 In range (1, 100].
992 """
994 axes: Annotated[Sequence[AxisId] | None, Field(examples=[("batch", "x", "y")])] = (
995 None
996 )
997 """The subset of axes to determine percentiles jointly,
999 i.e. axes to reduce to compute min/max from `min_percentile`/`max_percentile`.
1000 For example to clip 'batch', 'x' and 'y' jointly in a tensor ('batch', 'channel', 'y', 'x')
1001 resulting in a tensor of equal shape with clipped values per channel, specify `axes=('batch', 'x', 'y')`.
1002 To clip samples independently, leave out the 'batch' axis.
1004 Only valid if `min_percentile` and/or `max_percentile` are set.
1006 Default: Compute percentiles over all axes jointly."""
1008 @model_validator(mode="after")
1009 def _validate(self) -> Self:
1010 if (self.min is not None) and (self.min_percentile is not None):
1011 raise ValueError(
1012 "Only one of `min` and `min_percentile` may be set, not both."
1013 )
1014 if (self.max is not None) and (self.max_percentile is not None):
1015 raise ValueError(
1016 "Only one of `max` and `max_percentile` may be set, not both."
1017 )
1018 if (
1019 self.min is None
1020 and self.min_percentile is None
1021 and self.max is None
1022 and self.max_percentile is None
1023 ):
1024 raise ValueError(
1025 "At least one of `min`, `min_percentile`, `max`, or `max_percentile` must be set."
1026 )
1028 if (
1029 self.axes is not None
1030 and self.min_percentile is None
1031 and self.max_percentile is None
1032 ):
1033 raise ValueError(
1034 "If `axes` is set, at least one of `min_percentile` or `max_percentile` must be set."
1035 )
1037 return self
1040class ClipDescr(NodeWithExplicitlySetFields):
1041 """Set tensor values below min to min and above max to max.
1043 See `ScaleRangeDescr` for examples.
1044 """
1046 implemented_id: ClassVar[Literal["clip"]] = "clip"
1047 if TYPE_CHECKING:
1048 id: Literal["clip"] = "clip"
1049 else:
1050 id: Literal["clip"]
1052 kwargs: ClipKwargs
1055class EnsureDtypeKwargs(KwargsNode):
1056 """key word arguments for [EnsureDtypeDescr][]"""
1058 dtype: Literal[
1059 "float32",
1060 "float64",
1061 "uint8",
1062 "int8",
1063 "uint16",
1064 "int16",
1065 "uint32",
1066 "int32",
1067 "uint64",
1068 "int64",
1069 "bool",
1070 ]
1073class EnsureDtypeDescr(NodeWithExplicitlySetFields):
1074 """Cast the tensor data type to `EnsureDtypeKwargs.dtype` (if not matching).
1076 This can for example be used to ensure the inner neural network model gets a
1077 different input tensor data type than the fully described bioimage.io model does.
1079 Examples:
1080 The described bioimage.io model (incl. preprocessing) accepts any
1081 float32-compatible tensor, normalizes it with percentiles and clipping and then
1082 casts it to uint8, which is what the neural network in this example expects.
1083 - in YAML
1084 ```yaml
1085 inputs:
1086 - data:
1087 type: float32 # described bioimage.io model is compatible with any float32 input tensor
1088 preprocessing:
1089 - id: scale_range
1090 kwargs:
1091 axes: ['y', 'x']
1092 max_percentile: 99.8
1093 min_percentile: 5.0
1094 - id: clip
1095 kwargs:
1096 min: 0.0
1097 max: 1.0
1098 - id: ensure_dtype # the neural network of the model requires uint8
1099 kwargs:
1100 dtype: uint8
1101 ```
1102 - in Python:
1103 >>> preprocessing = [
1104 ... ScaleRangeDescr(
1105 ... kwargs=ScaleRangeKwargs(
1106 ... axes= (AxisId('y'), AxisId('x')),
1107 ... max_percentile= 99.8,
1108 ... min_percentile= 5.0,
1109 ... )
1110 ... ),
1111 ... ClipDescr(kwargs=ClipKwargs(min=0.0, max=1.0)),
1112 ... EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype="uint8")),
1113 ... ]
1114 """
1116 implemented_id: ClassVar[Literal["ensure_dtype"]] = "ensure_dtype"
1117 if TYPE_CHECKING:
1118 id: Literal["ensure_dtype"] = "ensure_dtype"
1119 else:
1120 id: Literal["ensure_dtype"]
1122 kwargs: EnsureDtypeKwargs
1125class ScaleLinearKwargs(KwargsNode):
1126 """Key word arguments for [ScaleLinearDescr][]"""
1128 gain: float = 1.0
1129 """multiplicative factor"""
1131 offset: float = 0.0
1132 """additive term"""
1134 @model_validator(mode="after")
1135 def _validate(self) -> Self:
1136 if self.gain == 1.0 and self.offset == 0.0:
1137 raise ValueError(
1138 "Redundant linear scaling not allowd. Set `gain` != 1.0 and/or `offset`"
1139 + " != 0.0."
1140 )
1142 return self
1145class ScaleLinearAlongAxisKwargs(KwargsNode):
1146 """Key word arguments for [ScaleLinearDescr][]"""
1148 axis: Annotated[NonBatchAxisId, Field(examples=["channel"])]
1149 """The axis of gain and offset values."""
1151 gain: float | NotEmpty[list[float]] = 1.0
1152 """multiplicative factor"""
1154 offset: float | NotEmpty[list[float]] = 0.0
1155 """additive term"""
1157 @model_validator(mode="after")
1158 def _validate(self) -> Self:
1159 if isinstance(self.gain, list):
1160 if isinstance(self.offset, list):
1161 if len(self.gain) != len(self.offset):
1162 raise ValueError(
1163 f"Size of `gain` ({len(self.gain)}) and `offset` ({len(self.offset)}) must match."
1164 )
1165 else:
1166 self.offset = [float(self.offset)] * len(self.gain)
1167 elif isinstance(self.offset, list):
1168 self.gain = [float(self.gain)] * len(self.offset)
1169 else:
1170 raise ValueError(
1171 "Do not specify an `axis` for scalar gain and offset values."
1172 )
1174 if all(g == 1.0 for g in self.gain) and all(off == 0.0 for off in self.offset):
1175 raise ValueError(
1176 "Redundant linear scaling not allowd. Set `gain` != 1.0 and/or `offset`"
1177 + " != 0.0."
1178 )
1180 return self
1183class ScaleLinearDescr(NodeWithExplicitlySetFields):
1184 """Fixed linear scaling.
1186 Examples:
1187 1. Scale with scalar gain and offset
1188 - in YAML
1189 ```yaml
1190 preprocessing:
1191 - id: scale_linear
1192 kwargs:
1193 gain: 2.0
1194 offset: 3.0
1195 ```
1196 - in Python:
1198 >>> preprocessing = [
1199 ... ScaleLinearDescr(kwargs=ScaleLinearKwargs(gain= 2.0, offset=3.0))
1200 ... ]
1202 2. Independent scaling along an axis
1203 - in YAML
1204 ```yaml
1205 preprocessing:
1206 - id: scale_linear
1207 kwargs:
1208 axis: 'channel'
1209 gain: [1.0, 2.0, 3.0]
1210 ```
1211 - in Python:
1213 >>> preprocessing = [
1214 ... ScaleLinearDescr(
1215 ... kwargs=ScaleLinearAlongAxisKwargs(
1216 ... axis=AxisId("channel"),
1217 ... gain=[1.0, 2.0, 3.0],
1218 ... )
1219 ... )
1220 ... ]
1222 """
1224 implemented_id: ClassVar[Literal["scale_linear"]] = "scale_linear"
1225 if TYPE_CHECKING:
1226 id: Literal["scale_linear"] = "scale_linear"
1227 else:
1228 id: Literal["scale_linear"]
1229 kwargs: ScaleLinearKwargs | ScaleLinearAlongAxisKwargs
1232class SigmoidDescr(NodeWithExplicitlySetFields):
1233 """The logistic sigmoid function, a.k.a. expit function.
1235 Examples:
1236 - in YAML
1237 ```yaml
1238 postprocessing:
1239 - id: sigmoid
1240 ```
1241 - in Python:
1243 >>> postprocessing = [SigmoidDescr()]
1244 """
1246 implemented_id: ClassVar[Literal["sigmoid"]] = "sigmoid"
1247 if TYPE_CHECKING:
1248 id: Literal["sigmoid"] = "sigmoid"
1249 else:
1250 id: Literal["sigmoid"]
1252 @property
1253 def kwargs(self) -> KwargsNode:
1254 """empty kwargs"""
1255 return KwargsNode()
1258class SoftmaxKwargs(KwargsNode):
1259 """key word arguments for [SoftmaxDescr][]"""
1261 axis: Annotated[NonBatchAxisId, Field(examples=["channel"])] = CHANNEL_AXIS_ID
1262 """The axis to apply the softmax function along.
1263 Note:
1264 Defaults to 'channel' axis
1265 (which may not exist, in which case
1266 a different axis id has to be specified).
1267 """
1270class SoftmaxDescr(NodeWithExplicitlySetFields):
1271 """The softmax function.
1273 Examples:
1274 - in YAML
1275 ```yaml
1276 postprocessing:
1277 - id: softmax
1278 kwargs:
1279 axis: channel
1280 ```
1281 - in Python:
1283 >>> postprocessing = [SoftmaxDescr(kwargs=SoftmaxKwargs(axis=AxisId("channel")))]
1284 """
1286 implemented_id: ClassVar[Literal["softmax"]] = "softmax"
1287 if TYPE_CHECKING:
1288 id: Literal["softmax"] = "softmax"
1289 else:
1290 id: Literal["softmax"]
1292 kwargs: SoftmaxKwargs = Field(default_factory=SoftmaxKwargs.model_construct)
1295class _StardistPostprocessingKwargsBase(KwargsNode):
1296 """key word arguments for [StardistPostprocessingDescr][]"""
1298 prob_threshold: float
1299 """The probability threshold for object candidate selection."""
1301 nms_threshold: float
1302 """The IoU threshold for non-maximum suppression."""
1304 n_rays: int
1305 """Number of radial lines (rays) cast from the center of an object to its boundary."""
1308class StardistPostprocessingKwargs2D(_StardistPostprocessingKwargsBase):
1309 grid: tuple[int, int]
1310 """Grid size of network predictions."""
1312 b: int | tuple[tuple[int, int], tuple[int, int]]
1313 """Border region in which object probability is set to zero."""
1316class StardistPostprocessingKwargs3D(_StardistPostprocessingKwargsBase):
1317 grid: tuple[int, int, int]
1318 """Grid size of network predictions."""
1320 b: int | tuple[tuple[int, int], tuple[int, int], tuple[int, int]]
1321 """Border region in which object probability is set to zero."""
1323 anisotropy: tuple[float, float, float]
1324 """Anisotropy factors for 3D star-convex polyhedra, i.e. the physical pixel size along each spatial axis."""
1326 overlap_label: int | None = None
1327 """Optional label to apply to any area of overlapping predicted objects."""
1330class StardistPostprocessingDescr(NodeWithExplicitlySetFields):
1331 """Stardist postprocessing including non-maximum suppression and converting polygon representations to instance labels
1333 as described in:
1334 - Uwe Schmidt, Martin Weigert, Coleman Broaddus, and Gene Myers.
1335 [*Cell Detection with Star-convex Polygons*](https://arxiv.org/abs/1806.03535).
1336 International Conference on Medical Image Computing and Computer-Assisted Intervention (MICCAI), Granada, Spain, September 2018.
1337 - Martin Weigert, Uwe Schmidt, Robert Haase, Ko Sugawara, and Gene Myers.
1338 [*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).
1339 The IEEE Winter Conference on Applications of Computer Vision (WACV), Snowmass Village, Colorado, March 2020.
1341 Note: Only available if the `stardist` package is installed.
1342 """
1344 implemented_id: ClassVar[Literal["stardist_postprocessing"]] = (
1345 "stardist_postprocessing"
1346 )
1347 if TYPE_CHECKING:
1348 id: Literal["stardist_postprocessing"] = "stardist_postprocessing"
1349 else:
1350 id: Literal["stardist_postprocessing"]
1352 kwargs: StardistPostprocessingKwargs2D | StardistPostprocessingKwargs3D
1355class CellposeFlowDynamicsKwargs(KwargsNode):
1356 """key word arguments for [CellposeFlowDynamicsDescr][]"""
1358 cellprob_threshold: float
1359 flow_threshold: float
1360 do_3D: bool
1361 min_size: int = 15
1362 """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."""
1363 output_dtype: Literal["uint16", "uint32"] = "uint16"
1366class CellposeFlowDynamicsDescr(NodeWithExplicitlySetFields):
1367 """Cellpose flow dynamics postprocessing as described in:
1368 - Carsen Stringer and Marius Pachitariu. [*Cellpose: a generalist algorithm for cellular segmentation*](https://www.nature.com/articles/s41592-020-01018-x). Nature Methods, 2021.
1370 Note: Only available if the `cellpose` package is installed.
1371 """
1373 implemented_id: ClassVar[Literal["cellpose_flow_dynamics"]] = (
1374 "cellpose_flow_dynamics"
1375 )
1376 if TYPE_CHECKING:
1377 id: Literal["cellpose_flow_dynamics"] = "cellpose_flow_dynamics"
1378 else:
1379 id: Literal["cellpose_flow_dynamics"]
1381 kwargs: CellposeFlowDynamicsKwargs
1384class CustomProcessingDescr(NodeWithExplicitlySetFields, FileDescr):
1385 """Custom (post)processing op — source file shipped inline with the model.
1387 Supports (post)processing that cannot be expressed by the built-in named
1388 operations (watershed, connected components, etc.)
1389 using a simple Python callable interface.
1391 The op is implemented in a ``.py`` file packaged alongside the model weights.
1392 Two styles are supported:
1394 *Callable class* — kwargs go to ``__init__``, tensors arrive in ``__call__``:
1396 .. code-block:: python
1398 # my_postprocess.py
1399 import numpy as np
1401 class my_postprocess:
1402 def __init__(self, threshold: float = 0.5) -> None:
1403 self.threshold = threshold
1404 def __call__(self, *arrays: np.ndarray) -> np.ndarray:
1405 # arrays = model output tensors in rdf.yaml declaration order
1406 return (arrays[0] > self.threshold).astype(np.uint8)
1408 *Factory function* — alternative closure style, identical runtime behaviour:
1410 .. code-block:: python
1412 # my_postprocess.py
1413 import numpy as np
1415 def my_postprocess(threshold: float = 0.5):
1416 def run(*arrays: np.ndarray) -> np.ndarray:
1417 return (arrays[0] > threshold).astype(np.uint8)
1418 return run
1420 Reference it in ``rdf.yaml`` with the source file included in the package:
1422 .. code-block:: yaml
1424 postprocessing:
1425 - id: custom
1426 callable: my_postprocess # class or function name in source
1427 source: my_postprocess.py # packaged alongside weights
1428 sha256: <hash> # sha256 of the source file
1429 kwargs: # forwarded to __init__ / factory
1430 threshold: 0.5
1432 **Security:** source files are SHA-256 verified before execution.
1433 Execution requires explicit opt-in in bioimageio.core and curator
1434 review before Zoo publication.
1435 """
1437 implemented_id: ClassVar[Literal["custom"]] = "custom"
1438 if TYPE_CHECKING:
1439 id: Literal["custom"] = "custom"
1440 else:
1441 id: Literal["custom"]
1443 callable: Annotated[
1444 str,
1445 Field(examples=["my_postprocess_factory", "MyPostprocessClass"]),
1446 ]
1447 """Name of the callable class or factory function defined in ``source``.
1449 At runtime: ``op = callable(**kwargs)``, then ``result = op(*output_tensors)``
1450 per image. Both a class with ``__call__`` and a factory function returning
1451 a callable satisfy this protocol."""
1453 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
1454 """Python source file (included when packaging the model)."""
1456 kwargs: dict[str, YamlValue] = Field(
1457 default_factory=cast(Callable[[], Dict[str, YamlValue]], dict)
1458 )
1459 """Keyword arguments forwarded to the callable (``__init__`` or factory)."""
1461 @model_serializer(mode="wrap", when_used="unless-none")
1462 def _serialize(
1463 self, nxt: SerializerFunctionWrapHandler, info: SerializationInfo
1464 ) -> dict[str, YamlValue]:
1465 return package_file_descr_serializer(self, nxt, info)
1468class FixedZeroMeanUnitVarianceKwargs(KwargsNode):
1469 """key word arguments for [FixedZeroMeanUnitVarianceDescr][]"""
1471 mean: float
1472 """The mean value to normalize with."""
1474 std: Annotated[float, Ge(1e-6)]
1475 """The standard deviation value to normalize with."""
1478class FixedZeroMeanUnitVarianceAlongAxisKwargs(KwargsNode):
1479 """key word arguments for [FixedZeroMeanUnitVarianceDescr][]"""
1481 mean: NotEmpty[list[float]]
1482 """The mean value(s) to normalize with."""
1484 std: NotEmpty[list[Annotated[float, Ge(1e-6)]]]
1485 """The standard deviation value(s) to normalize with.
1486 Size must match `mean` values."""
1488 axis: Annotated[NonBatchAxisId, Field(examples=["channel", "index"])]
1489 """The axis of the mean/std values to normalize each entry along that dimension
1490 separately."""
1492 @model_validator(mode="after")
1493 def _mean_and_std_match(self) -> Self:
1494 if len(self.mean) != len(self.std):
1495 raise ValueError(
1496 f"Size of `mean` ({len(self.mean)}) and `std` ({len(self.std)})"
1497 + " must match."
1498 )
1500 return self
1503class FixedZeroMeanUnitVarianceDescr(NodeWithExplicitlySetFields):
1504 """Subtract a given mean and divide by the standard deviation.
1506 Normalize with fixed, precomputed values for
1507 `FixedZeroMeanUnitVarianceKwargs.mean` and `FixedZeroMeanUnitVarianceKwargs.std`
1508 Use `FixedZeroMeanUnitVarianceAlongAxisKwargs` for independent scaling along given
1509 axes.
1511 Examples:
1512 1. scalar value for whole tensor
1513 - in YAML
1514 ```yaml
1515 preprocessing:
1516 - id: fixed_zero_mean_unit_variance
1517 kwargs:
1518 mean: 103.5
1519 std: 13.7
1520 ```
1521 - in Python
1522 >>> preprocessing = [FixedZeroMeanUnitVarianceDescr(
1523 ... kwargs=FixedZeroMeanUnitVarianceKwargs(mean=103.5, std=13.7)
1524 ... )]
1526 2. independently along an axis
1527 - in YAML
1528 ```yaml
1529 preprocessing:
1530 - id: fixed_zero_mean_unit_variance
1531 kwargs:
1532 axis: channel
1533 mean: [101.5, 102.5, 103.5]
1534 std: [11.7, 12.7, 13.7]
1535 ```
1536 - in Python
1537 >>> preprocessing = [FixedZeroMeanUnitVarianceDescr(
1538 ... kwargs=FixedZeroMeanUnitVarianceAlongAxisKwargs(
1539 ... axis=AxisId("channel"),
1540 ... mean=[101.5, 102.5, 103.5],
1541 ... std=[11.7, 12.7, 13.7],
1542 ... )
1543 ... )]
1544 """
1546 implemented_id: ClassVar[Literal["fixed_zero_mean_unit_variance"]] = (
1547 "fixed_zero_mean_unit_variance"
1548 )
1549 if TYPE_CHECKING:
1550 id: Literal["fixed_zero_mean_unit_variance"] = "fixed_zero_mean_unit_variance"
1551 else:
1552 id: Literal["fixed_zero_mean_unit_variance"]
1554 kwargs: FixedZeroMeanUnitVarianceKwargs | FixedZeroMeanUnitVarianceAlongAxisKwargs
1557class ZeroMeanUnitVarianceKwargs(KwargsNode):
1558 """key word arguments for [ZeroMeanUnitVarianceDescr][]"""
1560 axes: Annotated[Sequence[AxisId] | None, Field(examples=[("batch", "x", "y")])] = (
1561 None
1562 )
1563 """The subset of axes to normalize jointly, i.e. axes to reduce to compute mean/std.
1564 For example to normalize 'batch', 'x' and 'y' jointly in a tensor ('batch', 'channel', 'y', 'x')
1565 resulting in a tensor of equal shape normalized per channel, specify `axes=('batch', 'x', 'y')`.
1566 To normalize each sample independently leave out the 'batch' axis.
1567 Default: Scale all axes jointly."""
1569 eps: Annotated[float, Interval(gt=0, le=0.1)] = 1e-6
1570 """epsilon for numeric stability: `out = (tensor - mean) / (std + eps)`."""
1573class ZeroMeanUnitVarianceDescr(NodeWithExplicitlySetFields):
1574 """Subtract mean and divide by variance.
1576 Examples:
1577 Subtract tensor mean and variance
1578 - in YAML
1579 ```yaml
1580 preprocessing:
1581 - id: zero_mean_unit_variance
1582 ```
1583 - in Python
1584 >>> preprocessing = [ZeroMeanUnitVarianceDescr()]
1585 """
1587 implemented_id: ClassVar[Literal["zero_mean_unit_variance"]] = (
1588 "zero_mean_unit_variance"
1589 )
1590 if TYPE_CHECKING:
1591 id: Literal["zero_mean_unit_variance"] = "zero_mean_unit_variance"
1592 else:
1593 id: Literal["zero_mean_unit_variance"]
1595 kwargs: ZeroMeanUnitVarianceKwargs = Field(
1596 default_factory=ZeroMeanUnitVarianceKwargs.model_construct
1597 )
1600class ScaleRangeKwargs(KwargsNode):
1601 """key word arguments for [ScaleRangeDescr][]
1603 For `min_percentile`=0.0 (the default) and `max_percentile`=100 (the default)
1604 this processing step normalizes data to the [0, 1] intervall.
1605 For other percentiles the normalized values will partially be outside the [0, 1]
1606 intervall. Use `ScaleRange` followed by `ClipDescr` if you want to limit the
1607 normalized values to a range.
1608 """
1610 axes: Annotated[Sequence[AxisId] | None, Field(examples=[("batch", "x", "y")])] = (
1611 None
1612 )
1613 """The subset of axes to normalize jointly, i.e. axes to reduce to compute the min/max percentile value.
1614 For example to normalize 'batch', 'x' and 'y' jointly in a tensor ('batch', 'channel', 'y', 'x')
1615 resulting in a tensor of equal shape normalized per channel, specify `axes=('batch', 'x', 'y')`.
1616 To normalize samples independently, leave out the "batch" axis.
1617 Default: Scale all axes jointly."""
1619 min_percentile: Annotated[float, Interval(ge=0, lt=100)] = 0.0
1620 """The lower percentile used to determine the value to align with zero."""
1622 max_percentile: Annotated[float, Interval(gt=1, le=100)] = 100.0
1623 """The upper percentile used to determine the value to align with one.
1624 Has to be bigger than `min_percentile`.
1625 The range is 1 to 100 instead of 0 to 100 to avoid mistakenly
1626 accepting percentiles specified in the range 0.0 to 1.0."""
1628 eps: Annotated[float, Interval(gt=0, le=0.1)] = 1e-6
1629 """Epsilon for numeric stability.
1630 `out = (tensor - v_lower) / (v_upper - v_lower + eps)`;
1631 with `v_lower,v_upper` values at the respective percentiles."""
1633 reference_tensor: TensorId | None = None
1634 """ID of the unprocessed input tensor to compute the percentiles from.
1635 Default: The tensor itself.
1636 """
1638 @field_validator("max_percentile", mode="after")
1639 @classmethod
1640 def min_smaller_max(cls, value: float, info: ValidationInfo) -> float:
1641 if (min_p := info.data["min_percentile"]) >= value:
1642 raise ValueError(f"min_percentile {min_p} >= max_percentile {value}")
1644 return value
1647class ScaleRangeDescr(NodeWithExplicitlySetFields):
1648 """Scale with percentiles.
1650 Examples:
1651 1. Scale linearly to map 5th percentile to 0 and 99.8th percentile to 1.0
1652 - in YAML
1653 ```yaml
1654 preprocessing:
1655 - id: scale_range
1656 kwargs:
1657 axes: ['y', 'x']
1658 max_percentile: 99.8
1659 min_percentile: 5.0
1660 ```
1661 - in Python
1663 >>> preprocessing = [
1664 ... ScaleRangeDescr(
1665 ... kwargs=ScaleRangeKwargs(
1666 ... axes= (AxisId('y'), AxisId('x')),
1667 ... max_percentile= 99.8,
1668 ... min_percentile= 5.0,
1669 ... )
1670 ... )
1671 ... ]
1673 2. Combine the above scaling with additional clipping to clip values outside the range given by the percentiles.
1674 - in YAML
1675 ```yaml
1676 preprocessing:
1677 - id: scale_range
1678 kwargs:
1679 axes: ['y', 'x']
1680 max_percentile: 99.8
1681 min_percentile: 5.0
1682 - id: clip
1683 kwargs:
1684 min: 0.0
1685 max: 1.0
1686 ```
1687 - in Python
1689 >>> preprocessing = [
1690 ... ScaleRangeDescr(
1691 ... kwargs=ScaleRangeKwargs(
1692 ... axes= (AxisId('y'), AxisId('x')),
1693 ... max_percentile= 99.8,
1694 ... min_percentile= 5.0,
1695 ... )
1696 ... ),
1697 ... ClipDescr(
1698 ... kwargs=ClipKwargs(
1699 ... min=0.0,
1700 ... max=1.0,
1701 ... )
1702 ... ),
1703 ... ]
1705 """
1707 implemented_id: ClassVar[Literal["scale_range"]] = "scale_range"
1708 if TYPE_CHECKING:
1709 id: Literal["scale_range"] = "scale_range"
1710 else:
1711 id: Literal["scale_range"]
1712 kwargs: ScaleRangeKwargs = Field(default_factory=ScaleRangeKwargs.model_construct)
1715class ScaleMeanVarianceKwargs(KwargsNode):
1716 """key word arguments for [ScaleMeanVarianceKwargs][]"""
1718 reference_tensor: TensorId
1719 """ID of unprocessed input tensor to match."""
1721 axes: Annotated[Sequence[AxisId] | None, Field(examples=[("batch", "x", "y")])] = (
1722 None
1723 )
1724 """The subset of axes to normalize jointly, i.e. axes to reduce to compute mean/std.
1725 For example to normalize 'batch', 'x' and 'y' jointly in a tensor ('batch', 'channel', 'y', 'x')
1726 resulting in a tensor of equal shape normalized per channel, specify `axes=('batch', 'x', 'y')`.
1727 To normalize samples independently, leave out the 'batch' axis.
1728 Default: Scale all axes jointly."""
1730 eps: Annotated[float, Interval(gt=0, le=0.1)] = 1e-6
1731 """Epsilon for numeric stability:
1732 `out = (tensor - mean) / (std + eps) * (ref_std + eps) + ref_mean.`"""
1735class ScaleMeanVarianceDescr(NodeWithExplicitlySetFields):
1736 """Scale a tensor's data distribution to match another tensor's mean/std.
1737 `out = (tensor - mean) / (std + eps) * (ref_std + eps) + ref_mean.`
1738 """
1740 implemented_id: ClassVar[Literal["scale_mean_variance"]] = "scale_mean_variance"
1741 if TYPE_CHECKING:
1742 id: Literal["scale_mean_variance"] = "scale_mean_variance"
1743 else:
1744 id: Literal["scale_mean_variance"]
1745 kwargs: ScaleMeanVarianceKwargs
1748PreprocessingDescr = Annotated[
1749 Union[
1750 BinarizeDescr,
1751 ClipDescr,
1752 EnsureDtypeDescr,
1753 FixedZeroMeanUnitVarianceDescr,
1754 ScaleLinearDescr,
1755 ScaleRangeDescr,
1756 SigmoidDescr,
1757 SoftmaxDescr,
1758 ZeroMeanUnitVarianceDescr,
1759 ],
1760 Discriminator("id"),
1761]
1762PostprocessingDescr = Annotated[
1763 Union[
1764 BinarizeDescr,
1765 CellposeFlowDynamicsDescr,
1766 ClipDescr,
1767 CustomProcessingDescr,
1768 EnsureDtypeDescr,
1769 FixedZeroMeanUnitVarianceDescr,
1770 ScaleLinearDescr,
1771 ScaleMeanVarianceDescr,
1772 ScaleRangeDescr,
1773 SigmoidDescr,
1774 SoftmaxDescr,
1775 StardistPostprocessingDescr,
1776 ZeroMeanUnitVarianceDescr,
1777 ],
1778 Discriminator("id"),
1779]
1781IO_AxisT = TypeVar("IO_AxisT", InputAxis, OutputAxis)
1784class TensorDescrBase(Node, Generic[IO_AxisT]):
1785 id: TensorId
1786 """Tensor id. No duplicates are allowed."""
1788 description: Annotated[str, MaxLen(128)] = ""
1789 """free text description"""
1791 axes: NotEmpty[Sequence[IO_AxisT]]
1792 """tensor axes"""
1794 @property
1795 def shape(self):
1796 return tuple(a.size for a in self.axes)
1798 @field_validator("axes", mode="after", check_fields=False)
1799 @classmethod
1800 def _validate_axes(cls, axes: Sequence[AnyAxis]) -> Sequence[AnyAxis]:
1801 batch_axes = [a for a in axes if a.type == "batch"]
1802 if len(batch_axes) > 1:
1803 raise ValueError(
1804 f"Only one batch axis (per tensor) allowed, but got {batch_axes}"
1805 )
1807 seen_ids: set[AxisId] = set()
1808 duplicate_axes_ids: set[AxisId] = set()
1809 for a in axes:
1810 (duplicate_axes_ids if a.id in seen_ids else seen_ids).add(a.id)
1812 if duplicate_axes_ids:
1813 raise ValueError(f"Duplicate axis ids: {duplicate_axes_ids}")
1815 return axes
1817 test_tensor: FAIR[FileDescr_package | None] = None
1818 """An example tensor to use for testing.
1819 Using the model with the test input tensors is expected to yield the test output tensors.
1820 Each test tensor has be a an ndarray in the
1821 [numpy.lib file format](https://numpy.org/doc/stable/reference/generated/numpy.lib.format.html#module-numpy.lib.format).
1822 The file extension must be '.npy'."""
1824 sample_tensor: FAIR[FileDescr_package | None] = None
1825 """A sample tensor to illustrate a possible input/output for the model,
1826 The sample image primarily serves to inform a human user about an example use case
1827 and is typically stored as .hdf5, .png or .tiff.
1828 It has to be readable by the [imageio library](https://imageio.readthedocs.io/en/stable/formats/index.html#supported-formats)
1829 (numpy's `.npy` format is not supported).
1830 The image dimensionality has to match the number of axes specified in this tensor description.
1831 """
1833 @model_validator(mode="after")
1834 def _validate_sample_tensor(self) -> Self:
1835 if self.sample_tensor is None or not get_validation_context().perform_io_checks:
1836 return self
1838 reader = get_reader(self.sample_tensor.source, sha256=self.sample_tensor.sha256)
1839 tensor: NDArray[Any] = imread(
1840 reader.read(),
1841 extension=PurePosixPath(reader.original_file_name).suffix,
1842 )
1843 n_dims = len(tensor.squeeze().shape)
1844 n_dims_min = n_dims_max = len(self.axes)
1846 for a in self.axes:
1847 if isinstance(a, BatchAxis):
1848 n_dims_min -= 1
1849 elif isinstance(a.size, int):
1850 if a.size == 1:
1851 n_dims_min -= 1
1852 elif isinstance(a.size, (ParameterizedSize, DataDependentSize)):
1853 if a.size.min == 1:
1854 n_dims_min -= 1
1855 elif isinstance(a.size, SizeReference):
1856 if a.id != a.size.axis_id:
1857 raise ValueError(
1858 f"Axis {a.id} has a size reference to a different axis {a.size.axis_id}."
1859 )
1860 if a.size.offset < 2:
1861 # size reference may result in singleton axis
1862 n_dims_min -= 1
1863 else:
1864 assert_never(a.size)
1866 n_dims_min = max(0, n_dims_min)
1867 if n_dims < n_dims_min or n_dims > n_dims_max:
1868 raise ValueError(
1869 f"Expected sample tensor to have {n_dims_min} to"
1870 + f" {n_dims_max} dimensions, but found {n_dims} (shape: {tensor.shape})."
1871 )
1873 return self
1875 data: TensorDataDescr | NotEmpty[Sequence[TensorDataDescr]] = (
1876 IntervalOrRatioDataDescr()
1877 )
1878 """Description of the tensor's data values, optionally per channel.
1879 If specified per channel, the data `type` needs to match across channels."""
1881 @property
1882 def dtype(
1883 self,
1884 ) -> Literal[
1885 "float32",
1886 "float64",
1887 "uint8",
1888 "int8",
1889 "uint16",
1890 "int16",
1891 "uint32",
1892 "int32",
1893 "uint64",
1894 "int64",
1895 "bool",
1896 ]:
1897 """dtype as specified under `data.type` or `data[i].type`"""
1898 if isinstance(self.data, collections.abc.Sequence):
1899 return self.data[0].type
1900 else:
1901 return self.data.type
1903 @field_validator("data", mode="after")
1904 @classmethod
1905 def _check_data_type_across_channels(
1906 cls, value: TensorDataDescr | NotEmpty[Sequence[TensorDataDescr]]
1907 ) -> TensorDataDescr | NotEmpty[Sequence[TensorDataDescr]]:
1908 if not isinstance(value, list):
1909 return value
1911 dtypes = {t.type for t in value}
1912 if len(dtypes) > 1:
1913 raise ValueError(
1914 "Tensor data descriptions per channel need to agree in their data"
1915 + f" `type`, but found {dtypes}."
1916 )
1918 return value
1920 @model_validator(mode="after")
1921 def _check_data_matches_channelaxis(self) -> Self:
1922 if not isinstance(self.data, (list, tuple)):
1923 return self
1925 for a in self.axes:
1926 if isinstance(a, ChannelAxis):
1927 size = a.size
1928 assert isinstance(size, int)
1929 break
1930 else:
1931 return self
1933 if len(self.data) != size:
1934 raise ValueError(
1935 f"Got tensor data descriptions for {len(self.data)} channels, but"
1936 + f" '{a.id}' axis has size {size}."
1937 )
1939 return self
1941 def get_axis_sizes_for_array(self, array: NDArray[Any]) -> dict[AxisId, int]:
1942 if len(array.shape) != len(self.axes):
1943 raise ValueError(
1944 f"Dimension mismatch: array shape {array.shape} (#{len(array.shape)})"
1945 + f" incompatible with {len(self.axes)} axes."
1946 )
1947 return {a.id: array.shape[i] for i, a in enumerate(self.axes)}
1950class ConstantPadding(Node):
1951 mode: Literal["constant"] = "constant"
1952 value: int | float = 0
1955class EdgePadding(Node):
1956 mode: Literal["edge"] = "edge"
1959class ReflectPadding(Node):
1960 mode: Literal["reflect"] = "reflect"
1963class SymmetricPadding(Node):
1964 mode: Literal["symmetric"] = "symmetric"
1967Padding: TypeAlias = Union[
1968 ConstantPadding, EdgePadding, ReflectPadding, SymmetricPadding
1969]
1972class ModelId(ResourceId):
1973 pass
1976class InputTensorDescr(TensorDescrBase[InputAxis]):
1977 id: TensorId = TensorId("input")
1978 """Input tensor id.
1979 No duplicates are allowed across all inputs and outputs."""
1981 output_of: ModelId | None = None
1982 """If this input tensor is the output of another model, specify the model id here.
1983 This model's input id must match the output id of the referenced model.
1984 """
1986 @model_validator(mode="after")
1987 def _validate_output_of(self) -> Self:
1988 if self.output_of is None:
1989 return self
1991 try:
1992 with get_validation_context().replace(perform_io_checks=False):
1993 opened_ref_model = open_bioimageio_yaml(self.output_of)
1994 format_version = opened_ref_model.content["format_version"]
1995 assert isinstance(format_version, str)
1996 if format_version.startswith("0.4"):
1997 ref_model = _ModelDescr04.model_validate(opened_ref_model.content)
1998 else:
1999 ref_model = ModelDescr.model_validate(opened_ref_model.content)
2000 except Exception as e:
2001 raise ValueError(
2002 f"Failed to load model '{self.output_of}' referenced under output_of: {e}"
2003 )
2005 try:
2006 ref_model_outputs = {
2007 t.id if isinstance(t, OutputTensorDescr) else TensorId(t.name)
2008 for t in ref_model.outputs
2009 }
2010 except Exception as e:
2011 raise ValueError(
2012 f"Failed to read output IDs of model '{self.output_of}' referenced under output_of: {e}"
2013 )
2015 if self.id not in ref_model_outputs:
2016 raise ValueError(
2017 f"Input tensor '{self.id}' is specified as output of model '{self.output_of}', "
2018 + f"but that model's outputs are {ref_model_outputs}."
2019 )
2020 return self
2022 optional: bool = False
2023 """indicates that this tensor may be `None`"""
2025 pad: Padding | None = None
2026 """Explicitly specify how to pad this input tensor.
2028 Use `axes[i].pad` to specify padding width.
2030 Note:
2031 Non-blockwise sample prediction only applies padding for axes with a `pad` specification.
2032 """
2034 preprocessing: list[PreprocessingDescr] = Field(
2035 default_factory=cast(Callable[[], List[PreprocessingDescr]], list)
2036 )
2037 """Description of how this input should be preprocessed.
2039 notes:
2040 - If preprocessing does not start with an 'ensure_dtype' entry, it is added
2041 to ensure an input tensor's data type matches the input tensor's data description.
2042 - If preprocessing does not end with an 'ensure_dtype' or 'binarize' entry, an
2043 'ensure_dtype' step is added to ensure preprocessing steps are not unintentionally
2044 changing the data type.
2045 """
2047 @model_validator(mode="after")
2048 def _validate_preprocessing_kwargs(self) -> Self:
2049 axes_ids = [a.id for a in self.axes]
2050 for p in self.preprocessing:
2051 kwargs_axes: Sequence[Any] | None = p.kwargs.get("axes")
2052 if kwargs_axes is None:
2053 continue
2055 if not isinstance(kwargs_axes, collections.abc.Sequence):
2056 raise ValueError(
2057 f"Expected `preprocessing.i.kwargs.axes` to be a sequence, but got {type(kwargs_axes)}"
2058 )
2060 if any(a not in axes_ids for a in kwargs_axes):
2061 raise ValueError(
2062 "`preprocessing.i.kwargs.axes` needs to be subset of axes ids"
2063 )
2065 if isinstance(self.data, (NominalOrOrdinalDataDescr, IntervalOrRatioDataDescr)):
2066 dtype = self.data.type
2067 else:
2068 dtype = self.data[0].type
2070 # ensure `preprocessing` begins with `EnsureDtypeDescr`
2071 if not self.preprocessing or not isinstance(
2072 self.preprocessing[0], EnsureDtypeDescr
2073 ):
2074 self.preprocessing.insert(
2075 0, EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype=dtype))
2076 )
2078 # ensure `preprocessing` ends with `EnsureDtypeDescr` or `BinarizeDescr`
2079 if not isinstance(self.preprocessing[-1], (EnsureDtypeDescr, BinarizeDescr)):
2080 self.preprocessing.append(
2081 EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype=dtype))
2082 )
2084 return self
2087def convert_axes(
2088 axes: str,
2089 *,
2090 shape: Sequence[int] | _ParameterizedInputShape_v0_4 | _ImplicitOutputShape_v0_4,
2091 tensor_type: Literal["input", "output"],
2092 halo: Sequence[int] | None,
2093 size_refs: Mapping[_TensorName_v0_4, Mapping[str, int]],
2094):
2095 ret: list[AnyAxis] = []
2096 for i, a in enumerate(axes):
2097 axis_type = _AXIS_TYPE_MAP.get(a, a)
2098 if axis_type == "batch":
2099 ret.append(BatchAxis())
2100 continue
2102 scale = 1.0
2103 if isinstance(shape, _ParameterizedInputShape_v0_4):
2104 if shape.step[i] == 0:
2105 size = shape.min[i]
2106 else:
2107 size = ParameterizedSize(min=shape.min[i], step=shape.step[i])
2108 elif isinstance(shape, _ImplicitOutputShape_v0_4):
2109 ref_t = str(shape.reference_tensor)
2110 if ref_t.count(".") == 1:
2111 t_id, orig_a_id = ref_t.split(".")
2112 else:
2113 t_id = ref_t
2114 orig_a_id = a
2116 a_id = _AXIS_ID_MAP.get(orig_a_id, a)
2117 if not (orig_scale := shape.scale[i]):
2118 # old way to insert a new axis dimension
2119 size = int(2 * shape.offset[i])
2120 else:
2121 scale = 1 / orig_scale
2122 if axis_type in ("channel", "index"):
2123 # these axes no longer have a scale
2124 offset_from_scale = orig_scale * size_refs.get(
2125 _TensorName_v0_4(t_id), {}
2126 ).get(orig_a_id, 0)
2127 else:
2128 offset_from_scale = 0
2129 size = SizeReference(
2130 tensor_id=TensorId(t_id),
2131 axis_id=AxisId(a_id),
2132 offset=int(offset_from_scale + 2 * shape.offset[i]),
2133 )
2134 else:
2135 size = shape[i]
2137 if axis_type == "time":
2138 if tensor_type == "input":
2139 ret.append(TimeInputAxis(size=size, scale=scale))
2140 else:
2141 assert not isinstance(size, ParameterizedSize)
2142 if halo is None:
2143 ret.append(TimeOutputAxis(size=size, scale=scale))
2144 else:
2145 assert not isinstance(size, int)
2146 ret.append(
2147 TimeOutputAxisWithHalo(size=size, scale=scale, halo=halo[i])
2148 )
2150 elif axis_type == "index":
2151 if tensor_type == "input":
2152 ret.append(IndexInputAxis(size=size))
2153 else:
2154 if isinstance(size, ParameterizedSize):
2155 size = DataDependentSize(min=size.min)
2157 ret.append(IndexOutputAxis(size=size))
2158 elif axis_type == "channel":
2159 assert not isinstance(size, ParameterizedSize)
2160 if isinstance(size, SizeReference):
2161 warnings.warn(
2162 "Conversion of channel size from an implicit output shape may be"
2163 + " wrong"
2164 )
2165 ret.append(
2166 ChannelAxis(
2167 channel_names=[f"channel{i}" for i in range(size.offset)]
2168 )
2169 )
2170 else:
2171 ret.append(
2172 ChannelAxis(channel_names=[f"channel{i}" for i in range(size)])
2173 )
2174 elif axis_type == "space":
2175 if tensor_type == "input":
2176 ret.append(SpaceInputAxis(id=AxisId(a), size=size, scale=scale))
2177 else:
2178 assert not isinstance(size, ParameterizedSize)
2179 if halo is None or halo[i] == 0:
2180 ret.append(SpaceOutputAxis(id=AxisId(a), size=size, scale=scale))
2181 elif isinstance(size, int):
2182 raise NotImplementedError(
2183 f"output axis with halo and fixed size (here {size}) not allowed"
2184 )
2185 else:
2186 ret.append(
2187 SpaceOutputAxisWithHalo(
2188 id=AxisId(a), size=size, scale=scale, halo=halo[i]
2189 )
2190 )
2192 return ret
2195def _axes_letters_to_ids(
2196 axes: str | None,
2197) -> list[AxisId] | None:
2198 if axes is None:
2199 return None
2201 return [AxisId(a) for a in axes]
2204def _get_complement_v04_axis(
2205 tensor_axes: Sequence[str], axes: Sequence[str] | None
2206) -> AxisId | None:
2207 if axes is None:
2208 return None
2210 non_complement_axes = set(axes) | {"b"}
2211 complement_axes = [a for a in tensor_axes if a not in non_complement_axes]
2212 if len(complement_axes) > 1:
2213 raise ValueError(
2214 f"Expected none or a single complement axis, but axes '{axes}' "
2215 + f"for tensor dims '{tensor_axes}' leave '{complement_axes}'."
2216 )
2218 return None if not complement_axes else AxisId(complement_axes[0])
2221def _convert_proc(
2222 p: _PreprocessingDescr_v0_4 | _PostprocessingDescr_v0_4,
2223 tensor_axes: Sequence[str],
2224) -> PreprocessingDescr | PostprocessingDescr:
2225 if isinstance(p, _BinarizeDescr_v0_4):
2226 return BinarizeDescr(kwargs=BinarizeKwargs(threshold=p.kwargs.threshold))
2227 elif isinstance(p, _ClipDescr_v0_4):
2228 return ClipDescr(kwargs=ClipKwargs(min=p.kwargs.min, max=p.kwargs.max))
2229 elif isinstance(p, _SigmoidDescr_v0_4):
2230 return SigmoidDescr()
2231 elif isinstance(p, _ScaleLinearDescr_v0_4):
2232 axes = _axes_letters_to_ids(p.kwargs.axes)
2233 if p.kwargs.axes is None:
2234 axis = None
2235 else:
2236 axis = _get_complement_v04_axis(tensor_axes, p.kwargs.axes)
2238 if axis is None:
2239 assert not isinstance(p.kwargs.gain, list)
2240 assert not isinstance(p.kwargs.offset, list)
2241 kwargs = ScaleLinearKwargs(gain=p.kwargs.gain, offset=p.kwargs.offset)
2242 else:
2243 kwargs = ScaleLinearAlongAxisKwargs(
2244 axis=axis, gain=p.kwargs.gain, offset=p.kwargs.offset
2245 )
2246 return ScaleLinearDescr(kwargs=kwargs)
2247 elif isinstance(p, _ScaleMeanVarianceDescr_v0_4):
2248 return ScaleMeanVarianceDescr(
2249 kwargs=ScaleMeanVarianceKwargs(
2250 axes=_axes_letters_to_ids(p.kwargs.axes),
2251 reference_tensor=TensorId(str(p.kwargs.reference_tensor)),
2252 eps=p.kwargs.eps,
2253 )
2254 )
2255 elif isinstance(p, _ZeroMeanUnitVarianceDescr_v0_4):
2256 if p.kwargs.mode == "fixed":
2257 mean = p.kwargs.mean
2258 std = p.kwargs.std
2259 assert mean is not None
2260 assert std is not None
2262 axis = _get_complement_v04_axis(tensor_axes, p.kwargs.axes)
2264 if axis is None:
2265 if isinstance(mean, list):
2266 raise ValueError("Expected single float value for mean, not <list>")
2267 if isinstance(std, list):
2268 raise ValueError("Expected single float value for std, not <list>")
2269 return FixedZeroMeanUnitVarianceDescr(
2270 kwargs=FixedZeroMeanUnitVarianceKwargs.model_construct(
2271 mean=mean,
2272 std=std,
2273 )
2274 )
2275 else:
2276 if not isinstance(mean, list):
2277 mean = [float(mean)]
2278 if not isinstance(std, list):
2279 std = [float(std)]
2281 return FixedZeroMeanUnitVarianceDescr(
2282 kwargs=FixedZeroMeanUnitVarianceAlongAxisKwargs(
2283 axis=axis, mean=mean, std=std
2284 )
2285 )
2287 else:
2288 axes = _axes_letters_to_ids(p.kwargs.axes) or []
2289 if p.kwargs.mode == "per_dataset":
2290 axes = [AxisId("batch")] + axes
2291 if not axes:
2292 axes = None
2293 return ZeroMeanUnitVarianceDescr(
2294 kwargs=ZeroMeanUnitVarianceKwargs(axes=axes, eps=p.kwargs.eps)
2295 )
2297 elif isinstance(p, _ScaleRangeDescr_v0_4):
2298 return ScaleRangeDescr(
2299 kwargs=ScaleRangeKwargs(
2300 axes=_axes_letters_to_ids(p.kwargs.axes),
2301 min_percentile=p.kwargs.min_percentile,
2302 max_percentile=p.kwargs.max_percentile,
2303 eps=p.kwargs.eps,
2304 )
2305 )
2306 else:
2307 assert_never(p)
2310class _InputTensorConv(
2311 Converter[
2312 _InputTensorDescr_v0_4,
2313 InputTensorDescr,
2314 FileSource,
2315 Optional[FileSource],
2316 Mapping[_TensorName_v0_4, Mapping[str, int]],
2317 ]
2318):
2319 def _convert(
2320 self,
2321 src: _InputTensorDescr_v0_4,
2322 tgt: type[InputTensorDescr | dict[str, Any]],
2323 test_tensor: FileSource,
2324 sample_tensor: FileSource | None,
2325 size_refs: Mapping[_TensorName_v0_4, Mapping[str, int]],
2326 ) -> InputTensorDescr | dict[str, Any]:
2327 axes: list[InputAxis] = convert_axes( # pyright: ignore[reportAssignmentType]
2328 src.axes,
2329 shape=src.shape,
2330 tensor_type="input",
2331 halo=None,
2332 size_refs=size_refs,
2333 )
2334 prep: list[PreprocessingDescr] = []
2335 for p in src.preprocessing:
2336 cp = _convert_proc(p, src.axes)
2337 assert not isinstance(
2338 cp,
2339 (
2340 CellposeFlowDynamicsDescr,
2341 CustomProcessingDescr,
2342 ScaleMeanVarianceDescr,
2343 StardistPostprocessingDescr,
2344 ),
2345 )
2346 prep.append(cp)
2348 prep.append(EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype="float32")))
2350 return tgt(
2351 axes=axes,
2352 id=TensorId(str(src.name)),
2353 test_tensor=FileDescr(source=test_tensor),
2354 sample_tensor=(
2355 None if sample_tensor is None else FileDescr(source=sample_tensor)
2356 ),
2357 data={"type": src.data_type}, # pyright: ignore[reportArgumentType]
2358 preprocessing=prep,
2359 )
2362_input_tensor_conv = _InputTensorConv(_InputTensorDescr_v0_4, InputTensorDescr)
2365class OutputTensorDescr(TensorDescrBase[OutputAxis]):
2366 id: TensorId = TensorId("output")
2367 """Output tensor id.
2368 No duplicates are allowed across all inputs and outputs."""
2370 postprocessing: list[PostprocessingDescr] = Field(
2371 default_factory=cast(Callable[[], List[PostprocessingDescr]], list)
2372 )
2373 """Description of how this output should be postprocessed.
2375 note: `postprocessing` always ends with an 'ensure_dtype' operation.
2376 If not given this is added to cast to this tensor's `data.type`.
2377 """
2379 @model_validator(mode="after")
2380 def _validate_postprocessing_kwargs(self) -> Self:
2381 axes_ids = [a.id for a in self.axes]
2382 for p in self.postprocessing:
2383 kwargs_axes = p.kwargs.get("axes")
2384 if kwargs_axes is None:
2385 continue
2387 if not isinstance(kwargs_axes, collections.abc.Sequence):
2388 raise ValueError(
2389 f"expected `axes` sequence, but got {type(kwargs_axes)}"
2390 )
2392 kwargs_axes_seq: Sequence[Any] = cast(Sequence[Any], kwargs_axes)
2393 if any(a not in axes_ids for a in kwargs_axes_seq):
2394 raise ValueError("`kwargs.axes` needs to be subset of axes ids")
2396 if isinstance(self.data, (NominalOrOrdinalDataDescr, IntervalOrRatioDataDescr)):
2397 dtype = self.data.type
2398 else:
2399 dtype = self.data[0].type
2401 # ensure `postprocessing` ends with `EnsureDtypeDescr` or `BinarizeDescr`
2402 if not self.postprocessing or not isinstance(
2403 self.postprocessing[-1], (EnsureDtypeDescr, BinarizeDescr)
2404 ):
2405 self.postprocessing.append(
2406 EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype=dtype))
2407 )
2408 return self
2411class _OutputTensorConv(
2412 Converter[
2413 _OutputTensorDescr_v0_4,
2414 OutputTensorDescr,
2415 FileSource,
2416 Optional[FileSource],
2417 Mapping[_TensorName_v0_4, Mapping[str, int]],
2418 ]
2419):
2420 def _convert(
2421 self,
2422 src: _OutputTensorDescr_v0_4,
2423 tgt: type[OutputTensorDescr | dict[str, Any]],
2424 test_tensor: FileSource,
2425 sample_tensor: FileSource | None,
2426 size_refs: Mapping[_TensorName_v0_4, Mapping[str, int]],
2427 ) -> OutputTensorDescr | dict[str, Any]:
2428 # TODO: split convert_axes into convert_output_axes and convert_input_axes
2429 axes: list[OutputAxis] = convert_axes( # pyright: ignore[reportAssignmentType]
2430 src.axes,
2431 shape=src.shape,
2432 tensor_type="output",
2433 halo=src.halo,
2434 size_refs=size_refs,
2435 )
2436 data_descr: dict[str, Any] = {"type": src.data_type}
2437 if data_descr["type"] == "bool":
2438 data_descr["values"] = [False, True]
2440 return tgt(
2441 axes=axes,
2442 id=TensorId(str(src.name)),
2443 test_tensor=FileDescr(source=test_tensor),
2444 sample_tensor=(
2445 None if sample_tensor is None else FileDescr(source=sample_tensor)
2446 ),
2447 data=data_descr, # pyright: ignore[reportArgumentType]
2448 postprocessing=[_convert_proc(p, src.axes) for p in src.postprocessing],
2449 )
2452_output_tensor_conv = _OutputTensorConv(_OutputTensorDescr_v0_4, OutputTensorDescr)
2455TensorDescr: TypeAlias = Union[InputTensorDescr, OutputTensorDescr]
2458def get_halos(
2459 tensors: Mapping[TensorId, TensorDescr],
2460 /,
2461) -> dict[TensorId, dict[AxisId, tuple[int, int]]]:
2462 """Get all input and output halos from tensor descriptions.
2464 Note:
2465 - Input halos are to be padded
2466 - Output halos are to be cropped
2467 """
2468 halos: dict[TensorId, dict[AxisId, tuple[int, int]]] = {}
2469 for descr in tensors.values():
2470 if isinstance(descr, InputTensorDescr):
2471 continue
2472 for axis in descr.axes:
2473 if not isinstance(axis, WithHalo):
2474 continue
2476 ref_scale = next(
2477 a
2478 for a in tensors[axis.size.tensor_id].axes
2479 if a.id == axis.size.axis_id
2480 ).scale
2482 # set output halo (to be cropped)
2483 halos.setdefault(descr.id, {})[axis.id] = (axis.halo, axis.halo)
2484 # set input halo (to be padded)
2485 pad_width = int(axis.halo / axis.scale * ref_scale)
2486 halos.setdefault(axis.size.tensor_id, {})[axis.size.axis_id] = (
2487 pad_width,
2488 pad_width,
2489 )
2491 return halos
2494def validate_tensors(
2495 tensors: Mapping[TensorId, tuple[TensorDescr, NDArray[Any] | None]],
2496 tensor_origin: Literal[
2497 "source", "test_tensor"
2498 ] = "source", # for more precise error messages
2499 *,
2500 pad_inputs: bool | Literal["allow"] = True,
2501 crop_outputs: bool | Literal["allow"] = True,
2502):
2503 """Validate all inputs (and optionally output tensors) against their tensor descriptions.
2505 Args:
2506 tensors: Mapping of tensor id to a tuple of tensor description and optional numpy array.
2507 tensor_origin: String to use in error messages to indicate the origin of the tensors being validated.
2508 pad_inputs: Wether to apply/allow padding of inputs before shape comparison
2509 crop_outputs: Wether to apply/allow cropping of outputs before shape comparison.
2510 """
2511 all_tensor_axes: dict[TensorId, dict[AxisId, tuple[AnyAxis, int | None]]] = {}
2513 def e_msg_location(d: TensorDescr):
2514 return f"{'inputs' if isinstance(d, InputTensorDescr) else 'outputs'}[{d.id}]"
2516 for descr, array in tensors.values():
2517 if array is None:
2518 axis_sizes = {a.id: None for a in descr.axes}
2519 else:
2520 try:
2521 axis_sizes = descr.get_axis_sizes_for_array(array)
2522 except ValueError as e:
2523 raise ValueError(f"{e_msg_location(descr)} {e}")
2525 all_tensor_axes[descr.id] = {a.id: (a, axis_sizes[a.id]) for a in descr.axes}
2527 # get halos to be padded/cropped to validate against halo-adjusted sizes
2528 io_halos = get_halos({k: v[0] for k, v in tensors.items()})
2530 for descr, array in tensors.values():
2531 if array is None:
2532 continue
2534 if descr.dtype in ("float32", "float64"):
2535 invalid_test_tensor_dtype = array.dtype.name not in (
2536 "float32",
2537 "float64",
2538 "uint8",
2539 "int8",
2540 "uint16",
2541 "int16",
2542 "uint32",
2543 "int32",
2544 "uint64",
2545 "int64",
2546 )
2547 else:
2548 invalid_test_tensor_dtype = cast(bool, array.dtype.name != descr.dtype)
2550 if invalid_test_tensor_dtype:
2551 raise ValueError(
2552 f"{tensor_origin} data type '{array.dtype.name}' does not"
2553 + f" match described {e_msg_location(descr)}.dtype '{descr.dtype}'"
2554 )
2556 if array.min() > -1e-4 and array.max() < 1e-4:
2557 raise ValueError(
2558 "Output values are too small for reliable testing."
2559 + f" Values <-1e5 or >=1e5 must be present in {tensor_origin}"
2560 )
2562 for a in descr.axes:
2563 actual_size = all_tensor_axes[descr.id][a.id][1]
2565 if actual_size is None:
2566 continue
2568 if a.size is None:
2569 continue
2571 # add padding width to actual tensor size
2572 total_axis_halo = sum(io_halos.get(descr.id, {}).get(a.id, (0, 0)))
2573 if isinstance(descr, InputTensorDescr):
2574 # pad input halos
2575 actual_size_with_halo = actual_size + total_axis_halo
2576 if pad_inputs is True:
2577 check_sizes = {actual_size_with_halo}
2578 size_hint = " (after padding input halo)"
2579 elif pad_inputs == "allow":
2580 check_sizes = {actual_size, actual_size_with_halo}
2581 size_hint = " (with or without padding input halo)"
2582 elif pad_inputs is False:
2583 check_sizes = {actual_size}
2584 size_hint = ""
2585 else:
2586 assert_never(pad_inputs)
2588 elif isinstance(descr, OutputTensorDescr):
2589 # crop output halos
2590 actual_size_with_halo = max(0, actual_size - total_axis_halo)
2591 if crop_outputs is True:
2592 check_sizes = {actual_size_with_halo}
2593 size_hint = " (after cropping output halo)"
2594 elif crop_outputs == "allow":
2595 check_sizes = {actual_size, actual_size_with_halo}
2596 size_hint = " (with or without cropping output halo)"
2597 elif crop_outputs is False:
2598 check_sizes = {actual_size}
2599 size_hint = ""
2600 else:
2601 assert_never(crop_outputs)
2602 else:
2603 assert_never(descr)
2605 del actual_size # make sure we explicitly use unchanged or halo-adjusted size from here on
2607 if isinstance(a.size, int):
2608 if a.size not in check_sizes:
2609 raise ValueError(
2610 f"{e_msg_location(descr)}.axes[{a.id}]: {tensor_origin} axis "
2611 + f"has incompatible size {check_sizes}{size_hint}, expected {a.size}"
2612 )
2613 elif isinstance(a.size, (ParameterizedSize, DataDependentSize)):
2614 _ = try_all_raise_last(
2615 (partial(a.size.validate_size, s) for s in check_sizes),
2616 f"{e_msg_location(descr)}.axes[{a.id}]: {tensor_origin} axis ",
2617 )
2618 elif isinstance(a.size, SizeReference):
2619 ref_tensor_axes = all_tensor_axes.get(a.size.tensor_id)
2620 if ref_tensor_axes is None:
2621 raise ValueError(
2622 f"{e_msg_location(descr)}.axes[{a.id}].size.tensor_id: Unknown tensor"
2623 + f" reference '{a.size.tensor_id}', available: {list(all_tensor_axes)}"
2624 )
2626 ref_axis, ref_size = ref_tensor_axes.get(a.size.axis_id, (None, None))
2627 if ref_axis is None or ref_size is None:
2628 raise ValueError(
2629 f"{e_msg_location(descr)}.axes[{a.id}].size.axis_id: Unknown tensor axis"
2630 + f" reference '{a.size.tensor_id}.{a.size.axis_id}, available: {list(ref_tensor_axes)}"
2631 )
2633 if a.unit != ref_axis.unit:
2634 raise ValueError(
2635 f"{e_msg_location(descr)}.axes[{a.id}].size: `SizeReference` requires"
2636 + " axis and reference axis to have the same `unit`, but"
2637 + f" {a.unit}!={ref_axis.unit}"
2638 )
2640 if (
2641 expected_size := (
2642 ref_size * ref_axis.scale / a.scale + a.size.offset
2643 )
2644 ) not in check_sizes:
2645 raise ValueError(
2646 f"{e_msg_location(descr)}.{tensor_origin}: axis '{a.id}' of size"
2647 + f" {check_sizes} invalid for referenced size {ref_size};"
2648 + f" expected {expected_size}"
2649 )
2650 else:
2651 assert_never(a.size)
2654FileDescr_dependencies = Annotated[
2655 FileDescr_package,
2656 WithSuffix((".yaml", ".yml"), case_sensitive=True),
2657 Field(examples=[{"source": "environment.yaml"}]),
2658]
2661class _ArchitectureCallableDescr(Node):
2662 callable: Annotated[Identifier, Field(examples=["MyNetworkClass", "get_my_model"])]
2663 """Identifier of the callable that returns a torch.nn.Module instance."""
2665 kwargs: dict[str, YamlValue] = Field(
2666 default_factory=cast(Callable[[], Dict[str, YamlValue]], dict)
2667 )
2668 """key word arguments for the `callable`"""
2671class ArchitectureFromFileDescr(_ArchitectureCallableDescr, FileDescr):
2672 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
2673 """Architecture source file"""
2675 @model_serializer(mode="wrap", when_used="unless-none")
2676 def _serialize(self, nxt: SerializerFunctionWrapHandler, info: SerializationInfo):
2677 return package_file_descr_serializer(self, nxt, info)
2680class ArchitectureFromLibraryDescr(_ArchitectureCallableDescr):
2681 import_from: str
2682 """Where to import the callable from, i.e. `from <import_from> import <callable>`"""
2685class _ArchFileConv(
2686 Converter[
2687 _CallableFromFile_v0_4,
2688 ArchitectureFromFileDescr,
2689 Optional[Sha256],
2690 Dict[str, Any],
2691 ]
2692):
2693 def _convert(
2694 self,
2695 src: _CallableFromFile_v0_4,
2696 tgt: type[ArchitectureFromFileDescr | dict[str, Any]],
2697 sha256: Sha256 | None,
2698 kwargs: dict[str, Any],
2699 ) -> ArchitectureFromFileDescr | dict[str, Any]:
2700 if src.startswith("http") and src.count(":") == 2:
2701 http, source, callable_ = src.split(":")
2702 source = f"{http}:{source}"
2703 elif not src.startswith("http") and src.count(":") == 1:
2704 source, callable_ = src.split(":")
2705 else:
2706 source = str(src)
2707 callable_ = str(src)
2708 return tgt(
2709 callable=Identifier(callable_),
2710 source=cast(FileSource, source),
2711 sha256=sha256,
2712 kwargs=kwargs,
2713 )
2716_arch_file_conv = _ArchFileConv(_CallableFromFile_v0_4, ArchitectureFromFileDescr)
2719class _ArchLibConv(
2720 Converter[
2721 _CallableFromDepencency_v0_4, ArchitectureFromLibraryDescr, Dict[str, Any]
2722 ]
2723):
2724 def _convert(
2725 self,
2726 src: _CallableFromDepencency_v0_4,
2727 tgt: type[ArchitectureFromLibraryDescr | dict[str, Any]],
2728 kwargs: dict[str, Any],
2729 ) -> ArchitectureFromLibraryDescr | dict[str, Any]:
2730 *mods, callable_ = src.split(".")
2731 import_from = ".".join(mods)
2732 return tgt(
2733 import_from=import_from, callable=Identifier(callable_), kwargs=kwargs
2734 )
2737_arch_lib_conv = _ArchLibConv(
2738 _CallableFromDepencency_v0_4, ArchitectureFromLibraryDescr
2739)
2742class WeightsEntryDescrBase(FileDescr):
2743 type: ClassVar[WeightsFormat]
2744 weights_format_name: ClassVar[str] # human readable
2746 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
2747 """Source of the weights file."""
2749 authors: list[Author] | None = None
2750 """Authors
2751 Either the person(s) that have trained this model resulting in the original weights file.
2752 (If this is the initial weights entry, i.e. it does not have a `parent`)
2753 Or the person(s) who have converted the weights to this weights format.
2754 (If this is a child weight, i.e. it has a `parent` field)
2755 """
2757 parent: Annotated[WeightsFormat | None, Field(examples=["pytorch_state_dict"])] = (
2758 None
2759 )
2760 """The source weights these weights were converted from.
2761 For example, if a model's weights were converted from the `pytorch_state_dict` format to `torchscript`,
2762 The `pytorch_state_dict` weights entry has no `parent` and is the parent of the `torchscript` weights.
2763 All weight entries except one (the initial set of weights resulting from training the model),
2764 need to have this field."""
2766 comment: str = ""
2767 """A comment about this weights entry, for example how these weights were created."""
2769 @model_validator(mode="after")
2770 def _validate(self) -> Self:
2771 if self.type == self.parent:
2772 raise ValueError("Weights entry can't be it's own parent.")
2774 return self
2776 @model_serializer(mode="wrap", when_used="unless-none")
2777 def _serialize(self, nxt: SerializerFunctionWrapHandler, info: SerializationInfo):
2778 return package_file_descr_serializer(self, nxt, info)
2781class KerasHdf5WeightsDescr(WeightsEntryDescrBase):
2782 type: ClassVar[WeightsFormat] = "keras_hdf5"
2783 weights_format_name: ClassVar[str] = "Keras HDF5"
2784 tensorflow_version: Version
2785 """TensorFlow version used to create these weights."""
2788class KerasV3WeightsDescr(WeightsEntryDescrBase):
2789 type: ClassVar[WeightsFormat] = "keras_v3"
2790 weights_format_name: ClassVar[str] = "Keras v3"
2791 keras_version: Annotated[Version, Ge(Version(3))]
2792 """Keras version used to create these weights."""
2793 backend: tuple[Literal["tensorflow", "jax", "torch"], Version]
2794 """Keras backend used to create these weights."""
2795 source: Annotated[
2796 FileSource,
2797 AfterValidator(wo_special_file_name),
2798 WithSuffix(".keras", case_sensitive=True),
2799 ]
2800 """Source of the .keras weights file."""
2803FileDescr_external_data = Annotated[
2804 FileDescr_package,
2805 WithSuffix(".data", case_sensitive=True),
2806 Field(examples=[{"source": "weights.onnx.data"}]),
2807]
2810class OnnxWeightsDescr(WeightsEntryDescrBase):
2811 type: ClassVar[WeightsFormat] = "onnx"
2812 weights_format_name: ClassVar[str] = "ONNX"
2813 opset_version: Annotated[int, Ge(7)]
2814 """ONNX opset version"""
2816 external_data: FileDescr_external_data | None = None
2817 """Source of the external ONNX data file holding the weights.
2818 (If present **source** holds the ONNX architecture without weights)."""
2820 @model_validator(mode="after")
2821 def _validate_external_data_unique_file_name(self) -> Self:
2822 if self.external_data is not None and (
2823 extract_file_name(self.source)
2824 == extract_file_name(self.external_data.source)
2825 ):
2826 raise ValueError(
2827 f"ONNX `external_data` file name '{extract_file_name(self.external_data.source)}'"
2828 + " must be different from ONNX `source` file name."
2829 )
2831 return self
2834class PytorchStateDictWeightsDescr(WeightsEntryDescrBase):
2835 type: ClassVar[WeightsFormat] = "pytorch_state_dict"
2836 weights_format_name: ClassVar[str] = "Pytorch State Dict"
2837 architecture: ArchitectureFromFileDescr | ArchitectureFromLibraryDescr
2838 pytorch_version: Version
2839 """Version of the PyTorch library used.
2840 If `architecture.depencencies` is specified it has to include pytorch and any version pinning has to be compatible.
2841 """
2842 dependencies: FileDescr_dependencies | None = None
2843 """Custom depencies beyond pytorch described in a Conda environment file.
2844 Allows to specify custom dependencies, see conda docs:
2845 - [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)
2846 - [Creating an environment file manually](https://conda.io/projects/conda/en/latest/user-guide/tasks/manage-environments.html#creating-an-environment-file-manually)
2848 The conda environment file should include pytorch and any version pinning has to be compatible with
2849 **pytorch_version**.
2850 """
2851 strict: bool = True
2852 """Whether to allow missing or unexpected keys or to be strict about the architecture matching the state dict weights."""
2855class TensorflowJsWeightsDescr(WeightsEntryDescrBase):
2856 type: ClassVar[WeightsFormat] = "tensorflow_js"
2857 weights_format_name: ClassVar[str] = "Tensorflow.js"
2858 tensorflow_version: Version
2859 """Version of the TensorFlow library used."""
2861 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
2862 """The multi-file weights.
2863 All required files/folders should be a zip archive."""
2866class TensorflowSavedModelBundleWeightsDescr(WeightsEntryDescrBase):
2867 type: ClassVar[WeightsFormat] = "tensorflow_saved_model_bundle"
2868 weights_format_name: ClassVar[str] = "Tensorflow Saved Model"
2869 tensorflow_version: Version
2870 """Version of the TensorFlow library used."""
2872 dependencies: FileDescr_dependencies | None = None
2873 """Custom dependencies beyond tensorflow.
2874 Should include tensorflow and any version pinning has to be compatible with **tensorflow_version**."""
2876 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
2877 """The multi-file weights.
2878 All required files/folders should be a zip archive."""
2881class TorchscriptWeightsDescr(WeightsEntryDescrBase):
2882 type: ClassVar[WeightsFormat] = "torchscript"
2883 weights_format_name: ClassVar[str] = "TorchScript"
2884 pytorch_version: Version
2885 """Version of the PyTorch library used."""
2888SpecificWeightsDescr = Union[
2889 KerasHdf5WeightsDescr,
2890 KerasV3WeightsDescr,
2891 OnnxWeightsDescr,
2892 PytorchStateDictWeightsDescr,
2893 TensorflowJsWeightsDescr,
2894 TensorflowSavedModelBundleWeightsDescr,
2895 TorchscriptWeightsDescr,
2896]
2899class WeightsDescr(Node):
2900 keras_hdf5: KerasHdf5WeightsDescr | None = None
2901 keras_v3: KerasV3WeightsDescr | None = None
2902 onnx: OnnxWeightsDescr | None = None
2903 pytorch_state_dict: PytorchStateDictWeightsDescr | None = None
2904 tensorflow_js: TensorflowJsWeightsDescr | None = None
2905 tensorflow_saved_model_bundle: TensorflowSavedModelBundleWeightsDescr | None = None
2906 torchscript: TorchscriptWeightsDescr | None = None
2908 @model_validator(mode="after")
2909 def check_entries(self) -> Self:
2910 entries = {wtype for wtype, entry in self if entry is not None}
2912 if not entries:
2913 raise ValueError("Missing weights entry")
2915 entries_wo_parent = {
2916 wtype
2917 for wtype, entry in self
2918 if entry is not None and hasattr(entry, "parent") and entry.parent is None
2919 }
2920 if len(entries_wo_parent) != 1:
2921 issue_warning(
2922 "Exactly one weights entry may not specify the `parent` field (got"
2923 + " {value}). That entry is considered the original set of model weights."
2924 + " Other weight formats are created through conversion of the orignal or"
2925 + " already converted weights. They have to reference the weights format"
2926 + " they were converted from as their `parent`.",
2927 value=len(entries_wo_parent),
2928 field="weights",
2929 )
2931 for wtype, entry in self:
2932 if entry is None:
2933 continue
2935 assert hasattr(entry, "type")
2936 assert hasattr(entry, "parent")
2937 assert wtype == entry.type
2938 if (
2939 entry.parent is not None and entry.parent not in entries
2940 ): # self reference checked for `parent` field
2941 raise ValueError(
2942 f"`weights.{wtype}.parent={entry.parent} not in specified weight"
2943 + f" formats: {entries}"
2944 )
2946 return self
2948 def __getitem__(
2949 self,
2950 key: WeightsFormat,
2951 ):
2952 if key == "keras_hdf5":
2953 ret = self.keras_hdf5
2954 elif key == "keras_v3":
2955 ret = self.keras_v3
2956 elif key == "onnx":
2957 ret = self.onnx
2958 elif key == "pytorch_state_dict":
2959 ret = self.pytorch_state_dict
2960 elif key == "tensorflow_js":
2961 ret = self.tensorflow_js
2962 elif key == "tensorflow_saved_model_bundle":
2963 ret = self.tensorflow_saved_model_bundle
2964 elif key == "torchscript":
2965 ret = self.torchscript
2966 else:
2967 raise KeyError(key)
2969 if ret is None:
2970 raise KeyError(key)
2972 return ret
2974 @overload
2975 def __setitem__(
2976 self, key: Literal["keras_hdf5"], value: KerasHdf5WeightsDescr | None
2977 ) -> None: ...
2978 @overload
2979 def __setitem__(
2980 self, key: Literal["keras_v3"], value: KerasV3WeightsDescr | None
2981 ) -> None: ...
2982 @overload
2983 def __setitem__(
2984 self, key: Literal["onnx"], value: OnnxWeightsDescr | None
2985 ) -> None: ...
2986 @overload
2987 def __setitem__(
2988 self,
2989 key: Literal["pytorch_state_dict"],
2990 value: PytorchStateDictWeightsDescr | None,
2991 ) -> None: ...
2992 @overload
2993 def __setitem__(
2994 self, key: Literal["tensorflow_js"], value: TensorflowJsWeightsDescr | None
2995 ) -> None: ...
2996 @overload
2997 def __setitem__(
2998 self,
2999 key: Literal["tensorflow_saved_model_bundle"],
3000 value: TensorflowSavedModelBundleWeightsDescr | None,
3001 ) -> None: ...
3002 @overload
3003 def __setitem__(
3004 self, key: Literal["torchscript"], value: TorchscriptWeightsDescr | None
3005 ) -> None: ...
3007 def __setitem__(
3008 self,
3009 key: WeightsFormat,
3010 value: SpecificWeightsDescr | None,
3011 ):
3012 if key == "keras_hdf5":
3013 if value is not None and not isinstance(value, KerasHdf5WeightsDescr):
3014 raise TypeError(
3015 f"Expected KerasHdf5WeightsDescr or None for key 'keras_hdf5', got {type(value)}"
3016 )
3017 self.keras_hdf5 = value
3018 elif key == "keras_v3":
3019 if value is not None and not isinstance(value, KerasV3WeightsDescr):
3020 raise TypeError(
3021 f"Expected KerasV3WeightsDescr or None for key 'keras_v3', got {type(value)}"
3022 )
3023 self.keras_v3 = value
3024 elif key == "onnx":
3025 if value is not None and not isinstance(value, OnnxWeightsDescr):
3026 raise TypeError(
3027 f"Expected OnnxWeightsDescr or None for key 'onnx', got {type(value)}"
3028 )
3029 self.onnx = value
3030 elif key == "pytorch_state_dict":
3031 if value is not None and not isinstance(
3032 value, PytorchStateDictWeightsDescr
3033 ):
3034 raise TypeError(
3035 f"Expected PytorchStateDictWeightsDescr or None for key 'pytorch_state_dict', got {type(value)}"
3036 )
3037 self.pytorch_state_dict = value
3038 elif key == "tensorflow_js":
3039 if value is not None and not isinstance(value, TensorflowJsWeightsDescr):
3040 raise TypeError(
3041 f"Expected TensorflowJsWeightsDescr or None for key 'tensorflow_js', got {type(value)}"
3042 )
3043 self.tensorflow_js = value
3044 elif key == "tensorflow_saved_model_bundle":
3045 if value is not None and not isinstance(
3046 value, TensorflowSavedModelBundleWeightsDescr
3047 ):
3048 raise TypeError(
3049 f"Expected TensorflowSavedModelBundleWeightsDescr or None for key 'tensorflow_saved_model_bundle', got {type(value)}"
3050 )
3051 self.tensorflow_saved_model_bundle = value
3052 elif key == "torchscript":
3053 if value is not None and not isinstance(value, TorchscriptWeightsDescr):
3054 raise TypeError(
3055 f"Expected TorchscriptWeightsDescr or None for key 'torchscript', got {type(value)}"
3056 )
3057 self.torchscript = value
3058 else:
3059 raise KeyError(key)
3061 @property
3062 def available_formats(self) -> dict[WeightsFormat, SpecificWeightsDescr]:
3063 return {
3064 **({} if self.keras_hdf5 is None else {"keras_hdf5": self.keras_hdf5}),
3065 **({} if self.keras_v3 is None else {"keras_v3": self.keras_v3}),
3066 **({} if self.onnx is None else {"onnx": self.onnx}),
3067 **(
3068 {}
3069 if self.pytorch_state_dict is None
3070 else {"pytorch_state_dict": self.pytorch_state_dict}
3071 ),
3072 **(
3073 {}
3074 if self.tensorflow_js is None
3075 else {"tensorflow_js": self.tensorflow_js}
3076 ),
3077 **(
3078 {}
3079 if self.tensorflow_saved_model_bundle is None
3080 else {
3081 "tensorflow_saved_model_bundle": self.tensorflow_saved_model_bundle
3082 }
3083 ),
3084 **({} if self.torchscript is None else {"torchscript": self.torchscript}),
3085 }
3087 @property
3088 def missing_formats(self) -> set[WeightsFormat]:
3089 return {
3090 wf for wf in get_args(WeightsFormat) if wf not in self.available_formats
3091 }
3094class LinkedModel(LinkedResourceBase):
3095 """Reference to a bioimage.io model."""
3097 id: ModelId
3098 """A valid model `id` from the bioimage.io collection."""
3101class _DataDepSize(NamedTuple):
3102 min: StrictInt
3103 max: StrictInt | None
3106class _AxisSizes(NamedTuple):
3107 """the lenghts of all axes of model inputs and outputs"""
3109 inputs: dict[tuple[TensorId, AxisId], int]
3110 outputs: dict[tuple[TensorId, AxisId], int | _DataDepSize]
3113class _TensorSizes(NamedTuple):
3114 """_AxisSizes as nested dicts"""
3116 inputs: dict[TensorId, dict[AxisId, int]]
3117 outputs: dict[TensorId, dict[AxisId, int | _DataDepSize]]
3120_RT = TypeVar("_RT", RelativeTolerance, Literal[0])
3121_AT = TypeVar("_AT", AbsoluteTolerance, Literal[0])
3124class _ReproducibilityToleranceBase(Node, ABC, Generic[_RT, _AT], extra="allow"):
3125 """Describes what small numerical differences -- if any -- may be tolerated
3126 in the generated output when executing in different environments.
3128 A tensor element *output* is considered mismatched to the **test_tensor** if
3129 abs(*output* - **test_tensor**) > **absolute_tolerance** + **relative_tolerance** * abs(**test_tensor**).
3130 (Internally we call [numpy.testing.assert_allclose](https://numpy.org/doc/stable/reference/generated/numpy.testing.assert_allclose.html).)
3132 Motivation:
3133 For testing we can request the respective deep learning frameworks to be as
3134 reproducible as possible by setting seeds and chosing deterministic algorithms,
3135 but differences in operating systems, available hardware and installed drivers
3136 may still lead to numerical differences.
3137 """
3139 mismatched_elements_per_million: MismatchedElementsPerMillion = 100
3140 """Maximum number of mismatched elements/pixels per million to tolerate."""
3142 weights_formats: Sequence[WeightsFormat] = ()
3143 """Limits the weights formats these details apply to."""
3145 relative_tolerance: _RT
3146 absolute_tolerance: _AT
3148 @model_serializer(mode="wrap")
3149 def serialize(
3150 self,
3151 handler: SerializerFunctionWrapHandler,
3152 ) -> dict[str, Any]:
3153 """explicitly serialize tolerance values
3155 (even if they are default values and `exclude_defaults=True`).
3156 """
3157 data = handler(self)
3158 data["relative_tolerance"] = self.relative_tolerance
3159 data["absolute_tolerance"] = self.absolute_tolerance
3160 data["mismatched_elements_per_million"] = self.mismatched_elements_per_million
3161 return data
3164class ReproducibilityTolerance(
3165 _ReproducibilityToleranceBase[RelativeTolerance, AbsoluteTolerance]
3166):
3167 output_ids: Sequence[TensorId] = ()
3168 """Limits the output tensor IDs these reproducibility details apply to.
3170 If empty, the reproducibility details apply to all outputs."""
3172 relative_tolerance: RelativeTolerance = 1e-3
3173 """Maximum relative tolerance of reproduced test tensor.
3174 Needs to be 0 for discrete (integer and boolean) outputs.
3175 """
3177 absolute_tolerance: AbsoluteTolerance = 1e-3
3178 """Maximum absolute tolerance of reproduced test tensor.
3180 The `absolute_tolerance` may not be greater than 1% of the maximum absolute value of the test tensor
3181 (this serves to prevent false positives in the presence of very small numbers).
3182 Needs to be 0 for discrete (integer and boolean) outputs.
3183 """
3186class DiscreteReproducibilityTolerance(
3187 _ReproducibilityToleranceBase[Literal[0], Literal[0]]
3188):
3189 output_ids: Literal["any_discrete_output"] = "any_discrete_output"
3190 """Limits the output tensor IDs to all outputs of integer or boolean data type."""
3192 relative_tolerance: Literal[0] = 0
3193 """Relative tolerance does not apply to discrete outputs"""
3195 absolute_tolerance: Literal[0] = 0
3196 """Absolute tolerance does not apply to discrete outputs"""
3198 mismatched_elements_per_million: MismatchedElementsPerMillion = 1000
3199 """Maximum number of mismatched elements/pixels per million to tolerate.
3201 Note: Increased default mismatched elements per million compared to
3202 [`ReproducibilityTolerance`][ReproducibilityTolerance]/[`ContinuousReproducibilityTolerance`][ContinuousReproducibilityTolerance],
3203 since a single pixel mismatch can be more significant than for continuous outputs with numeric tolerance."""
3206class ContinuousReproducibilityTolerance(
3207 _ReproducibilityToleranceBase[RelativeTolerance, AbsoluteTolerance]
3208):
3209 output_ids: Literal["any_continuous_output"] = "any_continuous_output"
3210 """Limits the output tensor IDs to all outputs of float data type."""
3212 relative_tolerance: RelativeTolerance = 1e-3
3213 """Maximum relative tolerance of reproduced test tensor.
3214 Needs to be 0 for discrete (integer and boolean) outputs.
3215 """
3217 absolute_tolerance: AbsoluteTolerance = 1e-3
3218 """Maximum absolute tolerance of reproduced test tensor.
3220 The `absolute_tolerance` may not be greater than 1% of the maximum absolute value of the test tensor
3221 (this serves to prevent false positives in the presence of very small numbers).
3222 Needs to be 0 for discrete (integer and boolean) outputs.
3223 """
3226class BiasRisksLimitations(Node, extra="allow"):
3227 """Known biases, risks, technical limitations, and recommendations for model use."""
3229 known_biases: str = dedent("""\
3230 In general bioimage models may suffer from biases caused by:
3232 - Imaging protocol dependencies
3233 - Use of a specific cell type
3234 - Species-specific training data limitations
3236 """)
3237 """Biases in training data or model behavior."""
3239 risks: str = dedent("""\
3240 Common risks in bioimage analysis include:
3242 - Erroneously assuming generalization to unseen experimental conditions
3243 - Trusting (overconfident) model outputs without validation
3244 - Misinterpretation of results
3246 """)
3247 """Potential risks in the context of bioimage analysis."""
3249 limitations: str | None = None
3250 """Technical limitations and failure modes."""
3252 recommendations: str = "Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model."
3253 """Mitigation strategies regarding `known_biases`, `risks`, and `limitations`, as well as applicable best practices.
3255 Consider:
3256 - How to use a validation dataset?
3257 - How to manually validate?
3258 - Feasibility of domain adaptation for different experimental setups?
3260 """
3262 def format_md(self) -> str:
3263 if self.limitations is None:
3264 limitations_header = ""
3265 else:
3266 limitations_header = "## Limitations\n\n"
3268 return f"""# Bias, Risks, and Limitations
3270{self.known_biases}
3272{self.risks}
3274{limitations_header}{self.limitations or ""}
3276## Recommendations
3278{self.recommendations}
3280"""
3283class TrainingDetails(Node, extra="allow"):
3284 training_preprocessing: str | None = None
3285 """Detailed image preprocessing steps during model training:
3287 Mention:
3288 - *Normalization methods*
3289 - *Augmentation strategies*
3290 - *Resizing/resampling procedures*
3291 - *Artifact handling*
3293 """
3295 training_epochs: float | None = None
3296 """Number of training epochs."""
3298 training_batch_size: float | None = None
3299 """Batch size used in training."""
3301 initial_learning_rate: float | None = None
3302 """Initial learning rate used in training."""
3304 learning_rate_schedule: str | None = None
3305 """Learning rate schedule used in training."""
3307 loss_function: str | None = None
3308 """Loss function used in training, e.g. nn.MSELoss."""
3310 loss_function_kwargs: dict[str, YamlValue] = Field(
3311 default_factory=cast(Callable[[], Dict[str, YamlValue]], dict)
3312 )
3313 """key word arguments for the `loss_function`"""
3315 optimizer: str | None = None
3316 """optimizer, e.g. torch.optim.Adam"""
3318 optimizer_kwargs: dict[str, YamlValue] = Field(
3319 default_factory=cast(Callable[[], Dict[str, YamlValue]], dict)
3320 )
3321 """key word arguments for the `optimizer`"""
3323 regularization: str | None = None
3324 """Regularization techniques used during training, e.g. drop-out or weight decay."""
3326 training_duration: float | None = None
3327 """Total training duration in hours."""
3330class Evaluation(Node, extra="allow"):
3331 model_id: ModelId | None = None
3332 """Model being evaluated."""
3334 dataset_id: DatasetId
3335 """Dataset used for evaluation."""
3337 dataset_source: HttpUrl
3338 """Source of the dataset."""
3340 dataset_role: Literal["train", "validation", "test", "independent", "unknown"]
3341 """Role of the dataset used for evaluation.
3343 - `train`: dataset was (part of) the training data
3344 - `validation`: dataset was (part of) the validation data used during training, e.g. used for model selection or hyperparameter tuning
3345 - `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
3346 - `independent`: dataset is entirely independent test data; not used during training or validation, and acquired from a different source/distribution than training data
3347 - `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.
3348 """
3350 sample_count: int
3351 """Number of evaluated samples."""
3353 evaluation_factors: list[Annotated[str, MaxLen(16)]]
3354 """(Abbreviations of) each evaluation factor.
3356 Evaluation factors are criteria along which model performance is evaluated, e.g. different image conditions
3357 like 'low SNR', 'high cell density', or different biological conditions like 'cell type A', 'cell type B'.
3358 An 'overall' factor may be included to summarize performance across all conditions.
3359 """
3361 evaluation_factors_long: list[str]
3362 """Descriptions (long form) of each evaluation factor."""
3364 metrics: list[Annotated[str, MaxLen(16)]]
3365 """(Abbreviations of) metrics used for evaluation."""
3367 metrics_long: list[str]
3368 """Description of each metric used."""
3370 @model_validator(mode="after")
3371 def _validate_list_lengths(self) -> Self:
3372 if len(self.evaluation_factors) != len(self.evaluation_factors_long):
3373 raise ValueError(
3374 "`evaluation_factors` and `evaluation_factors_long` must have the same length"
3375 )
3377 if len(self.metrics) != len(self.metrics_long):
3378 raise ValueError("`metrics` and `metrics_long` must have the same length")
3380 if len(self.results) != len(self.metrics):
3381 raise ValueError("`results` must have the same number of rows as `metrics`")
3383 for row in self.results:
3384 if len(row) != len(self.evaluation_factors):
3385 raise ValueError(
3386 "`results` must have the same number of columns (in every row) as `evaluation_factors`"
3387 )
3389 return self
3391 results: list[list[str | float | int]]
3392 """Results for each metric (rows; outer list) and each evaluation factor (columns; inner list)."""
3394 results_summary: str | None = None
3395 """Interpretation of results for general audience.
3397 Consider:
3398 - Overall model performance
3399 - Comparison to existing methods
3400 - Limitations and areas for improvement
3402"""
3404 def format_md(self):
3405 results_header = ["Metric"] + self.evaluation_factors
3406 results_table_cells = [results_header, ["---"] * len(results_header)] + [
3407 [metric] + [str(r) for r in row]
3408 for metric, row in zip(self.metrics, self.results)
3409 ]
3411 results_table = "".join(
3412 "| " + " | ".join(row) + " |\n" for row in results_table_cells
3413 )
3414 factors = "".join(
3415 f"\n - {ef}: {efl}"
3416 for ef, efl in zip(self.evaluation_factors, self.evaluation_factors_long)
3417 )
3418 metrics = "".join(
3419 f"\n - {em}: {eml}" for em, eml in zip(self.metrics, self.metrics_long)
3420 )
3422 return f"""## Testing Data, Factors & Metrics
3424Evaluation of {self.model_id or "this"} model on the {self.dataset_id} dataset (dataset role: {self.dataset_role}).
3426### Testing Data
3428- **Source:** [{self.dataset_id}]({self.dataset_source})
3429- **Size:** {self.sample_count} evaluated samples
3431### Factors
3432{factors}
3434### Metrics
3435{metrics}
3437## Results
3439### Quantitative Results
3441{results_table}
3443### Summary
3445{self.results_summary or "missing"}
3447"""
3450class EnvironmentalImpact(Node, extra="allow"):
3451 """Environmental considerations for model training and deployment.
3453 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).
3454 """
3456 hardware_type: str | None = None
3457 """GPU/CPU specifications"""
3459 hours_used: float | None = None
3460 """Total compute hours"""
3462 cloud_provider: str | None = None
3463 """If applicable"""
3465 compute_region: str | None = None
3466 """Geographic location"""
3468 co2_emitted: float | None = None
3469 """kg CO2 equivalent
3471 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).
3472 """
3474 def format_md(self):
3475 """Filled Markdown template section following [Hugging Face Model Card Template](https://huggingface.co/docs/hub/en/model-card-annotated)."""
3476 if self == self.__class__():
3477 return ""
3479 ret = "# Environmental Impact\n\n"
3480 if self.hardware_type is not None:
3481 ret += f"- **Hardware Type:** {self.hardware_type}\n"
3482 if self.hours_used is not None:
3483 ret += f"- **Hours used:** {self.hours_used}\n"
3484 if self.cloud_provider is not None:
3485 ret += f"- **Cloud Provider:** {self.cloud_provider}\n"
3486 if self.compute_region is not None:
3487 ret += f"- **Compute Region:** {self.compute_region}\n"
3488 if self.co2_emitted is not None:
3489 ret += f"- **Carbon Emitted:** {self.co2_emitted} kg CO2e\n"
3491 return ret + "\n"
3494class BioimageioConfig(Node, extra="allow"):
3495 reproducibility_tolerance: Sequence[
3496 ReproducibilityTolerance
3497 | ContinuousReproducibilityTolerance
3498 | DiscreteReproducibilityTolerance
3499 ] = Field(
3500 default_factory=lambda: (
3501 ContinuousReproducibilityTolerance.model_construct(),
3502 DiscreteReproducibilityTolerance.model_construct(),
3503 )
3504 )
3505 """Tolerances to allow when reproducing the model's test outputs
3506 from the model's test inputs.
3507 Only the first entry matching tensor id and weights format is considered.
3508 For defaults see [`ContinuousReproducibilityTolerance`][ContinuousReproducibilityTolerance]/[`DiscreteReproducibilityTolerance`][DiscreteReproducibilityTolerance]
3509 """
3511 @model_serializer(mode="wrap")
3512 def _serialize(
3513 self,
3514 handler: SerializerFunctionWrapHandler,
3515 ) -> dict[str, Any]:
3516 """explicitly serialize reproducibility_tolerance
3518 (even if they are default values and `exclude_defaults=True`).
3519 """
3520 data = handler(self)
3521 data["reproducibility_tolerance"] = self.reproducibility_tolerance
3522 return data
3524 funded_by: str | None = None
3525 """Funding agency, grant number if applicable"""
3527 architecture_type: Annotated[str, MaxLen(32)] | None = (
3528 None # TODO: add to differentiated tags
3529 )
3530 """Model architecture type, e.g., 3D U-Net, ResNet, transformer"""
3532 architecture_description: str | None = None
3533 """Text description of model architecture."""
3535 modality: str | None = None # TODO: add to differentiated tags
3536 """Input modality, e.g., fluorescence microscopy, electron microscopy"""
3538 target_structure: list[str] = Field( # TODO: add to differentiated tags
3539 default_factory=cast(Callable[[], List[str]], list)
3540 )
3541 """Biological structure(s) the model is designed to analyze, e.g., nuclei, mitochondria, cells"""
3543 task: str | None = None # TODO: add to differentiated tags
3544 """Bioimage-specific task type, e.g., segmentation, classification, detection, denoising"""
3546 new_version: ModelId | None = None
3547 """A new version of this model exists with a different model id."""
3549 out_of_scope_use: str | None = None
3550 """Describe how the model may be misused in bioimage analysis contexts and what users should **not** do with the model."""
3552 bias_risks_limitations: BiasRisksLimitations = Field(
3553 default_factory=BiasRisksLimitations.model_construct
3554 )
3555 """Description of known bias, risks, and technical limitations for in-scope model use."""
3557 model_parameter_count: int | None = None
3558 """Total number of model parameters."""
3560 training: TrainingDetails = Field(default_factory=TrainingDetails.model_construct)
3561 """Details on how the model was trained."""
3563 inference_time: str | None = None
3564 """Average inference time per image/tile. Specify hardware and image size. Multiple examples can be given."""
3566 memory_requirements_inference: str | None = None
3567 """GPU memory needed for inference. Multiple examples with different image size can be given."""
3569 memory_requirements_training: str | None = None
3570 """GPU memory needed for training. Multiple examples with different image/batch sizes can be given."""
3572 evaluations: list[Evaluation] = Field(
3573 default_factory=cast(Callable[[], List[Evaluation]], list)
3574 )
3575 """Quantitative model evaluations.
3577 Note:
3578 At the moment we recommend to include only a single test dataset
3579 (with evaluation factors that may mark subsets of the dataset)
3580 to avoid confusion and make the presentation of results cleaner.
3581 """
3583 environmental_impact: EnvironmentalImpact = Field(
3584 default_factory=EnvironmentalImpact.model_construct
3585 )
3586 """Environmental considerations for model training and deployment"""
3589class Config(Node, extra="allow"):
3590 bioimageio: BioimageioConfig = Field(
3591 default_factory=BioimageioConfig.model_construct
3592 )
3593 stardist: YamlValue = None
3596class ModelDescr(GenericModelDescrBase):
3597 """Specification of the fields used in a bioimage.io-compliant RDF to describe AI models with pretrained weights.
3598 These fields are typically stored in a YAML file which we call a model resource description file (model RDF).
3599 """
3601 implemented_format_version: ClassVar[Literal["0.5.14"]] = "0.5.14"
3602 if TYPE_CHECKING:
3603 format_version: Literal["0.5.14"] = "0.5.14"
3604 else:
3605 format_version: Literal["0.5.14"]
3606 """Version of the bioimage.io model description specification used.
3607 When creating a new model always use the latest micro/patch version described here.
3608 The `format_version` is important for any consumer software to understand how to parse the fields.
3609 """
3611 implemented_type: ClassVar[Literal["model"]] = "model"
3612 if TYPE_CHECKING:
3613 type: Literal["model"] = "model"
3614 else:
3615 type: Literal["model"]
3616 """Specialized resource type 'model'"""
3618 id: ModelId | None = None
3619 """bioimage.io-wide unique resource identifier
3620 assigned by bioimage.io; version **un**specific."""
3622 authors: FAIR[list[Author]] = Field(
3623 default_factory=cast(Callable[[], List[Author]], list)
3624 )
3625 """The authors are the creators of the model RDF and the primary points of contact."""
3627 documentation: FAIR[FileDescr_documentation | None] = None
3628 """Additional model documentation.
3629 The recommended documentation source file name is `README.md`. An `.md` suffix is mandatory.
3630 The documentation should include a '#[#] Validation' (sub)section
3631 with details on how to quantitatively validate the model on unseen data."""
3633 @field_validator("documentation", mode="after")
3634 @classmethod
3635 def _validate_documentation(cls, value: FileDescr | None) -> FileDescr | None:
3636 if not get_validation_context().perform_io_checks or value is None:
3637 return value
3639 doc_reader = get_reader(value)
3640 doc_content = doc_reader.read().decode(encoding="utf-8")
3641 if not re.search("#.*[vV]alidation", doc_content):
3642 issue_warning(
3643 "No '# Validation' (sub)section found in {value}.",
3644 value=value,
3645 field="documentation",
3646 )
3648 return value
3650 inputs: NotEmpty[Sequence[InputTensorDescr]]
3651 """Describes the input tensors expected by this model."""
3653 @field_validator("inputs", mode="after")
3654 @classmethod
3655 def _validate_input_axes(
3656 cls, inputs: Sequence[InputTensorDescr]
3657 ) -> Sequence[InputTensorDescr]:
3658 input_size_refs = cls._get_axes_with_independent_size(inputs)
3660 for i, ipt in enumerate(inputs):
3661 valid_independent_refs: dict[
3662 tuple[TensorId, AxisId],
3663 tuple[TensorDescr, AnyAxis, int | ParameterizedSize],
3664 ] = {
3665 **{
3666 (ipt.id, a.id): (ipt, a, a.size)
3667 for a in ipt.axes
3668 if not isinstance(a, BatchAxis)
3669 and isinstance(a.size, (int, ParameterizedSize))
3670 },
3671 **input_size_refs,
3672 }
3673 for a, ax in enumerate(ipt.axes):
3674 cls._validate_axis(
3675 "inputs",
3676 i=i,
3677 tensor_id=ipt.id,
3678 a=a,
3679 axis=ax,
3680 valid_independent_refs=valid_independent_refs,
3681 )
3682 return inputs
3684 @staticmethod
3685 def _validate_axis(
3686 field_name: str,
3687 i: int,
3688 tensor_id: TensorId,
3689 a: int,
3690 axis: AnyAxis,
3691 valid_independent_refs: dict[
3692 tuple[TensorId, AxisId],
3693 tuple[TensorDescr, AnyAxis, int | ParameterizedSize],
3694 ],
3695 ):
3696 if isinstance(axis, BatchAxis) or isinstance(
3697 axis.size, (int, ParameterizedSize, DataDependentSize)
3698 ):
3699 return
3700 elif not isinstance(axis.size, SizeReference):
3701 assert_never(axis.size)
3703 # validate axis.size SizeReference
3704 ref = (axis.size.tensor_id, axis.size.axis_id)
3705 if ref not in valid_independent_refs:
3706 raise ValueError(
3707 "Invalid tensor axis reference at"
3708 + f" {field_name}[{i}].axes[{a}].size: {axis.size}."
3709 )
3710 if ref == (tensor_id, axis.id):
3711 raise ValueError(
3712 "Self-referencing not allowed for"
3713 + f" {field_name}[{i}].axes[{a}].size: {axis.size}"
3714 )
3715 if axis.type == "channel":
3716 if valid_independent_refs[ref][1].type != "channel":
3717 raise ValueError(
3718 "A channel axis' size may only reference another fixed size"
3719 + " channel axis."
3720 )
3721 if isinstance(axis.channel_names, str) and "{i}" in axis.channel_names:
3722 ref_size = valid_independent_refs[ref][2]
3723 assert isinstance(ref_size, int), (
3724 "channel axis ref (another channel axis) has to specify fixed"
3725 + " size"
3726 )
3727 generated_channel_names = [
3728 axis.channel_names.format(i=i) for i in range(1, ref_size + 1)
3729 ]
3730 axis.channel_names = generated_channel_names
3732 if (ax_unit := getattr(axis, "unit", None)) != (
3733 ref_unit := getattr(valid_independent_refs[ref][1], "unit", None)
3734 ):
3735 raise ValueError(
3736 "The units of an axis and its reference axis need to match, but"
3737 + f" '{ax_unit}' != '{ref_unit}'."
3738 )
3739 ref_axis = valid_independent_refs[ref][1]
3740 if isinstance(ref_axis, BatchAxis):
3741 raise ValueError(
3742 f"Invalid reference axis '{ref_axis.id}' for {tensor_id}.{axis.id}"
3743 + " (a batch axis is not allowed as reference)."
3744 )
3746 if isinstance(axis, WithHalo):
3747 min_size = axis.size.get_size(axis, ref_axis, n=0)
3748 if (min_size - 2 * axis.halo) < 1:
3749 raise ValueError(
3750 f"axis {axis.id} with minimum size {min_size} is too small for halo"
3751 + f" {axis.halo}."
3752 )
3754 ref_halo = axis.halo * axis.scale / ref_axis.scale
3755 if ref_halo != int(ref_halo):
3756 raise ValueError(
3757 f"Inferred halo for {'.'.join(ref)} is not an integer ({ref_halo} ="
3758 + f" {tensor_id}.{axis.id}.halo {axis.halo}"
3759 + f" * {tensor_id}.{axis.id}.scale {axis.scale}"
3760 + f" / {'.'.join(ref)}.scale {ref_axis.scale})."
3761 )
3763 def validate_input_tensors(
3764 self,
3765 sources: Sequence[NDArray[Any]] | Mapping[TensorId, NDArray[Any] | None],
3766 *,
3767 pad_inputs: bool | Literal["allow"] = True,
3768 crop_outputs: bool | Literal["allow"] = True,
3769 ) -> Mapping[TensorId, NDArray[Any] | None]:
3770 """Check if the given input tensors match the model's input tensor descriptions.
3771 This includes checks of tensor shapes and dtypes, but not of the actual values.
3772 """
3773 if not isinstance(sources, collections.abc.Mapping):
3774 sources = {descr.id: tensor for descr, tensor in zip(self.inputs, sources)}
3776 tensors = {
3777 **{descr.id: (descr, sources.get(descr.id)) for descr in self.inputs},
3778 **{ # outputs are required for halo
3779 descr.id: (descr, None) for descr in self.outputs
3780 },
3781 }
3782 validate_tensors(tensors, pad_inputs=pad_inputs, crop_outputs=crop_outputs)
3784 return sources
3786 @model_validator(mode="after")
3787 def _validate_test_tensors(self) -> Self:
3788 if not get_validation_context().perform_io_checks:
3789 return self
3791 test_inputs = {
3792 descr.id: (
3793 descr,
3794 None if descr.test_tensor is None else load_array(descr.test_tensor),
3795 )
3796 for descr in self.inputs
3797 }
3798 test_outputs = {
3799 descr.id: (
3800 descr,
3801 None if descr.test_tensor is None else load_array(descr.test_tensor),
3802 )
3803 for descr in self.outputs
3804 }
3806 validate_tensors(
3807 {**test_inputs, **test_outputs},
3808 tensor_origin="test_tensor",
3809 pad_inputs="allow",
3810 crop_outputs="allow",
3811 )
3813 seen_output_ids: set[TensorId] = (
3814 set()
3815 ) # reproducibility tolerance is only applicable to the first matching output tensor id
3816 for rep_tol in self.config.bioimageio.reproducibility_tolerance:
3817 if not rep_tol.output_ids:
3818 out_arrays: dict[TensorId, NDArray[Any] | None] = {
3819 k: v[1] for k, v in test_outputs.items() if k not in seen_output_ids
3820 }
3821 else:
3822 out_arrays = {}
3823 for k, v in test_outputs.items():
3824 if k in seen_output_ids:
3825 continue
3827 if rep_tol.output_ids == "any_discrete_output":
3828 if not (
3829 v[0].dtype.startswith("int")
3830 or v[0].dtype.startswith("uint")
3831 or v[0].dtype.startswith("bool")
3832 ):
3833 continue
3834 elif rep_tol.output_ids == "any_continuous_output":
3835 if not v[0].dtype.startswith("float"):
3836 continue
3837 elif k not in rep_tol.output_ids:
3838 continue
3840 out_arrays[k] = v[1]
3842 for out_id, array in out_arrays.items():
3843 seen_output_ids.add(out_id)
3844 if array is None:
3845 continue
3847 if rep_tol.absolute_tolerance > (max_test_value := array.max()) * 0.01:
3848 raise ValueError(
3849 "config.bioimageio.reproducibility_tolerance.absolute_tolerance="
3850 + f"{rep_tol.absolute_tolerance} > 0.01*{max_test_value}"
3851 + f" (1% of the maximum value of the test tensor '{out_id}')"
3852 )
3854 return self
3856 @model_validator(mode="after")
3857 def _validate_tensor_references_in_proc_kwargs(self, info: ValidationInfo) -> Self:
3858 ipt_refs = {t.id for t in self.inputs}
3859 missing_refs = [
3860 k["reference_tensor"]
3861 for k in [p.kwargs for ipt in self.inputs for p in ipt.preprocessing]
3862 + [p.kwargs for out in self.outputs for p in out.postprocessing]
3863 if "reference_tensor" in k
3864 and k["reference_tensor"] is not None
3865 and k["reference_tensor"] not in ipt_refs
3866 ]
3868 if missing_refs:
3869 raise ValueError(
3870 f"`reference_tensor`s {missing_refs} not found. Valid input tensor"
3871 + f" references are: {ipt_refs}."
3872 )
3874 return self
3876 name: Annotated[
3877 str,
3878 RestrictCharacters(string.ascii_letters + string.digits + "_+- ()"),
3879 MinLen(5),
3880 MaxLen(128),
3881 warn(Annotated[str, MaxLen(64)], "Name longer than 64 characters.", INFO),
3882 ]
3883 """A human-readable name of this model.
3884 It should be no longer than 64 characters
3885 and may only contain letter, number, underscore, minus, parentheses and spaces.
3886 We recommend to chose a name that refers to the model's task and image modality.
3887 """
3889 outputs: NotEmpty[Sequence[OutputTensorDescr]]
3890 """Describes the output tensors."""
3892 @field_validator("outputs", mode="after")
3893 @classmethod
3894 def _validate_tensor_ids(
3895 cls, outputs: Sequence[OutputTensorDescr], info: ValidationInfo
3896 ) -> Sequence[OutputTensorDescr]:
3897 tensor_ids = [
3898 t.id for t in info.data.get("inputs", []) + info.data.get("outputs", [])
3899 ]
3900 duplicate_tensor_ids: list[str] = []
3901 seen: set[str] = set()
3902 for t in tensor_ids:
3903 if t in seen:
3904 duplicate_tensor_ids.append(t)
3906 seen.add(t)
3908 if duplicate_tensor_ids:
3909 raise ValueError(f"Duplicate tensor ids: {duplicate_tensor_ids}")
3911 return outputs
3913 @staticmethod
3914 def _get_axes_with_parameterized_size(
3915 io: Sequence[InputTensorDescr] | Sequence[OutputTensorDescr],
3916 ):
3917 return {
3918 f"{t.id}.{a.id}": (t, a, a.size)
3919 for t in io
3920 for a in t.axes
3921 if not isinstance(a, BatchAxis) and isinstance(a.size, ParameterizedSize)
3922 }
3924 @staticmethod
3925 def _get_axes_with_independent_size(
3926 io: Sequence[InputTensorDescr] | Sequence[OutputTensorDescr],
3927 ):
3928 return {
3929 (t.id, a.id): (t, a, a.size)
3930 for t in io
3931 for a in t.axes
3932 if not isinstance(a, BatchAxis)
3933 and isinstance(a.size, (int, ParameterizedSize))
3934 }
3936 @field_validator("outputs", mode="after")
3937 @classmethod
3938 def _validate_output_axes(
3939 cls, outputs: list[OutputTensorDescr], info: ValidationInfo
3940 ) -> list[OutputTensorDescr]:
3941 input_size_refs = cls._get_axes_with_independent_size(
3942 info.data.get("inputs", [])
3943 )
3944 output_size_refs = cls._get_axes_with_independent_size(outputs)
3946 for i, out in enumerate(outputs):
3947 valid_independent_refs: dict[
3948 tuple[TensorId, AxisId],
3949 tuple[TensorDescr, AnyAxis, int | ParameterizedSize],
3950 ] = {
3951 **{
3952 (out.id, a.id): (out, a, a.size)
3953 for a in out.axes
3954 if not isinstance(a, BatchAxis)
3955 and isinstance(a.size, (int, ParameterizedSize))
3956 },
3957 **input_size_refs,
3958 **output_size_refs,
3959 }
3960 for a, ax in enumerate(out.axes):
3961 cls._validate_axis(
3962 "outputs",
3963 i,
3964 out.id,
3965 a,
3966 ax,
3967 valid_independent_refs=valid_independent_refs,
3968 )
3970 return outputs
3972 packaged_by: list[Author] = Field(
3973 default_factory=cast(Callable[[], List[Author]], list)
3974 )
3975 """The persons that have packaged and uploaded this model.
3976 Only required if those persons differ from the `authors`."""
3978 parent: LinkedModel | None = None
3979 """The model from which this model is derived, e.g. by fine-tuning the weights."""
3981 @model_validator(mode="after")
3982 def _validate_parent_is_not_self(self) -> Self:
3983 if self.parent is not None and self.parent.id == self.id:
3984 raise ValueError("A model description may not reference itself as parent.")
3986 return self
3988 run_mode: Annotated[
3989 RunMode | None,
3990 warn(None, "Run mode '{value}' has limited support across consumer softwares."),
3991 ] = None
3992 """Custom run mode for this model: for more complex prediction procedures like test time
3993 data augmentation that currently cannot be expressed in the specification.
3994 No standard run modes are defined yet."""
3996 timestamp: Datetime = Field(default_factory=Datetime.now)
3997 """Timestamp in [ISO 8601](#https://en.wikipedia.org/wiki/ISO_8601) format
3998 with a few restrictions listed [here](https://docs.python.org/3/library/datetime.html#datetime.datetime.fromisoformat).
3999 (In Python a datetime object is valid, too)."""
4001 training_data: Annotated[
4002 None | LinkedDataset | DatasetDescr | DatasetDescr02,
4003 Field(union_mode="left_to_right"),
4004 ] = None
4005 """The dataset used to train this model"""
4007 weights: Annotated[WeightsDescr, WrapSerializer(package_weights)]
4008 """The weights for this model.
4009 Weights can be given for different formats, but should otherwise be equivalent.
4010 The available weight formats determine which consumers can use this model."""
4012 config: Config = Field(default_factory=Config.model_construct)
4014 @model_validator(mode="after")
4015 def _validate_reproducibility_tolerance_output_ids(self) -> Self:
4016 seen_output_ids: set[TensorId] = set()
4017 for rep_tol in self.config.bioimageio.reproducibility_tolerance:
4018 if not rep_tol.output_ids:
4019 continue
4021 if rep_tol.output_ids == "any_discrete_output":
4022 applicable_output_ids = [
4023 t.id
4024 for t in self.outputs
4025 if t.dtype.startswith("int")
4026 or t.dtype.startswith("uint")
4027 or t.dtype.startswith("bool")
4028 ]
4029 elif rep_tol.output_ids == "any_continuous_output":
4030 applicable_output_ids = [
4031 t.id for t in self.outputs if t.dtype.startswith("float")
4032 ]
4033 else:
4034 applicable_output_ids = rep_tol.output_ids
4036 for out_id in applicable_output_ids:
4037 if out_id in seen_output_ids:
4038 continue
4039 else:
4040 seen_output_ids.add(out_id)
4042 discrete_values = next(
4043 (
4044 t.dtype.startswith("int")
4045 or t.dtype.startswith("uint")
4046 or t.dtype.startswith("bool")
4047 for t in self.outputs
4048 if t.id == out_id
4049 ),
4050 None,
4051 )
4052 if discrete_values is None:
4053 raise ValueError(
4054 f"config.bioimageio.reproducibility_tolerance.output_ids contains '{out_id}', which is not a valid output tensor id."
4055 )
4056 if discrete_values and rep_tol.relative_tolerance != 0:
4057 raise ValueError(
4058 f"config.bioimageio.reproducibility_tolerance.output_ids contains '{out_id}', which is a discrete output tensor, but config.bioimageio.reproducibility_tolerance.relative_tolerance={rep_tol.relative_tolerance} is not 0."
4059 )
4060 if discrete_values and rep_tol.absolute_tolerance != 0:
4061 raise ValueError(
4062 f"config.bioimageio.reproducibility_tolerance.output_ids contains '{out_id}', which is a discrete output tensor, but config.bioimageio.reproducibility_tolerance.absolute_tolerance={rep_tol.absolute_tolerance} is not 0."
4063 )
4065 return self
4067 @model_validator(mode="after")
4068 def _add_default_cover(self) -> Self:
4069 if not get_validation_context().perform_io_checks or self.covers:
4070 return self
4072 try:
4073 generated_covers = generate_covers(
4074 [
4075 (t, load_array(t.test_tensor))
4076 for t in self.inputs
4077 if t.test_tensor is not None
4078 ],
4079 [
4080 (t, load_array(t.test_tensor))
4081 for t in self.outputs
4082 if t.test_tensor is not None
4083 ],
4084 )
4085 except Exception as e:
4086 issue_warning(
4087 "Failed to generate cover image(s): {e}",
4088 value=self.covers,
4089 msg_context={"e": e},
4090 field="covers",
4091 )
4092 else:
4093 self.covers.extend(generated_covers)
4095 return self
4097 def get_input_test_arrays(self) -> list[NDArray[Any]]:
4098 return self._get_test_arrays(self.inputs)
4100 def get_output_test_arrays(self) -> list[NDArray[Any]]:
4101 return self._get_test_arrays(self.outputs)
4103 @staticmethod
4104 def _get_test_arrays(
4105 io_descr: Sequence[InputTensorDescr] | Sequence[OutputTensorDescr],
4106 ):
4107 ts: list[FileDescr] = []
4108 for d in io_descr:
4109 if d.test_tensor is None:
4110 raise ValueError(
4111 f"Failed to get test arrays: description of '{d.id}' is missing a `test_tensor`."
4112 )
4113 ts.append(d.test_tensor)
4115 data = [load_array(t) for t in ts]
4116 assert all(isinstance(d, np.ndarray) for d in data)
4117 return data
4119 @staticmethod
4120 def get_batch_size(tensor_sizes: Mapping[TensorId, Mapping[AxisId, int]]) -> int:
4121 batch_size = 1
4122 tensor_with_batchsize: TensorId | None = None
4123 for tid in tensor_sizes:
4124 for aid, s in tensor_sizes[tid].items():
4125 if aid != BATCH_AXIS_ID or s == 1 or s == batch_size:
4126 continue
4128 if batch_size != 1:
4129 assert tensor_with_batchsize is not None
4130 raise ValueError(
4131 f"batch size mismatch for tensors '{tensor_with_batchsize}' ({batch_size}) and '{tid}' ({s})"
4132 )
4134 batch_size = s
4135 tensor_with_batchsize = tid
4137 return batch_size
4139 def get_output_tensor_sizes(
4140 self, input_sizes: Mapping[TensorId, Mapping[AxisId, int]]
4141 ) -> dict[TensorId, dict[AxisId, int]]:
4142 """Returns the tensor output sizes for given **input_sizes**.
4143 Data-dependent output sizes are set to -1.
4145 Note:
4146 A warning is logged if the calculated output sizes do not match the output tensor sizes
4147 inferred from the model descriptions with parameters inferred fromthe given input sizes.
4148 This is likely the case if the input sizes are not valid according to the model description.
4149 """
4150 # get tensor sizes for valid inputs that can be parameterized by batch size and axis size parameters `n` (for each parameterized axis)
4151 batch_size = self.get_batch_size(input_sizes)
4152 ns = self.get_ns(input_sizes)
4153 valid_tensor_sizes = self.get_tensor_sizes(ns, batch_size=batch_size)
4155 input_sizes_refs = {
4156 ipt.id: {
4157 ax.id: (ax, input_sizes[ipt.id][ax.id])
4158 for ax in ipt.axes
4159 if not isinstance(ax, BatchAxis) and ax.id in input_sizes[ipt.id]
4160 }
4161 for ipt in self.inputs
4162 if ipt.id in input_sizes
4163 }
4164 output_sizes: dict[TensorId, dict[AxisId, int]] = {
4165 out.id: {} for out in self.outputs
4166 }
4167 for out in self.outputs:
4168 for ax in out.axes:
4169 if isinstance(ax, BatchAxis):
4170 s = batch_size
4171 elif isinstance(ax.size, int):
4172 s = ax.size
4173 elif isinstance(ax.size, SizeReference):
4174 if (
4175 ax.size.tensor_id not in input_sizes_refs
4176 or ax.size.axis_id not in input_sizes_refs[ax.size.tensor_id]
4177 ):
4178 logger.warning(
4179 "Output axis {}.{} references input axis {}.{}, which is not present in the provided input sizes."
4180 + " Assuming the output is optional and not produced for this input.",
4181 out.id,
4182 ax.id,
4183 ax.size.tensor_id,
4184 ax.size.axis_id,
4185 )
4186 continue
4188 ref_ax, ref_size = input_sizes_refs[ax.size.tensor_id][
4189 ax.size.axis_id
4190 ]
4191 s = ax.size.get_size(ax, ref_ax, ref_size=ref_size)
4192 elif isinstance(ax.size, DataDependentSize):
4193 s = -1
4194 else:
4195 assert_never(ax.size)
4197 output_sizes[out.id][ax.id] = s
4199 if valid_tensor_sizes.outputs != output_sizes:
4200 if valid_tensor_sizes.inputs != input_sizes:
4201 logger.trace(
4202 "Input sizes do not match the model's input tensor descriptions which may lead to inconsistent output sizes."
4203 + " Expected: {}, got: {}.",
4204 valid_tensor_sizes.inputs,
4205 input_sizes,
4206 )
4207 else:
4208 logger.error(
4209 "Directly calculated output sizes do not match the output tensor sizes inferred from the model descriptions with the given input size parameters."
4210 + " Expected: {}, got: {}.",
4211 valid_tensor_sizes.outputs,
4212 output_sizes,
4213 )
4215 return output_sizes
4217 def get_ns(self, input_sizes: Mapping[TensorId, Mapping[AxisId, int]]):
4218 """get parameter `n` for each parameterized axis
4219 such that the valid input size is >= the given input size"""
4220 ret: dict[tuple[TensorId, AxisId], ParameterizedSize_N] = {}
4221 axes_per_member = {t.id: {a.id: a for a in t.axes} for t in self.inputs}
4222 for tid in input_sizes:
4223 axes = axes_per_member.get(tid)
4224 if axes is None:
4225 continue
4226 for aid, s in input_sizes[tid].items():
4227 axis_descr = axes.get(aid)
4228 if axis_descr is None:
4229 continue
4230 if isinstance(axis_descr.size, ParameterizedSize):
4231 ret[(tid, aid)] = axis_descr.size.get_n(s)
4232 elif axis_descr.size is None or isinstance(
4233 axis_descr.size, (int, SizeReference)
4234 ):
4235 pass
4236 else:
4237 assert_never(axis_descr.size)
4239 return ret
4241 def get_tensor_sizes(
4242 self,
4243 ns: Mapping[tuple[TensorId, AxisId], ParameterizedSize_N] | ParameterizedSize_N,
4244 batch_size: int,
4245 max_input_shape: Mapping[TensorId, Mapping[AxisId, int]] | None = None,
4246 ) -> _TensorSizes:
4247 max_axis_sizes: dict[tuple[TensorId, AxisId], int] = {}
4248 for m, this_max_axis_sizes in (max_input_shape or {}).items():
4249 for a, s in this_max_axis_sizes.items():
4250 max_axis_sizes[(m, a)] = s
4252 axis_sizes = self.get_axis_sizes(
4253 ns, batch_size=batch_size, max_input_shape=max_axis_sizes
4254 )
4255 return _TensorSizes(
4256 {
4257 t: {
4258 aa: axis_sizes.inputs[(tt, aa)]
4259 for tt, aa in axis_sizes.inputs
4260 if tt == t
4261 }
4262 for t in {tt for tt, _ in axis_sizes.inputs}
4263 },
4264 {
4265 t: {
4266 aa: axis_sizes.outputs[(tt, aa)]
4267 for tt, aa in axis_sizes.outputs
4268 if tt == t
4269 }
4270 for t in {tt for tt, _ in axis_sizes.outputs}
4271 },
4272 )
4274 def get_axis_sizes(
4275 self,
4276 ns: Mapping[tuple[TensorId, AxisId], ParameterizedSize_N] | ParameterizedSize_N,
4277 batch_size: int | None = None,
4278 *,
4279 max_input_shape: Mapping[tuple[TensorId, AxisId], int] | None = None,
4280 ) -> _AxisSizes:
4281 """Determine input and output block shape for scale factors **ns**
4282 of parameterized input sizes.
4284 Args:
4285 ns: Scale factor `n` for each axis (keyed by (tensor_id, axis_id))
4286 that is parameterized as `size = min + n * step`.
4287 batch_size: The desired size of the batch dimension.
4288 If given **batch_size** overwrites any batch size present in
4289 **max_input_shape**. Default 1.
4290 max_input_shape: Limits the derived block shapes.
4291 Each axis for which the input size, parameterized by `n`, is larger
4292 than **max_input_shape** is set to the minimal value `n_min` for which
4293 this is still true.
4294 Use this for small input samples or large values of **ns**.
4295 Or simply whenever you know the full input shape.
4297 Returns:
4298 Resolved axis sizes for model inputs and outputs.
4299 """
4300 max_input_shape = max_input_shape or {}
4301 if batch_size is None:
4302 for (_t_id, a_id), s in max_input_shape.items():
4303 if a_id == BATCH_AXIS_ID:
4304 batch_size = s
4305 break
4306 else:
4307 batch_size = 1
4309 all_axes = {
4310 t.id: {a.id: a for a in t.axes} for t in chain(self.inputs, self.outputs)
4311 }
4313 inputs: dict[tuple[TensorId, AxisId], int] = {}
4314 outputs: dict[tuple[TensorId, AxisId], int | _DataDepSize] = {}
4316 def get_axis_size(a: InputAxis | OutputAxis):
4317 if isinstance(a, BatchAxis):
4318 if not isinstance(ns, int) and (t_descr.id, a.id) in ns:
4319 logger.warning(
4320 "Ignoring unexpected size increment factor (n) for batch axis"
4321 + " of tensor '{}'.",
4322 t_descr.id,
4323 )
4324 return batch_size
4325 elif isinstance(a.size, int):
4326 if not isinstance(ns, int) and (t_descr.id, a.id) in ns:
4327 logger.warning(
4328 "Ignoring unexpected size increment factor (n) for fixed size"
4329 + " axis '{}' of tensor '{}'.",
4330 a.id,
4331 t_descr.id,
4332 )
4333 return a.size
4334 elif isinstance(a.size, ParameterizedSize):
4335 n = ns if isinstance(ns, int) else ns.get((t_descr.id, a.id))
4336 if n is None:
4337 raise ValueError(
4338 "Size increment factor (n) missing for parametrized axis"
4339 + f" '{a.id}' of tensor '{t_descr.id}'."
4340 )
4341 s_max = max_input_shape.get((t_descr.id, a.id))
4342 if s_max is not None:
4343 n = min(n, a.size.get_n(s_max))
4345 return a.size.get_size(n)
4347 elif isinstance(a.size, SizeReference):
4348 if not isinstance(ns, int) and (t_descr.id, a.id) in ns:
4349 logger.warning(
4350 "Ignoring unexpected size increment factor (n) for axis '{}'"
4351 + " of tensor '{}' with size reference.",
4352 a.id,
4353 t_descr.id,
4354 )
4355 assert not isinstance(a, BatchAxis)
4356 ref_axis = all_axes[a.size.tensor_id][a.size.axis_id]
4357 assert not isinstance(ref_axis, BatchAxis)
4358 ref_key = (a.size.tensor_id, a.size.axis_id)
4359 ref_size = inputs.get(ref_key, outputs.get(ref_key))
4360 assert ref_size is not None, ref_key
4361 assert not isinstance(ref_size, _DataDepSize), ref_key
4362 return a.size.get_size(
4363 axis=a,
4364 ref_axis=ref_axis,
4365 ref_size=ref_size,
4366 )
4367 elif isinstance(a.size, DataDependentSize):
4368 if not isinstance(ns, int) and (t_descr.id, a.id) in ns:
4369 logger.warning(
4370 "Ignoring unexpected increment factor (n) for data dependent"
4371 + " size axis '{}' of tensor '{}'.",
4372 a.id,
4373 t_descr.id,
4374 )
4375 return _DataDepSize(a.size.min, a.size.max)
4376 else:
4377 assert_never(a.size)
4379 # first resolve all , but the `SizeReference` input sizes
4380 for t_descr in self.inputs:
4381 for a in t_descr.axes:
4382 if not isinstance(a.size, SizeReference):
4383 s = get_axis_size(a)
4384 assert not isinstance(s, _DataDepSize)
4385 inputs[t_descr.id, a.id] = s
4387 # resolve all other input axis sizes
4388 for t_descr in self.inputs:
4389 for a in t_descr.axes:
4390 if isinstance(a.size, SizeReference):
4391 s = get_axis_size(a)
4392 assert not isinstance(s, _DataDepSize)
4393 inputs[t_descr.id, a.id] = s
4395 # resolve all output axis sizes
4396 for t_descr in self.outputs:
4397 for a in t_descr.axes:
4398 assert not isinstance(a.size, ParameterizedSize)
4399 s = get_axis_size(a)
4400 outputs[t_descr.id, a.id] = s
4402 return _AxisSizes(inputs=inputs, outputs=outputs)
4404 @model_validator(mode="before")
4405 @classmethod
4406 def _convert(cls, data: dict[str, Any]) -> dict[str, Any]:
4407 cls.convert_from_old_format_wo_validation(data)
4408 return data
4410 @classmethod
4411 def convert_from_old_format_wo_validation(cls, data: dict[str, Any]) -> None:
4412 """Convert metadata following an older format version to this classes' format
4413 without validating the result.
4414 """
4415 if (
4416 data.get("type") == "model"
4417 and isinstance(fv := data.get("format_version"), str)
4418 and fv.count(".") == 2
4419 ):
4420 fv_parts = fv.split(".")
4421 if any(not p.isdigit() for p in fv_parts):
4422 return
4424 fv_tuple = tuple(map(int, fv_parts))
4426 assert cls.implemented_format_version_tuple[0:2] == (0, 5)
4427 if fv_tuple[:2] in ((0, 3), (0, 4)):
4428 m04 = _ModelDescr_v0_4.load(data)
4429 if isinstance(m04, InvalidDescr):
4430 try:
4431 updated = _model_conv.convert_as_dict(
4432 m04 # pyright: ignore[reportArgumentType]
4433 )
4434 except Exception as e:
4435 logger.error(
4436 "Failed to convert from invalid model 0.4 description."
4437 + f"\nerror: {e}"
4438 + "\nProceeding with model 0.5 validation without conversion."
4439 )
4440 updated = None
4441 else:
4442 updated = _model_conv.convert_as_dict(m04)
4444 if updated is not None:
4445 data.clear()
4446 data.update(updated)
4448 elif fv_tuple[:2] == (0, 5):
4449 # bump patch version
4450 data["format_version"] = cls.implemented_format_version
4452 if fv_tuple[:2] in ((0, 3), (0, 4)) or (
4453 fv_tuple[:2] == (0, 5) and fv_tuple[2] < 11
4454 ):
4455 convert_plain_covers_and_docs_and_icon(data)
4458class _ModelConv(Converter[_ModelDescr_v0_4, ModelDescr]):
4459 def _convert(
4460 self, src: _ModelDescr_v0_4, tgt: type[ModelDescr | dict[str, Any]]
4461 ) -> ModelDescr | dict[str, Any]:
4462 name = "".join(
4463 c if c in string.ascii_letters + string.digits + "_+- ()" else " "
4464 for c in src.name
4465 )
4467 def conv_authors(auths: Sequence[_Author_v0_4] | None):
4468 conv = (
4469 _author_conv.convert if TYPE_CHECKING else _author_conv.convert_as_dict
4470 )
4471 return None if auths is None else [conv(a) for a in auths]
4473 if TYPE_CHECKING:
4474 arch_file_conv = _arch_file_conv.convert
4475 arch_lib_conv = _arch_lib_conv.convert
4476 else:
4477 arch_file_conv = _arch_file_conv.convert_as_dict
4478 arch_lib_conv = _arch_lib_conv.convert_as_dict
4480 input_size_refs = {
4481 ipt.name: {
4482 a: s
4483 for a, s in zip(
4484 ipt.axes,
4485 (
4486 ipt.shape.min
4487 if isinstance(ipt.shape, _ParameterizedInputShape_v0_4)
4488 else ipt.shape
4489 ),
4490 )
4491 }
4492 for ipt in src.inputs
4493 if ipt.shape
4494 }
4495 output_size_refs = {
4496 **{
4497 out.name: {a: s for a, s in zip(out.axes, out.shape)}
4498 for out in src.outputs
4499 if not isinstance(out.shape, _ImplicitOutputShape_v0_4)
4500 },
4501 **input_size_refs,
4502 }
4504 return tgt(
4505 attachments=(
4506 []
4507 if src.attachments is None
4508 else [FileDescr(source=f) for f in src.attachments.files]
4509 ),
4510 authors=[_author_conv.convert_as_dict(a) for a in src.authors], # pyright: ignore[reportArgumentType]
4511 cite=[{"text": c.text, "doi": c.doi, "url": c.url} for c in src.cite], # pyright: ignore[reportArgumentType]
4512 config=src.config, # pyright: ignore[reportArgumentType]
4513 covers=[{"source": c} for c in src.covers], # pyright: ignore[reportArgumentType]
4514 description=src.description,
4515 documentation={"source": src.documentation} if src.documentation else None, # pyright: ignore[reportArgumentType]
4516 format_version="0.5.14",
4517 git_repo=src.git_repo, # pyright: ignore[reportArgumentType]
4518 icon={"source": src.icon} if src.icon else None, # pyright: ignore[reportArgumentType]
4519 id=None if src.id is None else ModelId(src.id),
4520 id_emoji=src.id_emoji,
4521 license=src.license, # type: ignore
4522 links=src.links,
4523 maintainers=[_maintainer_conv.convert_as_dict(m) for m in src.maintainers], # pyright: ignore[reportArgumentType]
4524 name=name,
4525 tags=src.tags,
4526 type=src.type,
4527 uploader=src.uploader,
4528 version=src.version,
4529 inputs=[ # pyright: ignore[reportArgumentType]
4530 _input_tensor_conv.convert_as_dict(ipt, tt, st, input_size_refs)
4531 for ipt, tt, st in zip(
4532 src.inputs,
4533 src.test_inputs,
4534 src.sample_inputs or [None] * len(src.test_inputs),
4535 )
4536 ],
4537 outputs=[ # pyright: ignore[reportArgumentType]
4538 _output_tensor_conv.convert_as_dict(out, tt, st, output_size_refs)
4539 for out, tt, st in zip(
4540 src.outputs,
4541 src.test_outputs,
4542 src.sample_outputs or [None] * len(src.test_outputs),
4543 )
4544 ],
4545 parent=(
4546 None
4547 if src.parent is None
4548 else LinkedModel(
4549 id=ModelId(
4550 str(src.parent.id)
4551 + (
4552 ""
4553 if src.parent.version_number is None
4554 else f"/{src.parent.version_number}"
4555 )
4556 )
4557 )
4558 ),
4559 training_data=(
4560 None
4561 if src.training_data is None
4562 else (
4563 LinkedDataset(
4564 id=DatasetId(
4565 str(src.training_data.id)
4566 + (
4567 ""
4568 if src.training_data.version_number is None
4569 else f"/{src.training_data.version_number}"
4570 )
4571 )
4572 )
4573 if isinstance(src.training_data, LinkedDataset02)
4574 else src.training_data
4575 )
4576 ),
4577 packaged_by=[_author_conv.convert_as_dict(a) for a in src.packaged_by], # pyright: ignore[reportArgumentType]
4578 run_mode=src.run_mode,
4579 timestamp=src.timestamp,
4580 weights=(WeightsDescr if TYPE_CHECKING else dict)(
4581 keras_hdf5=(w := src.weights.keras_hdf5)
4582 and (KerasHdf5WeightsDescr if TYPE_CHECKING else dict)(
4583 authors=conv_authors(w.authors),
4584 source=w.source,
4585 tensorflow_version=w.tensorflow_version or Version("1.15"),
4586 parent=w.parent,
4587 ),
4588 onnx=(w := src.weights.onnx)
4589 and (OnnxWeightsDescr if TYPE_CHECKING else dict)(
4590 source=w.source,
4591 authors=conv_authors(w.authors),
4592 parent=w.parent,
4593 opset_version=w.opset_version or 15,
4594 ),
4595 pytorch_state_dict=(w := src.weights.pytorch_state_dict)
4596 and (PytorchStateDictWeightsDescr if TYPE_CHECKING else dict)(
4597 source=w.source,
4598 authors=conv_authors(w.authors),
4599 parent=w.parent,
4600 architecture=(
4601 arch_file_conv(
4602 w.architecture,
4603 w.architecture_sha256,
4604 w.kwargs,
4605 )
4606 if isinstance(w.architecture, _CallableFromFile_v0_4)
4607 else arch_lib_conv(w.architecture, w.kwargs)
4608 ),
4609 pytorch_version=w.pytorch_version or Version("1.10"),
4610 dependencies=(
4611 None
4612 if w.dependencies is None
4613 else (FileDescr if TYPE_CHECKING else dict)(
4614 source=cast(
4615 FileSource,
4616 str(deps := w.dependencies)[
4617 (
4618 len("conda:")
4619 if str(deps).startswith("conda:")
4620 else 0
4621 ) :
4622 ],
4623 )
4624 )
4625 ),
4626 ),
4627 tensorflow_js=(w := src.weights.tensorflow_js)
4628 and (TensorflowJsWeightsDescr if TYPE_CHECKING else dict)(
4629 source=w.source,
4630 authors=conv_authors(w.authors),
4631 parent=w.parent,
4632 tensorflow_version=w.tensorflow_version or Version("1.15"),
4633 ),
4634 tensorflow_saved_model_bundle=(
4635 w := src.weights.tensorflow_saved_model_bundle
4636 )
4637 and (TensorflowSavedModelBundleWeightsDescr if TYPE_CHECKING else dict)(
4638 authors=conv_authors(w.authors),
4639 parent=w.parent,
4640 source=w.source,
4641 tensorflow_version=w.tensorflow_version or Version("1.15"),
4642 dependencies=(
4643 None
4644 if w.dependencies is None
4645 else (FileDescr if TYPE_CHECKING else dict)(
4646 source=cast(
4647 FileSource,
4648 (
4649 str(w.dependencies)[len("conda:") :]
4650 if str(w.dependencies).startswith("conda:")
4651 else str(w.dependencies)
4652 ),
4653 )
4654 )
4655 ),
4656 ),
4657 torchscript=(w := src.weights.torchscript)
4658 and (TorchscriptWeightsDescr if TYPE_CHECKING else dict)(
4659 source=w.source,
4660 authors=conv_authors(w.authors),
4661 parent=w.parent,
4662 pytorch_version=w.pytorch_version or Version("1.10"),
4663 ),
4664 ),
4665 )
4668_model_conv = _ModelConv(_ModelDescr_v0_4, ModelDescr)
4671# create better cover images for 3d data and non-image outputs
4672def generate_covers(
4673 inputs: Sequence[tuple[InputTensorDescr, NDArray[Any]]],
4674 outputs: Sequence[tuple[OutputTensorDescr, NDArray[Any]]],
4675) -> list[FileDescr]:
4676 def squeeze(
4677 data: NDArray[Any], axes: Sequence[AnyAxis]
4678 ) -> tuple[NDArray[Any], list[AnyAxis]]:
4679 """apply numpy.ndarray.squeeze while keeping track of the axis descriptions remaining"""
4680 if data.ndim != len(axes):
4681 raise ValueError(
4682 f"tensor shape {data.shape} does not match described axes"
4683 + f" {[a.id for a in axes]}"
4684 )
4686 axes = [deepcopy(a) for a, s in zip(axes, data.shape) if s != 1]
4687 return data.squeeze(), axes
4689 def normalize(
4690 data: NDArray[Any], axis: tuple[int, ...] | None, eps: float = 1e-7
4691 ) -> NDArray[np.float32]:
4692 data = data.astype("float32")
4693 data -= data.min(axis=axis, keepdims=True)
4694 data /= data.max(axis=axis, keepdims=True) + eps
4695 return data
4697 def to_2d_image(data: NDArray[Any], axes: Sequence[AnyAxis]):
4698 original_shape = data.shape
4699 original_axes = list(axes)
4700 data, axes = squeeze(data, axes)
4702 # take slice fom any batch or index axis if needed
4703 # and convert the first channel axis and take a slice from any additional channel axes
4704 slices: tuple[slice, ...] = ()
4705 ndim = data.ndim
4706 ndim_need = 3 if any(isinstance(a, ChannelAxis) for a in axes) else 2
4707 has_c_axis = False
4708 for i, a in enumerate(axes):
4709 s = data.shape[i]
4710 assert s > 1
4711 if (
4712 isinstance(a, (BatchAxis, IndexInputAxis, IndexOutputAxis))
4713 and ndim > ndim_need
4714 ):
4715 data = data[slices + (slice(s // 2 - 1, s // 2),)]
4716 ndim -= 1
4717 elif isinstance(a, ChannelAxis):
4718 if has_c_axis:
4719 # second channel axis
4720 data = data[slices + (slice(0, 1),)]
4721 ndim -= 1
4722 else:
4723 has_c_axis = True
4724 if s == 2:
4725 # visualize two channels with cyan and magenta
4726 data = np.concatenate(
4727 [
4728 data[slices + (slice(1, 2),)],
4729 data[slices + (slice(0, 1),)],
4730 (
4731 data[slices + (slice(0, 1),)]
4732 + data[slices + (slice(1, 2),)]
4733 )
4734 / 2, # TODO: take maximum instead?
4735 ],
4736 axis=i,
4737 )
4738 elif data.shape[i] == 3:
4739 pass # visualize 3 channels as RGB
4740 else:
4741 # visualize first 3 channels as RGB
4742 data = data[slices + (slice(3),)]
4744 assert data.shape[i] == 3
4746 slices += (slice(None),)
4748 data, axes = squeeze(data, axes)
4749 assert len(axes) == ndim
4750 # take slice from z axis if needed
4751 slices = ()
4752 if ndim > ndim_need:
4753 for i, a in enumerate(axes):
4754 s = data.shape[i]
4755 if a.id == AxisId("z"):
4756 data = data[slices + (slice(s // 2 - 1, s // 2),)]
4757 data, axes = squeeze(data, axes)
4758 ndim -= 1
4759 break
4761 slices += (slice(None),)
4763 # take slice from any space or time axis
4764 slices = ()
4766 for i, a in enumerate(axes):
4767 if ndim <= ndim_need:
4768 break
4770 s = data.shape[i]
4771 assert s > 1
4772 if isinstance(
4773 a, (SpaceInputAxis, SpaceOutputAxis, TimeInputAxis, TimeOutputAxis)
4774 ):
4775 data = data[slices + (slice(s // 2 - 1, s // 2),)]
4776 ndim -= 1
4778 slices += (slice(None),)
4780 del slices
4781 data, axes = squeeze(data, axes)
4782 assert len(axes) == ndim
4784 if (has_c_axis and ndim != 3) or (not has_c_axis and ndim != 2):
4785 raise ValueError(
4786 f"Failed to construct cover image from shape {original_shape} with axes {[a.id for a in original_axes]}."
4787 )
4789 if not has_c_axis:
4790 assert ndim == 2
4791 data = np.repeat(data[:, :, None], 3, axis=2)
4792 axes.append(ChannelAxis(channel_names=list("RGB")))
4793 ndim += 1
4795 assert ndim == 3
4797 # transpose axis order such that longest axis comes first...
4798 axis_order: list[int] = [int(i) for i in np.argsort(list(data.shape))]
4799 axis_order.reverse()
4800 # ... and channel axis is last
4801 c = next(i for i in range(3) if isinstance(axes[i], ChannelAxis))
4802 axis_order.append(axis_order.pop(c))
4803 axes = [axes[ao] for ao in axis_order]
4804 data = cast(NDArray[Any], data.transpose(axis_order))
4806 # h, w = data.shape[:2]
4807 # if h / w in (1.0 or 2.0):
4808 # pass
4809 # elif h / w < 2:
4810 # TODO: enforce 2:1 or 1:1 aspect ratio for generated cover images
4812 norm_along = (
4813 tuple(i for i, a in enumerate(axes) if a.type in ("space", "time")) or None
4814 )
4815 # normalize the data and map to 8 bit
4816 data = normalize(data, norm_along)
4817 data = (data * 255).astype("uint8")
4819 return data
4821 def create_diagonal_split_image(im0: NDArray[Any], im1: NDArray[Any]):
4822 assert im0.dtype == im1.dtype == np.uint8
4823 assert im0.shape == im1.shape
4824 assert im0.ndim == 3
4825 N, M, C = im0.shape
4826 assert C == 3
4827 out = np.ones((N, M, C), dtype="uint8")
4828 for c in range(C):
4829 outc = np.tril(im0[..., c])
4830 mask = outc == 0
4831 outc[mask] = np.triu(im1[..., c])[mask]
4832 out[..., c] = outc
4834 return out
4836 if not inputs:
4837 raise ValueError("Missing test input tensor for cover generation.")
4839 if not outputs:
4840 raise ValueError("Missing test output tensor for cover generation.")
4842 ipt_descr, ipt = inputs[0]
4843 out_descr, out = outputs[0]
4845 ipt_img = to_2d_image(ipt, ipt_descr.axes)
4846 out_img = to_2d_image(out, out_descr.axes)
4848 cover_folder = Path(mkdtemp())
4849 if ipt_img.shape == out_img.shape:
4850 covers = [cover_folder / "cover.png"]
4851 imwrite(covers[0], create_diagonal_split_image(ipt_img, out_img))
4852 else:
4853 covers = [cover_folder / "input.png", cover_folder / "output.png"]
4854 imwrite(covers[0], ipt_img)
4855 imwrite(covers[1], out_img)
4857 return [FileDescr(source=c) for c in covers]