Coverage for src/bioimageio/spec/model/v0_5.py: 72%
1710 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-18 09:17 +0000
1from __future__ import annotations
3import collections.abc
4import re
5import string
6import warnings
7from copy import deepcopy
8from functools import partial
9from itertools import chain
10from math import ceil
11from pathlib import Path, PurePosixPath
12from tempfile import mkdtemp
13from textwrap import dedent
14from typing import (
15 TYPE_CHECKING,
16 Any,
17 Callable,
18 ClassVar,
19 Dict,
20 Generic,
21 List,
22 Literal,
23 Mapping,
24 NamedTuple,
25 Optional,
26 Sequence,
27 TypeVar,
28 Union,
29 cast,
30 overload,
31)
33import numpy as np
34from annotated_types import Ge, Gt, Interval, MaxLen, MinLen, Predicate
35from imageio.v3 import imread, imwrite # pyright: ignore[reportUnknownVariableType]
36from loguru import logger
37from numpy.typing import NDArray
38from pydantic import (
39 AfterValidator,
40 Discriminator,
41 Field,
42 RootModel,
43 SerializationInfo,
44 SerializerFunctionWrapHandler,
45 StrictInt,
46 Tag,
47 ValidationInfo,
48 WrapSerializer,
49 field_validator,
50 model_serializer,
51 model_validator,
52)
53from pydantic_extra_types.color import Color
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_mapping, is_sequence
82from .._internal.types import (
83 FAIR,
84 AbsoluteTolerance,
85 LowerCaseIdentifier,
86 LowerCaseIdentifierAnno,
87 MismatchedElementsPerMillion,
88 RelativeTolerance,
89 validate_identifier,
90 validate_is_not_keyword,
91)
92from .._internal.types import Datetime as Datetime
93from .._internal.types import Identifier as Identifier
94from .._internal.types import NotEmpty as NotEmpty
95from .._internal.types import SiUnit as SiUnit
96from .._internal.url import HttpUrl as HttpUrl
97from .._internal.utils import try_all_raise_last
98from .._internal.validation_context import get_validation_context
99from .._internal.validator_annotations import RestrictCharacters
100from .._internal.version_type import Version as Version
101from .._internal.warning_levels import INFO
102from ..dataset.v0_2 import DatasetDescr as DatasetDescr02
103from ..dataset.v0_2 import LinkedDataset as LinkedDataset02
104from ..dataset.v0_3 import DatasetDescr as DatasetDescr
105from ..dataset.v0_3 import DatasetId as DatasetId
106from ..dataset.v0_3 import LinkedDataset as LinkedDataset
107from ..dataset.v0_3 import Uploader as Uploader
108from ..generic._v0_3_converter import convert_plain_covers_and_docs_and_icon
109from ..generic.v0_3 import (
110 VALID_COVER_IMAGE_EXTENSIONS as VALID_COVER_IMAGE_EXTENSIONS,
111)
112from ..generic.v0_3 import Author as Author
113from ..generic.v0_3 import BadgeDescr as BadgeDescr
114from ..generic.v0_3 import CiteEntry as CiteEntry
115from ..generic.v0_3 import DeprecatedLicenseId as DeprecatedLicenseId
116from ..generic.v0_3 import Doi as Doi
117from ..generic.v0_3 import (
118 FileDescr_documentation,
119 GenericModelDescrBase,
120 LinkedResourceBase,
121 _author_conv, # pyright: ignore[reportPrivateUsage]
122 _maintainer_conv, # pyright: ignore[reportPrivateUsage]
123)
124from ..generic.v0_3 import LicenseId as LicenseId
125from ..generic.v0_3 import LinkedResource as LinkedResource
126from ..generic.v0_3 import Maintainer as Maintainer
127from ..generic.v0_3 import OrcidId as OrcidId
128from ..generic.v0_3 import RelativeFilePath as RelativeFilePath
129from ..generic.v0_3 import ResourceId as ResourceId
130from .v0_4 import Author as _Author_v0_4
131from .v0_4 import BinarizeDescr as _BinarizeDescr_v0_4
132from .v0_4 import CallableFromDepencency as CallableFromDepencency
133from .v0_4 import CallableFromDepencency as _CallableFromDepencency_v0_4
134from .v0_4 import CallableFromFile as _CallableFromFile_v0_4
135from .v0_4 import ClipDescr as _ClipDescr_v0_4
136from .v0_4 import ImplicitOutputShape as _ImplicitOutputShape_v0_4
137from .v0_4 import InputTensorDescr as _InputTensorDescr_v0_4
138from .v0_4 import KnownRunMode as KnownRunMode
139from .v0_4 import ModelDescr as _ModelDescr04
140from .v0_4 import ModelDescr as _ModelDescr_v0_4
141from .v0_4 import OutputTensorDescr as _OutputTensorDescr_v0_4
142from .v0_4 import ParameterizedInputShape as _ParameterizedInputShape_v0_4
143from .v0_4 import PostprocessingDescr as _PostprocessingDescr_v0_4
144from .v0_4 import PreprocessingDescr as _PreprocessingDescr_v0_4
145from .v0_4 import RunMode as RunMode
146from .v0_4 import ScaleLinearDescr as _ScaleLinearDescr_v0_4
147from .v0_4 import ScaleMeanVarianceDescr as _ScaleMeanVarianceDescr_v0_4
148from .v0_4 import ScaleRangeDescr as _ScaleRangeDescr_v0_4
149from .v0_4 import SigmoidDescr as _SigmoidDescr_v0_4
150from .v0_4 import TensorName as _TensorName_v0_4
151from .v0_4 import ZeroMeanUnitVarianceDescr as _ZeroMeanUnitVarianceDescr_v0_4
152from .v0_4 import package_weights
154SpaceUnit = Literal[
155 "attometer",
156 "angstrom",
157 "centimeter",
158 "decimeter",
159 "exameter",
160 "femtometer",
161 "foot",
162 "gigameter",
163 "hectometer",
164 "inch",
165 "kilometer",
166 "megameter",
167 "meter",
168 "micrometer",
169 "mile",
170 "millimeter",
171 "nanometer",
172 "parsec",
173 "petameter",
174 "picometer",
175 "terameter",
176 "yard",
177 "yoctometer",
178 "yottameter",
179 "zeptometer",
180 "zettameter",
181]
182"""Space unit compatible to the [OME-Zarr axes specification 0.5](https://ngff.openmicroscopy.org/0.5/#axes-md)"""
184TimeUnit = Literal[
185 "attosecond",
186 "centisecond",
187 "day",
188 "decisecond",
189 "exasecond",
190 "femtosecond",
191 "gigasecond",
192 "hectosecond",
193 "hour",
194 "kilosecond",
195 "megasecond",
196 "microsecond",
197 "millisecond",
198 "minute",
199 "nanosecond",
200 "petasecond",
201 "picosecond",
202 "second",
203 "terasecond",
204 "yoctosecond",
205 "yottasecond",
206 "zeptosecond",
207 "zettasecond",
208]
209"""Time unit compatible to the [OME-Zarr axes specification 0.5](https://ngff.openmicroscopy.org/0.5/#axes-md)"""
211AxisType = Literal["batch", "channel", "index", "time", "space"]
213_AXIS_TYPE_MAP: Mapping[str, AxisType] = {
214 "b": "batch",
215 "t": "time",
216 "i": "index",
217 "c": "channel",
218 "x": "space",
219 "y": "space",
220 "z": "space",
221}
223_AXIS_ID_MAP = {
224 "b": "batch",
225 "t": "time",
226 "i": "index",
227 "c": "channel",
228 "s": "channel",
229}
231WeightsFormat = Literal[
232 "keras_hdf5",
233 "keras_v3",
234 "onnx",
235 "pytorch_state_dict",
236 "tensorflow_js",
237 "tensorflow_saved_model_bundle",
238 "torchscript",
239]
242class TensorId(LowerCaseIdentifier):
243 root_model: ClassVar[type[RootModel[Any]]] = RootModel[
244 Annotated[LowerCaseIdentifierAnno, MaxLen(32)]
245 ]
248def _normalize_axis_id(a: str):
249 b = str(a).lower()
250 normalized = _AXIS_ID_MAP.get(b, b)
251 if a != normalized:
252 logger.opt(depth=3).warning(
253 "Normalized axis id from '{}' to '{}'.", a, normalized
254 )
255 return normalized
258class AxisId(LowerCaseIdentifier):
259 root_model: ClassVar[type[RootModel[Any]]] = RootModel[
260 Annotated[
261 NotEmpty[str],
262 AfterValidator(_normalize_axis_id),
263 MaxLen(16),
264 AfterValidator(validate_identifier),
265 AfterValidator(validate_is_not_keyword),
266 ]
267 ]
270def _is_batch(a: str) -> bool:
271 return str(a) == "batch"
274def _is_not_batch(a: str) -> bool:
275 return not _is_batch(a)
278NonBatchAxisId = Annotated[AxisId, Predicate(_is_not_batch)]
280PreprocessingId = Literal[
281 "binarize",
282 "clip",
283 "ensure_dtype",
284 "fixed_zero_mean_unit_variance",
285 "scale_linear",
286 "scale_range",
287 "sigmoid",
288 "softmax",
289]
290PostprocessingId = Literal[
291 "binarize",
292 "clip",
293 "custom",
294 "ensure_dtype",
295 "fixed_zero_mean_unit_variance",
296 "scale_linear",
297 "scale_mean_variance",
298 "scale_range",
299 "sigmoid",
300 "softmax",
301 "zero_mean_unit_variance",
302]
305SAME_AS_TYPE = "<same as type>"
308ParameterizedSize_N: TypeAlias = int
309"""
310Annotates an integer to calculate a concrete axis size from a `ParameterizedSize`.
311"""
314class ParameterizedSize(Node):
315 """Describes a range of valid tensor axis sizes as `size = min + n*step`.
317 - **min** and **step** are given by the model description.
318 - All blocksize paramters n = 0,1,2,... yield a valid `size`.
319 - A greater blocksize paramter n = 0,1,2,... results in a greater **size**.
320 This allows to adjust the axis size more generically.
321 """
323 N: ClassVar[type[int]] = ParameterizedSize_N
324 """Positive integer to parameterize this axis"""
326 min: Annotated[int, Gt(0)]
327 step: Annotated[int, Gt(0)]
329 def validate_size(self, size: int, msg_prefix: str = "") -> int:
330 if size < self.min:
331 raise ValueError(
332 f"{msg_prefix}size {size} < {self.min} (minimum axis size)"
333 )
334 if (size - self.min) % self.step != 0:
335 raise ValueError(
336 f"{msg_prefix}size {size} is not parameterized by `min + n*step` ="
337 + f" `{self.min} + n*{self.step}`"
338 )
340 return size
342 def get_size(self, n: ParameterizedSize_N) -> int:
343 return self.min + self.step * n
345 def get_n(self, s: int) -> ParameterizedSize_N:
346 """return smallest n parameterizing a size greater or equal than `s`"""
347 return ceil((s - self.min) / self.step)
350class DataDependentSize(Node):
351 min: Annotated[int, Gt(0)] = 1
352 max: Annotated[int | None, Gt(1)] = None
354 @model_validator(mode="after")
355 def _validate_max_gt_min(self):
356 if self.max is not None and self.min >= self.max:
357 raise ValueError(f"expected `min` < `max`, but got {self.min}, {self.max}")
359 return self
361 def validate_size(self, size: int, msg_prefix: str = "") -> int:
362 if size < self.min:
363 raise ValueError(f"{msg_prefix}size {size} < {self.min}")
365 if self.max is not None and size > self.max:
366 raise ValueError(f"{msg_prefix}size {size} > {self.max}")
368 return size
371class SizeReference(Node):
372 """A tensor axis size (extent in pixels/frames) defined in relation to a reference axis.
374 `axis.size = reference.size * reference.scale / axis.scale + offset`
376 Note:
377 1. The axis and the referenced axis need to have the same unit (or no unit).
378 2. Batch axes may not be referenced.
379 3. Fractions are rounded down.
380 4. If the reference axis is `concatenable` the referencing axis is assumed to be
381 `concatenable` as well with the same block order.
383 Example:
384 An unisotropic input image of w*h=100*49 pixels depicts a phsical space of 200*196mm².
385 Let's assume that we want to express the image height h in relation to its width w
386 instead of only accepting input images of exactly 100*49 pixels
387 (for example to express a range of valid image shapes by parametrizing w, see `ParameterizedSize`).
389 >>> w = SpaceInputAxis(id=AxisId("w"), size=100, unit="millimeter", scale=2)
390 >>> h = SpaceInputAxis(
391 ... id=AxisId("h"),
392 ... size=SizeReference(tensor_id=TensorId("input"), axis_id=AxisId("w"), offset=-1),
393 ... unit="millimeter",
394 ... scale=4,
395 ... )
396 >>> print(h.size.get_size(h, w))
397 49
399 ⇒ h = w * w.scale / h.scale + offset = 100 * 2mm / 4mm - 1 = 49
400 """
402 tensor_id: TensorId
403 """tensor id of the reference axis"""
405 axis_id: AxisId
406 """axis id of the reference axis"""
408 offset: StrictInt = 0
410 def get_size(
411 self,
412 axis: ChannelAxis
413 | IndexInputAxis
414 | IndexOutputAxis
415 | TimeInputAxis
416 | SpaceInputAxis
417 | TimeOutputAxis
418 | TimeOutputAxisWithHalo
419 | SpaceOutputAxis
420 | SpaceOutputAxisWithHalo,
421 ref_axis: ChannelAxis
422 | IndexInputAxis
423 | IndexOutputAxis
424 | TimeInputAxis
425 | SpaceInputAxis
426 | TimeOutputAxis
427 | TimeOutputAxisWithHalo
428 | SpaceOutputAxis
429 | SpaceOutputAxisWithHalo,
430 n: ParameterizedSize_N = 0,
431 ref_size: int | None = None,
432 ):
433 """Compute the concrete size for a given axis and its reference axis.
435 Args:
436 axis: The axis this [SizeReference][] is the size of.
437 ref_axis: The reference axis to compute the size from.
438 n: If the **ref_axis** is parameterized (of type `ParameterizedSize`)
439 and no fixed **ref_size** is given,
440 **n** is used to compute the size of the parameterized **ref_axis**.
441 ref_size: Overwrite the reference size instead of deriving it from
442 **ref_axis**
443 (**ref_axis.scale** is still used; any given **n** is ignored).
444 """
445 assert axis.size == self, (
446 "Given `axis.size` is not defined by this `SizeReference`"
447 )
449 assert ref_axis.id == self.axis_id, (
450 f"Expected `ref_axis.id` to be {self.axis_id}, but got {ref_axis.id}."
451 )
453 assert axis.unit == ref_axis.unit, (
454 "`SizeReference` requires `axis` and `ref_axis` to have the same `unit`,"
455 f" but {axis.unit}!={ref_axis.unit}"
456 )
457 if ref_size is None:
458 if isinstance(ref_axis.size, (int, float)):
459 ref_size = ref_axis.size
460 elif isinstance(ref_axis.size, ParameterizedSize):
461 ref_size = ref_axis.size.get_size(n)
462 elif isinstance(ref_axis.size, DataDependentSize):
463 raise ValueError(
464 "Reference axis referenced in `SizeReference` may not be a `DataDependentSize`."
465 )
466 elif isinstance(ref_axis.size, SizeReference):
467 raise ValueError(
468 "Reference axis referenced in `SizeReference` may not be sized by a"
469 + " `SizeReference` itself."
470 )
471 else:
472 assert_never(ref_axis.size)
474 return int(ref_size * ref_axis.scale / axis.scale + self.offset)
476 @staticmethod
477 def _get_unit(
478 axis: ChannelAxis
479 | IndexInputAxis
480 | IndexOutputAxis
481 | TimeInputAxis
482 | SpaceInputAxis
483 | TimeOutputAxis
484 | TimeOutputAxisWithHalo
485 | SpaceOutputAxis
486 | SpaceOutputAxisWithHalo,
487 ):
488 return axis.unit
491class AxisBase(NodeWithExplicitlySetFields):
492 id: AxisId
493 """An axis id unique across all axes of one tensor."""
495 description: Annotated[str, MaxLen(128)] = ""
496 """A short description of this axis beyond its type and id."""
499class WithHalo(Node):
500 halo: Annotated[int, Ge(1)]
501 """The halo should be cropped from the output tensor to avoid boundary effects.
502 It is to be cropped from both sides, i.e. `size_after_crop = size - 2 * halo`.
503 To document a halo that is already cropped by the model use `size.offset` instead."""
505 size: Annotated[
506 SizeReference,
507 Field(examples=[{"tensor_id": "t", "axis_id": "a", "offset": 5}]),
508 ]
509 """reference to another axis with an optional offset (see [SizeReference][])"""
512BATCH_AXIS_ID = AxisId("batch")
513CHANNEL_AXIS_ID = AxisId("channel")
514DEFAULT_SPACE_AXIS_ID = AxisId("x")
515DEFAULT_INDEX_AXIS_ID = AxisId("index")
516DEFAULT_TIME_AXIS_ID = AxisId("time")
519class BatchAxis(AxisBase):
520 implemented_type: ClassVar[Literal["batch"]] = "batch"
521 if TYPE_CHECKING:
522 type: Literal["batch"] = "batch"
523 else:
524 type: Literal["batch"]
526 id: Annotated[AxisId, Predicate(_is_batch)] = BATCH_AXIS_ID
527 size: Literal[1] | None = None
528 """The batch size may be fixed to 1,
529 otherwise (the default) it may be chosen arbitrarily depending on available memory"""
531 @property
532 def scale(self):
533 return 1.0
535 @property
536 def concatenable(self):
537 return True
539 @property
540 def unit(self):
541 return None
544class ChannelAxis(AxisBase):
545 implemented_type: ClassVar[Literal["channel"]] = "channel"
546 if TYPE_CHECKING:
547 type: Literal["channel"] = "channel"
548 else:
549 type: Literal["channel"]
551 id: NonBatchAxisId = CHANNEL_AXIS_ID
553 channel_names: NotEmpty[list[str]]
554 """Name/label for each channel. The number of channels is given by `len(channel_names)`."""
556 channel_colors: NotEmpty[list[Color]] = Field(
557 default_factory=cast(Callable[[], list[Color]], list)
558 ) # real default is set by _set_default_channel_colors()
559 """Colors for each channel for visualization purposes.
560 If not given, a default color palette is used:
562 - For < 8 channels: colorblind-friendly palette from https://www.nature.com/articles/nmeth.1618 (without black)
563 - For < 21 channels: discrete matplotlib colormap 'tab20b' (redistributed for more even color distribution < 20 channels)
564 - For >= 21 channels: sample colors from continuous matplotlib colormap 'cividis'
566 """
568 @model_validator(mode="before")
569 @classmethod
570 def _set_default_channel_colors(cls, data: Any):
571 if (
572 is_mapping(data)
573 and not data.get("channel_colors")
574 and is_sequence(channel_names := data.get("channel_names"))
575 ):
576 n_channels = len(channel_names)
577 data = dict(data)
578 if n_channels < 8:
579 # use colorblind-friendly palette from https://www.nature.com/articles/nmeth.1618
580 # (without black)
581 data["channel_colors"] = [
582 "#E69F00",
583 "#56B4E9",
584 "#009E73",
585 "#F0E442",
586 "#0072B2",
587 "#D55E00",
588 "#CC79A7",
589 ][:n_channels]
590 elif n_channels < 21:
591 # use discrete matplotlib colormap 'tab20b'
592 # (redistributed for more even color distribution < 20 channels)
593 data["channel_colors"] = [
594 "#393b79",
595 "#8ca252",
596 "#e7ba52",
597 "#e7969c",
598 "#7b4173",
599 "#5254a3",
600 "#b5cf6b",
601 "#e7cb94",
602 "#843c39",
603 "#a55194",
604 "#6b6ecf",
605 "#cedb9c",
606 "#8c6d31",
607 "#d6616b",
608 "#ce6dbd",
609 "#9c9ede",
610 "#637939",
611 "#bd9e39",
612 "#ad494a",
613 "#de9ed6",
614 ][:n_channels]
615 else:
616 # sample colors from continuous matplotlib colormap 'cividis'
617 import matplotlib.colors
618 import matplotlib.pyplot as plt
620 cmap = plt.colormaps["cividis"].resampled(n_channels)
621 data["channel_colors"] = [
622 matplotlib.colors.to_hex(cmap(i)) for i in range(n_channels)
623 ]
625 return data
627 @property
628 def size(self) -> int:
629 return len(self.channel_names)
631 @property
632 def concatenable(self):
633 return False
635 @property
636 def scale(self) -> float:
637 return 1.0
639 @property
640 def unit(self):
641 return None
644class _WithInputAxisSize(Node):
645 size: Annotated[
646 Annotated[int, Gt(0)] | ParameterizedSize | SizeReference,
647 Field(
648 examples=[
649 10,
650 ParameterizedSize(min=32, step=16).model_dump(mode="json"),
651 {"tensor_id": "t", "axis_id": "a", "offset": 5},
652 ]
653 ),
654 ]
655 """The size/length of this axis can be specified as
656 - fixed integer
657 - parameterized series of valid sizes ([ParameterizedSize][])
658 - reference to another axis with an optional offset ([SizeReference][])
659 """
662class IndexAxisBase(AxisBase):
663 implemented_type: ClassVar[Literal["index"]] = "index"
664 if TYPE_CHECKING:
665 type: Literal["index"] = "index"
666 else:
667 type: Literal["index"]
669 id: NonBatchAxisId = DEFAULT_INDEX_AXIS_ID
671 @property
672 def scale(self) -> float:
673 return 1.0
675 @property
676 def unit(self):
677 return None
680class IndexInputAxis(IndexAxisBase, _WithInputAxisSize):
681 concatenable: bool = False
682 """If a model has a `concatenable` input axis, it can be processed blockwise,
683 splitting a longer sample axis into blocks matching its input tensor description.
684 Output axes are concatenable if they have a [SizeReference][] to a concatenable
685 input axis.
686 """
689class IndexOutputAxis(IndexAxisBase):
690 size: Annotated[
691 Annotated[int, Gt(0)] | SizeReference | DataDependentSize,
692 Field(examples=[10, {"tensor_id": "t", "axis_id": "a", "offset": 5}]),
693 ]
694 """The size/length of this axis can be specified as
695 - fixed integer
696 - reference to another axis with an optional offset ([SizeReference][])
697 - data dependent size using [DataDependentSize][] (size is only known after model inference)
698 """
701class TimeAxisBase(AxisBase):
702 implemented_type: ClassVar[Literal["time"]] = "time"
703 if TYPE_CHECKING:
704 type: Literal["time"] = "time"
705 else:
706 type: Literal["time"]
708 id: NonBatchAxisId = DEFAULT_TIME_AXIS_ID
709 unit: TimeUnit | None = None
710 scale: Annotated[float, Gt(0)] = 1.0
713class TimeInputAxis(TimeAxisBase, _WithInputAxisSize):
714 concatenable: bool = False
715 """If a model has a `concatenable` input axis, it can be processed blockwise,
716 splitting a longer sample axis into blocks matching its input tensor description.
717 Output axes are concatenable if they have a [SizeReference][] to a concatenable
718 input axis.
719 """
722class SpaceAxisBase(AxisBase):
723 implemented_type: ClassVar[Literal["space"]] = "space"
724 if TYPE_CHECKING:
725 type: Literal["space"] = "space"
726 else:
727 type: Literal["space"]
729 id: Annotated[NonBatchAxisId, Field(examples=["x", "y", "z"])] = (
730 DEFAULT_SPACE_AXIS_ID
731 )
732 unit: SpaceUnit | None = None
733 scale: Annotated[float, Gt(0)] = 1.0
736class SpaceInputAxis(SpaceAxisBase, _WithInputAxisSize):
737 concatenable: bool = False
738 """If a model has a `concatenable` input axis, it can be processed blockwise,
739 splitting a longer sample axis into blocks matching its input tensor description.
740 Output axes are concatenable if they have a [SizeReference][] to a concatenable
741 input axis.
742 """
745INPUT_AXIS_TYPES = (
746 BatchAxis,
747 ChannelAxis,
748 IndexInputAxis,
749 TimeInputAxis,
750 SpaceInputAxis,
751)
752"""intended for isinstance comparisons in py<3.10"""
754_InputAxisUnion = Union[
755 BatchAxis, ChannelAxis, IndexInputAxis, TimeInputAxis, SpaceInputAxis
756]
757InputAxis = Annotated[_InputAxisUnion, Discriminator("type")]
760class _WithOutputAxisSize(Node):
761 size: Annotated[
762 Annotated[int, Gt(0)] | SizeReference,
763 Field(examples=[10, {"tensor_id": "t", "axis_id": "a", "offset": 5}]),
764 ]
765 """The size/length of this axis can be specified as
766 - fixed integer
767 - reference to another axis with an optional offset (see [SizeReference][])
768 """
771class TimeOutputAxis(TimeAxisBase, _WithOutputAxisSize):
772 pass
775class TimeOutputAxisWithHalo(TimeAxisBase, WithHalo):
776 pass
779def _get_halo_axis_discriminator_value(v: Any) -> Literal["with_halo", "wo_halo"]:
780 if isinstance(v, dict):
781 return "with_halo" if "halo" in v else "wo_halo"
782 else:
783 return "with_halo" if hasattr(v, "halo") else "wo_halo"
786_TimeOutputAxisUnion = Annotated[
787 Union[
788 Annotated[TimeOutputAxis, Tag("wo_halo")],
789 Annotated[TimeOutputAxisWithHalo, Tag("with_halo")],
790 ],
791 Discriminator(_get_halo_axis_discriminator_value),
792]
795class SpaceOutputAxis(SpaceAxisBase, _WithOutputAxisSize):
796 pass
799class SpaceOutputAxisWithHalo(SpaceAxisBase, WithHalo):
800 pass
803_SpaceOutputAxisUnion = Annotated[
804 Union[
805 Annotated[SpaceOutputAxis, Tag("wo_halo")],
806 Annotated[SpaceOutputAxisWithHalo, Tag("with_halo")],
807 ],
808 Discriminator(_get_halo_axis_discriminator_value),
809]
812_OutputAxisUnion = Union[
813 BatchAxis, ChannelAxis, IndexOutputAxis, _TimeOutputAxisUnion, _SpaceOutputAxisUnion
814]
815OutputAxis = Annotated[_OutputAxisUnion, Discriminator("type")]
817OUTPUT_AXIS_TYPES = (
818 BatchAxis,
819 ChannelAxis,
820 IndexOutputAxis,
821 TimeOutputAxis,
822 TimeOutputAxisWithHalo,
823 SpaceOutputAxis,
824 SpaceOutputAxisWithHalo,
825)
826"""intended for isinstance comparisons in py<3.10"""
829AnyAxis = Union[InputAxis, OutputAxis]
831ANY_AXIS_TYPES = INPUT_AXIS_TYPES + OUTPUT_AXIS_TYPES
832"""intended for isinstance comparisons in py<3.10"""
834TVs = Union[
835 NotEmpty[List[int]],
836 NotEmpty[List[float]],
837 NotEmpty[List[bool]],
838 NotEmpty[List[str]],
839]
842NominalOrOrdinalDType = Literal[
843 "float32",
844 "float64",
845 "uint8",
846 "int8",
847 "uint16",
848 "int16",
849 "uint32",
850 "int32",
851 "uint64",
852 "int64",
853 "bool",
854]
857class NominalOrOrdinalDataDescr(Node):
858 values: TVs
859 """A fixed set of nominal or an ascending sequence of ordinal values.
860 In this case `data.type` is required to be an unsigend integer type, e.g. 'uint8'.
861 String `values` are interpreted as labels for tensor values 0, ..., N.
862 Note: as YAML 1.2 does not natively support a "set" datatype,
863 nominal values should be given as a sequence (aka list/array) as well.
864 """
866 type: Annotated[
867 NominalOrOrdinalDType,
868 Field(
869 examples=[
870 "float32",
871 "uint8",
872 "uint16",
873 "int64",
874 "bool",
875 ],
876 ),
877 ] = "uint8"
879 @model_validator(mode="after")
880 def _validate_values_match_type(
881 self,
882 ) -> Self:
883 incompatible: list[Any] = []
884 for v in self.values:
885 if self.type == "bool":
886 if not isinstance(v, bool):
887 incompatible.append(v)
888 elif self.type in DTYPE_LIMITS:
889 if (
890 isinstance(v, (int, float))
891 and (
892 v < DTYPE_LIMITS[self.type].min
893 or v > DTYPE_LIMITS[self.type].max
894 )
895 or (isinstance(v, str) and "uint" not in self.type)
896 or (isinstance(v, float) and "int" in self.type)
897 ):
898 incompatible.append(v)
899 else:
900 incompatible.append(v)
902 if len(incompatible) == 5:
903 incompatible.append("...")
904 break
906 if incompatible:
907 raise ValueError(
908 f"data type '{self.type}' incompatible with values {incompatible}"
909 )
911 return self
913 unit: Literal["arbitrary unit"] | SiUnit | None = None
915 @property
916 def range(self):
917 if isinstance(self.values[0], str):
918 return 0, len(self.values) - 1
919 else:
920 return min(self.values), max(self.values)
923IntervalOrRatioDType = Literal[
924 "float32",
925 "float64",
926 "uint8",
927 "int8",
928 "uint16",
929 "int16",
930 "uint32",
931 "int32",
932 "uint64",
933 "int64",
934]
937class IntervalOrRatioDataDescr(Node):
938 type: Annotated[ # TODO: rename to dtype
939 IntervalOrRatioDType,
940 Field(
941 examples=["float32", "float64", "uint8", "uint16"],
942 ),
943 ] = "float32"
944 range: tuple[float | None, float | None] = (
945 None,
946 None,
947 )
948 """Tuple `(minimum, maximum)` specifying the allowed range of the data in this tensor.
949 `None` corresponds to min/max of what can be expressed by **type**."""
950 unit: Literal["arbitrary unit"] | SiUnit = "arbitrary unit"
951 scale: float = 1.0
952 """Scale for data on an interval (or ratio) scale."""
953 offset: float | None = None
954 """Offset for data on a ratio scale."""
956 @model_validator(mode="before")
957 def _replace_inf(cls, data: Any):
958 if is_dict(data) and "range" in data and is_sequence(data["range"]):
959 forbidden = (
960 "inf",
961 "-inf",
962 ".inf",
963 "-.inf",
964 float("inf"),
965 float("-inf"),
966 )
967 if any(v in forbidden for v in data["range"]):
968 issue_warning("replaced 'inf' value", value=data["range"])
970 data["range"] = tuple(
971 (None if v in forbidden else v) for v in data["range"]
972 )
974 return data
977TensorDataDescr = Union[NominalOrOrdinalDataDescr, IntervalOrRatioDataDescr]
980class BinarizeKwargs(KwargsNode):
981 """key word arguments for [BinarizeDescr][]"""
983 threshold: float
984 """The fixed threshold"""
987class BinarizeAlongAxisKwargs(KwargsNode):
988 """key word arguments for [BinarizeDescr][]"""
990 threshold: NotEmpty[list[float]]
991 """The fixed threshold values along `axis`"""
993 axis: Annotated[NonBatchAxisId, Field(examples=["channel"])]
994 """The `threshold` axis"""
997class BinarizeDescr(NodeWithExplicitlySetFields):
998 """Binarize the tensor with a fixed threshold.
1000 Values above [BinarizeKwargs.threshold][]/[BinarizeAlongAxisKwargs.threshold][]
1001 will be set to one, values below the threshold to zero.
1003 Examples:
1004 - in YAML
1005 ```yaml
1006 postprocessing:
1007 - id: binarize
1008 kwargs:
1009 axis: 'channel'
1010 threshold: [0.25, 0.5, 0.75]
1011 ```
1012 - in Python:
1014 >>> postprocessing = [BinarizeDescr(
1015 ... kwargs=BinarizeAlongAxisKwargs(
1016 ... axis=AxisId('channel'),
1017 ... threshold=[0.25, 0.5, 0.75],
1018 ... )
1019 ... )]
1020 """
1022 implemented_id: ClassVar[Literal["binarize"]] = "binarize"
1023 if TYPE_CHECKING:
1024 id: Literal["binarize"] = "binarize"
1025 else:
1026 id: Literal["binarize"]
1027 kwargs: BinarizeKwargs | BinarizeAlongAxisKwargs
1030class ClipKwargs(KwargsNode):
1031 """key word arguments for [ClipDescr][]"""
1033 min: float | None = None
1034 """Minimum value for clipping.
1036 Exclusive with [min_percentile][]
1037 """
1038 min_percentile: Annotated[float, Interval(ge=0, lt=100)] | None = None
1039 """Minimum percentile for clipping.
1041 Exclusive with [min][].
1043 In range [0, 100).
1044 """
1046 max: float | None = None
1047 """Maximum value for clipping.
1049 Exclusive with `max_percentile`.
1050 """
1051 max_percentile: Annotated[float, Interval(gt=1, le=100)] | None = None
1052 """Maximum percentile for clipping.
1054 Exclusive with `max`.
1056 In range (1, 100].
1057 """
1059 axes: Annotated[Sequence[AxisId] | None, Field(examples=[("batch", "x", "y")])] = (
1060 None
1061 )
1062 """The subset of axes to determine percentiles jointly,
1064 i.e. axes to reduce to compute min/max from `min_percentile`/`max_percentile`.
1065 For example to clip 'batch', 'x' and 'y' jointly in a tensor ('batch', 'channel', 'y', 'x')
1066 resulting in a tensor of equal shape with clipped values per channel, specify `axes=('batch', 'x', 'y')`.
1067 To clip samples independently, leave out the 'batch' axis.
1069 Only valid if `min_percentile` and/or `max_percentile` are set.
1071 Default: Compute percentiles over all axes jointly."""
1073 @model_validator(mode="after")
1074 def _validate(self) -> Self:
1075 if (self.min is not None) and (self.min_percentile is not None):
1076 raise ValueError(
1077 "Only one of `min` and `min_percentile` may be set, not both."
1078 )
1079 if (self.max is not None) and (self.max_percentile is not None):
1080 raise ValueError(
1081 "Only one of `max` and `max_percentile` may be set, not both."
1082 )
1083 if (
1084 self.min is None
1085 and self.min_percentile is None
1086 and self.max is None
1087 and self.max_percentile is None
1088 ):
1089 raise ValueError(
1090 "At least one of `min`, `min_percentile`, `max`, or `max_percentile` must be set."
1091 )
1093 if (
1094 self.axes is not None
1095 and self.min_percentile is None
1096 and self.max_percentile is None
1097 ):
1098 raise ValueError(
1099 "If `axes` is set, at least one of `min_percentile` or `max_percentile` must be set."
1100 )
1102 return self
1105class ClipDescr(NodeWithExplicitlySetFields):
1106 """Set tensor values below min to min and above max to max.
1108 See `ScaleRangeDescr` for examples.
1109 """
1111 implemented_id: ClassVar[Literal["clip"]] = "clip"
1112 if TYPE_CHECKING:
1113 id: Literal["clip"] = "clip"
1114 else:
1115 id: Literal["clip"]
1117 kwargs: ClipKwargs
1120class EnsureDtypeKwargs(KwargsNode):
1121 """key word arguments for [EnsureDtypeDescr][]"""
1123 dtype: Literal[
1124 "float32",
1125 "float64",
1126 "uint8",
1127 "int8",
1128 "uint16",
1129 "int16",
1130 "uint32",
1131 "int32",
1132 "uint64",
1133 "int64",
1134 "bool",
1135 ]
1138class EnsureDtypeDescr(NodeWithExplicitlySetFields):
1139 """Cast the tensor data type to `EnsureDtypeKwargs.dtype` (if not matching).
1141 This can for example be used to ensure the inner neural network model gets a
1142 different input tensor data type than the fully described bioimage.io model does.
1144 Examples:
1145 The described bioimage.io model (incl. preprocessing) accepts any
1146 float32-compatible tensor, normalizes it with percentiles and clipping and then
1147 casts it to uint8, which is what the neural network in this example expects.
1148 - in YAML
1149 ```yaml
1150 inputs:
1151 - data:
1152 type: float32 # described bioimage.io model is compatible with any float32 input tensor
1153 preprocessing:
1154 - id: scale_range
1155 kwargs:
1156 axes: ['y', 'x']
1157 max_percentile: 99.8
1158 min_percentile: 5.0
1159 - id: clip
1160 kwargs:
1161 min: 0.0
1162 max: 1.0
1163 - id: ensure_dtype # the neural network of the model requires uint8
1164 kwargs:
1165 dtype: uint8
1166 ```
1167 - in Python:
1168 >>> preprocessing = [
1169 ... ScaleRangeDescr(
1170 ... kwargs=ScaleRangeKwargs(
1171 ... axes= (AxisId('y'), AxisId('x')),
1172 ... max_percentile= 99.8,
1173 ... min_percentile= 5.0,
1174 ... )
1175 ... ),
1176 ... ClipDescr(kwargs=ClipKwargs(min=0.0, max=1.0)),
1177 ... EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype="uint8")),
1178 ... ]
1179 """
1181 implemented_id: ClassVar[Literal["ensure_dtype"]] = "ensure_dtype"
1182 if TYPE_CHECKING:
1183 id: Literal["ensure_dtype"] = "ensure_dtype"
1184 else:
1185 id: Literal["ensure_dtype"]
1187 kwargs: EnsureDtypeKwargs
1190class ScaleLinearKwargs(KwargsNode):
1191 """Key word arguments for [ScaleLinearDescr][]"""
1193 gain: float = 1.0
1194 """multiplicative factor"""
1196 offset: float = 0.0
1197 """additive term"""
1199 @model_validator(mode="after")
1200 def _validate(self) -> Self:
1201 if self.gain == 1.0 and self.offset == 0.0:
1202 raise ValueError(
1203 "Redundant linear scaling not allowd. Set `gain` != 1.0 and/or `offset`"
1204 + " != 0.0."
1205 )
1207 return self
1210class ScaleLinearAlongAxisKwargs(KwargsNode):
1211 """Key word arguments for [ScaleLinearDescr][]"""
1213 axis: Annotated[NonBatchAxisId, Field(examples=["channel"])]
1214 """The axis of gain and offset values."""
1216 gain: float | NotEmpty[list[float]] = 1.0
1217 """multiplicative factor"""
1219 offset: float | NotEmpty[list[float]] = 0.0
1220 """additive term"""
1222 @model_validator(mode="after")
1223 def _validate(self) -> Self:
1224 if isinstance(self.gain, list):
1225 if isinstance(self.offset, list):
1226 if len(self.gain) != len(self.offset):
1227 raise ValueError(
1228 f"Size of `gain` ({len(self.gain)}) and `offset` ({len(self.offset)}) must match."
1229 )
1230 else:
1231 self.offset = [float(self.offset)] * len(self.gain)
1232 elif isinstance(self.offset, list):
1233 self.gain = [float(self.gain)] * len(self.offset)
1234 else:
1235 raise ValueError(
1236 "Do not specify an `axis` for scalar gain and offset values."
1237 )
1239 if all(g == 1.0 for g in self.gain) and all(off == 0.0 for off in self.offset):
1240 raise ValueError(
1241 "Redundant linear scaling not allowd. Set `gain` != 1.0 and/or `offset`"
1242 + " != 0.0."
1243 )
1245 return self
1248class ScaleLinearDescr(NodeWithExplicitlySetFields):
1249 """Fixed linear scaling.
1251 Examples:
1252 1. Scale with scalar gain and offset
1253 - in YAML
1254 ```yaml
1255 preprocessing:
1256 - id: scale_linear
1257 kwargs:
1258 gain: 2.0
1259 offset: 3.0
1260 ```
1261 - in Python:
1263 >>> preprocessing = [
1264 ... ScaleLinearDescr(kwargs=ScaleLinearKwargs(gain= 2.0, offset=3.0))
1265 ... ]
1267 2. Independent scaling along an axis
1268 - in YAML
1269 ```yaml
1270 preprocessing:
1271 - id: scale_linear
1272 kwargs:
1273 axis: 'channel'
1274 gain: [1.0, 2.0, 3.0]
1275 ```
1276 - in Python:
1278 >>> preprocessing = [
1279 ... ScaleLinearDescr(
1280 ... kwargs=ScaleLinearAlongAxisKwargs(
1281 ... axis=AxisId("channel"),
1282 ... gain=[1.0, 2.0, 3.0],
1283 ... )
1284 ... )
1285 ... ]
1287 """
1289 implemented_id: ClassVar[Literal["scale_linear"]] = "scale_linear"
1290 if TYPE_CHECKING:
1291 id: Literal["scale_linear"] = "scale_linear"
1292 else:
1293 id: Literal["scale_linear"]
1294 kwargs: ScaleLinearKwargs | ScaleLinearAlongAxisKwargs
1297class SigmoidDescr(NodeWithExplicitlySetFields):
1298 """The logistic sigmoid function, a.k.a. expit function.
1300 Examples:
1301 - in YAML
1302 ```yaml
1303 postprocessing:
1304 - id: sigmoid
1305 ```
1306 - in Python:
1308 >>> postprocessing = [SigmoidDescr()]
1309 """
1311 implemented_id: ClassVar[Literal["sigmoid"]] = "sigmoid"
1312 if TYPE_CHECKING:
1313 id: Literal["sigmoid"] = "sigmoid"
1314 else:
1315 id: Literal["sigmoid"]
1317 @property
1318 def kwargs(self) -> KwargsNode:
1319 """empty kwargs"""
1320 return KwargsNode()
1323class SoftmaxKwargs(KwargsNode):
1324 """key word arguments for [SoftmaxDescr][]"""
1326 axis: Annotated[NonBatchAxisId, Field(examples=["channel"])] = CHANNEL_AXIS_ID
1327 """The axis to apply the softmax function along.
1328 Note:
1329 Defaults to 'channel' axis
1330 (which may not exist, in which case
1331 a different axis id has to be specified).
1332 """
1335class SoftmaxDescr(NodeWithExplicitlySetFields):
1336 """The softmax function.
1338 Examples:
1339 - in YAML
1340 ```yaml
1341 postprocessing:
1342 - id: softmax
1343 kwargs:
1344 axis: channel
1345 ```
1346 - in Python:
1348 >>> postprocessing = [SoftmaxDescr(kwargs=SoftmaxKwargs(axis=AxisId("channel")))]
1349 """
1351 implemented_id: ClassVar[Literal["softmax"]] = "softmax"
1352 if TYPE_CHECKING:
1353 id: Literal["softmax"] = "softmax"
1354 else:
1355 id: Literal["softmax"]
1357 kwargs: SoftmaxKwargs = Field(default_factory=SoftmaxKwargs.model_construct)
1360class _StardistPostprocessingKwargsBase(KwargsNode):
1361 """key word arguments for [StardistPostprocessingDescr][]"""
1363 prob_threshold: float
1364 """The probability threshold for object candidate selection."""
1366 nms_threshold: float
1367 """The IoU threshold for non-maximum suppression."""
1369 n_rays: int
1370 """Number of radial lines (rays) cast from the center of an object to its boundary."""
1373class StardistPostprocessingKwargs2D(_StardistPostprocessingKwargsBase):
1374 grid: tuple[int, int]
1375 """Grid size of network predictions."""
1377 b: int | tuple[tuple[int, int], tuple[int, int]]
1378 """Border region in which object probability is set to zero."""
1381class StardistPostprocessingKwargs3D(_StardistPostprocessingKwargsBase):
1382 grid: tuple[int, int, int]
1383 """Grid size of network predictions."""
1385 b: int | tuple[tuple[int, int], tuple[int, int], tuple[int, int]]
1386 """Border region in which object probability is set to zero."""
1388 anisotropy: tuple[float, float, float]
1389 """Anisotropy factors for 3D star-convex polyhedra, i.e. the physical pixel size along each spatial axis."""
1391 overlap_label: int | None = None
1392 """Optional label to apply to any area of overlapping predicted objects."""
1395class StardistPostprocessingDescr(NodeWithExplicitlySetFields):
1396 """Stardist postprocessing including non-maximum suppression and converting polygon representations to instance labels
1398 as described in:
1399 - Uwe Schmidt, Martin Weigert, Coleman Broaddus, and Gene Myers.
1400 [*Cell Detection with Star-convex Polygons*](https://arxiv.org/abs/1806.03535).
1401 International Conference on Medical Image Computing and Computer-Assisted Intervention (MICCAI), Granada, Spain, September 2018.
1402 - Martin Weigert, Uwe Schmidt, Robert Haase, Ko Sugawara, and Gene Myers.
1403 [*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).
1404 The IEEE Winter Conference on Applications of Computer Vision (WACV), Snowmass Village, Colorado, March 2020.
1406 Note: Only available if the `stardist` package is installed.
1407 """
1409 implemented_id: ClassVar[Literal["stardist_postprocessing"]] = (
1410 "stardist_postprocessing"
1411 )
1412 if TYPE_CHECKING:
1413 id: Literal["stardist_postprocessing"] = "stardist_postprocessing"
1414 else:
1415 id: Literal["stardist_postprocessing"]
1417 kwargs: StardistPostprocessingKwargs2D | StardistPostprocessingKwargs3D
1420class CellposeFlowDynamicsKwargs(KwargsNode):
1421 """key word arguments for [CellposeFlowDynamicsDescr][]"""
1423 cellprob_threshold: float
1424 flow_threshold: float
1425 do_3D: bool
1426 min_size: int = 15
1427 """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."""
1428 output_dtype: Literal["uint16", "uint32"] = "uint16"
1431class CellposeFlowDynamicsDescr(NodeWithExplicitlySetFields):
1432 """Cellpose flow dynamics postprocessing as described in:
1433 - Carsen Stringer and Marius Pachitariu. [*Cellpose: a generalist algorithm for cellular segmentation*](https://www.nature.com/articles/s41592-020-01018-x). Nature Methods, 2021.
1435 Note: Only available if the `cellpose` package is installed.
1436 """
1438 implemented_id: ClassVar[Literal["cellpose_flow_dynamics"]] = (
1439 "cellpose_flow_dynamics"
1440 )
1441 if TYPE_CHECKING:
1442 id: Literal["cellpose_flow_dynamics"] = "cellpose_flow_dynamics"
1443 else:
1444 id: Literal["cellpose_flow_dynamics"]
1446 kwargs: CellposeFlowDynamicsKwargs
1449class CustomProcessingDescr(NodeWithExplicitlySetFields, FileDescr):
1450 """Custom (post)processing op — source file shipped inline with the model.
1452 Supports (post)processing that cannot be expressed by the built-in named
1453 operations (watershed, connected components, etc.)
1454 using a simple Python callable interface.
1456 The op is implemented in a ``.py`` file packaged alongside the model weights.
1457 Two styles are supported:
1459 *Callable class* — kwargs go to ``__init__``, tensors arrive in ``__call__``:
1461 .. code-block:: python
1463 # my_postprocess.py
1464 import numpy as np
1466 class my_postprocess:
1467 def __init__(self, threshold: float = 0.5) -> None:
1468 self.threshold = threshold
1469 def __call__(self, *arrays: np.ndarray) -> np.ndarray:
1470 # arrays = model output tensors in rdf.yaml declaration order
1471 return (arrays[0] > self.threshold).astype(np.uint8)
1473 *Factory function* — alternative closure style, identical runtime behaviour:
1475 .. code-block:: python
1477 # my_postprocess.py
1478 import numpy as np
1480 def my_postprocess(threshold: float = 0.5):
1481 def run(*arrays: np.ndarray) -> np.ndarray:
1482 return (arrays[0] > threshold).astype(np.uint8)
1483 return run
1485 Reference it in ``rdf.yaml`` with the source file included in the package:
1487 .. code-block:: yaml
1489 postprocessing:
1490 - id: custom
1491 callable: my_postprocess # class or function name in source
1492 source: my_postprocess.py # packaged alongside weights
1493 sha256: <hash> # sha256 of the source file
1494 kwargs: # forwarded to __init__ / factory
1495 threshold: 0.5
1497 **Security:** source files are SHA-256 verified before execution.
1498 Execution requires explicit opt-in in bioimageio.core and curator
1499 review before Zoo publication.
1500 """
1502 implemented_id: ClassVar[Literal["custom"]] = "custom"
1503 if TYPE_CHECKING:
1504 id: Literal["custom"] = "custom"
1505 else:
1506 id: Literal["custom"]
1508 callable: Annotated[
1509 str,
1510 Field(examples=["my_postprocess_factory", "MyPostprocessClass"]),
1511 ]
1512 """Name of the callable class or factory function defined in ``source``.
1514 At runtime: ``op = callable(**kwargs)``, then ``result = op(*output_tensors)``
1515 per image. Both a class with ``__call__`` and a factory function returning
1516 a callable satisfy this protocol."""
1518 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
1519 """Python source file (included when packaging the model)."""
1521 kwargs: dict[str, YamlValue] = Field(
1522 default_factory=cast(Callable[[], Dict[str, YamlValue]], dict)
1523 )
1524 """Keyword arguments forwarded to the callable (``__init__`` or factory)."""
1526 @model_serializer(mode="wrap", when_used="unless-none")
1527 def _serialize(
1528 self, nxt: SerializerFunctionWrapHandler, info: SerializationInfo
1529 ) -> dict[str, YamlValue]:
1530 return package_file_descr_serializer(self, nxt, info)
1533class FixedZeroMeanUnitVarianceKwargs(KwargsNode):
1534 """key word arguments for [FixedZeroMeanUnitVarianceDescr][]"""
1536 mean: float
1537 """The mean value to normalize with."""
1539 std: Annotated[float, Ge(1e-6)]
1540 """The standard deviation value to normalize with."""
1543class FixedZeroMeanUnitVarianceAlongAxisKwargs(KwargsNode):
1544 """key word arguments for [FixedZeroMeanUnitVarianceDescr][]"""
1546 mean: NotEmpty[list[float]]
1547 """The mean value(s) to normalize with."""
1549 std: NotEmpty[list[Annotated[float, Ge(1e-6)]]]
1550 """The standard deviation value(s) to normalize with.
1551 Size must match `mean` values."""
1553 axis: Annotated[NonBatchAxisId, Field(examples=["channel", "index"])]
1554 """The axis of the mean/std values to normalize each entry along that dimension
1555 separately."""
1557 @model_validator(mode="after")
1558 def _mean_and_std_match(self) -> Self:
1559 if len(self.mean) != len(self.std):
1560 raise ValueError(
1561 f"Size of `mean` ({len(self.mean)}) and `std` ({len(self.std)})"
1562 + " must match."
1563 )
1565 return self
1568class FixedZeroMeanUnitVarianceDescr(NodeWithExplicitlySetFields):
1569 """Subtract a given mean and divide by the standard deviation.
1571 Normalize with fixed, precomputed values for
1572 `FixedZeroMeanUnitVarianceKwargs.mean` and `FixedZeroMeanUnitVarianceKwargs.std`
1573 Use `FixedZeroMeanUnitVarianceAlongAxisKwargs` for independent scaling along given
1574 axes.
1576 Examples:
1577 1. scalar value for whole tensor
1578 - in YAML
1579 ```yaml
1580 preprocessing:
1581 - id: fixed_zero_mean_unit_variance
1582 kwargs:
1583 mean: 103.5
1584 std: 13.7
1585 ```
1586 - in Python
1587 >>> preprocessing = [FixedZeroMeanUnitVarianceDescr(
1588 ... kwargs=FixedZeroMeanUnitVarianceKwargs(mean=103.5, std=13.7)
1589 ... )]
1591 2. independently along an axis
1592 - in YAML
1593 ```yaml
1594 preprocessing:
1595 - id: fixed_zero_mean_unit_variance
1596 kwargs:
1597 axis: channel
1598 mean: [101.5, 102.5, 103.5]
1599 std: [11.7, 12.7, 13.7]
1600 ```
1601 - in Python
1602 >>> preprocessing = [FixedZeroMeanUnitVarianceDescr(
1603 ... kwargs=FixedZeroMeanUnitVarianceAlongAxisKwargs(
1604 ... axis=AxisId("channel"),
1605 ... mean=[101.5, 102.5, 103.5],
1606 ... std=[11.7, 12.7, 13.7],
1607 ... )
1608 ... )]
1609 """
1611 implemented_id: ClassVar[Literal["fixed_zero_mean_unit_variance"]] = (
1612 "fixed_zero_mean_unit_variance"
1613 )
1614 if TYPE_CHECKING:
1615 id: Literal["fixed_zero_mean_unit_variance"] = "fixed_zero_mean_unit_variance"
1616 else:
1617 id: Literal["fixed_zero_mean_unit_variance"]
1619 kwargs: FixedZeroMeanUnitVarianceKwargs | FixedZeroMeanUnitVarianceAlongAxisKwargs
1622class ZeroMeanUnitVarianceKwargs(KwargsNode):
1623 """key word arguments for [ZeroMeanUnitVarianceDescr][]"""
1625 axes: Annotated[Sequence[AxisId] | None, Field(examples=[("batch", "x", "y")])] = (
1626 None
1627 )
1628 """The subset of axes to normalize jointly, i.e. axes to reduce to compute mean/std.
1629 For example to normalize 'batch', 'x' and 'y' jointly in a tensor ('batch', 'channel', 'y', 'x')
1630 resulting in a tensor of equal shape normalized per channel, specify `axes=('batch', 'x', 'y')`.
1631 To normalize each sample independently leave out the 'batch' axis.
1632 Default: Scale all axes jointly."""
1634 eps: Annotated[float, Interval(gt=0, le=0.1)] = 1e-6
1635 """epsilon for numeric stability: `out = (tensor - mean) / (std + eps)`."""
1638class ZeroMeanUnitVarianceDescr(NodeWithExplicitlySetFields):
1639 """Subtract mean and divide by variance.
1641 Examples:
1642 Subtract tensor mean and variance
1643 - in YAML
1644 ```yaml
1645 preprocessing:
1646 - id: zero_mean_unit_variance
1647 ```
1648 - in Python
1649 >>> preprocessing = [ZeroMeanUnitVarianceDescr()]
1650 """
1652 implemented_id: ClassVar[Literal["zero_mean_unit_variance"]] = (
1653 "zero_mean_unit_variance"
1654 )
1655 if TYPE_CHECKING:
1656 id: Literal["zero_mean_unit_variance"] = "zero_mean_unit_variance"
1657 else:
1658 id: Literal["zero_mean_unit_variance"]
1660 kwargs: ZeroMeanUnitVarianceKwargs = Field(
1661 default_factory=ZeroMeanUnitVarianceKwargs.model_construct
1662 )
1665class ScaleRangeKwargs(KwargsNode):
1666 """key word arguments for [ScaleRangeDescr][]
1668 For `min_percentile`=0.0 (the default) and `max_percentile`=100 (the default)
1669 this processing step normalizes data to the [0, 1] intervall.
1670 For other percentiles the normalized values will partially be outside the [0, 1]
1671 intervall. Use `ScaleRange` followed by `ClipDescr` if you want to limit the
1672 normalized values to a range.
1673 """
1675 axes: Annotated[Sequence[AxisId] | None, Field(examples=[("batch", "x", "y")])] = (
1676 None
1677 )
1678 """The subset of axes to normalize jointly, i.e. axes to reduce to compute the min/max percentile value.
1679 For example to normalize 'batch', 'x' and 'y' jointly in a tensor ('batch', 'channel', 'y', 'x')
1680 resulting in a tensor of equal shape normalized per channel, specify `axes=('batch', 'x', 'y')`.
1681 To normalize samples independently, leave out the "batch" axis.
1682 Default: Scale all axes jointly."""
1684 min_percentile: Annotated[float, Interval(ge=0, lt=100)] = 0.0
1685 """The lower percentile used to determine the value to align with zero."""
1687 max_percentile: Annotated[float, Interval(gt=1, le=100)] = 100.0
1688 """The upper percentile used to determine the value to align with one.
1689 Has to be bigger than `min_percentile`.
1690 The range is 1 to 100 instead of 0 to 100 to avoid mistakenly
1691 accepting percentiles specified in the range 0.0 to 1.0."""
1693 eps: Annotated[float, Interval(gt=0, le=0.1)] = 1e-6
1694 """Epsilon for numeric stability.
1695 `out = (tensor - v_lower) / (v_upper - v_lower + eps)`;
1696 with `v_lower,v_upper` values at the respective percentiles."""
1698 reference_tensor: TensorId | None = None
1699 """ID of the unprocessed input tensor to compute the percentiles from.
1700 Default: The tensor itself.
1701 """
1703 @field_validator("max_percentile", mode="after")
1704 @classmethod
1705 def min_smaller_max(cls, value: float, info: ValidationInfo) -> float:
1706 if (min_p := info.data["min_percentile"]) >= value:
1707 raise ValueError(f"min_percentile {min_p} >= max_percentile {value}")
1709 return value
1712class ScaleRangeDescr(NodeWithExplicitlySetFields):
1713 """Scale with percentiles.
1715 Examples:
1716 1. Scale linearly to map 5th percentile to 0 and 99.8th percentile to 1.0
1717 - in YAML
1718 ```yaml
1719 preprocessing:
1720 - id: scale_range
1721 kwargs:
1722 axes: ['y', 'x']
1723 max_percentile: 99.8
1724 min_percentile: 5.0
1725 ```
1726 - in Python
1728 >>> preprocessing = [
1729 ... ScaleRangeDescr(
1730 ... kwargs=ScaleRangeKwargs(
1731 ... axes= (AxisId('y'), AxisId('x')),
1732 ... max_percentile= 99.8,
1733 ... min_percentile= 5.0,
1734 ... )
1735 ... )
1736 ... ]
1738 2. Combine the above scaling with additional clipping to clip values outside the range given by the percentiles.
1739 - in YAML
1740 ```yaml
1741 preprocessing:
1742 - id: scale_range
1743 kwargs:
1744 axes: ['y', 'x']
1745 max_percentile: 99.8
1746 min_percentile: 5.0
1747 - id: clip
1748 kwargs:
1749 min: 0.0
1750 max: 1.0
1751 ```
1752 - in Python
1754 >>> preprocessing = [
1755 ... ScaleRangeDescr(
1756 ... kwargs=ScaleRangeKwargs(
1757 ... axes= (AxisId('y'), AxisId('x')),
1758 ... max_percentile= 99.8,
1759 ... min_percentile= 5.0,
1760 ... )
1761 ... ),
1762 ... ClipDescr(
1763 ... kwargs=ClipKwargs(
1764 ... min=0.0,
1765 ... max=1.0,
1766 ... )
1767 ... ),
1768 ... ]
1770 """
1772 implemented_id: ClassVar[Literal["scale_range"]] = "scale_range"
1773 if TYPE_CHECKING:
1774 id: Literal["scale_range"] = "scale_range"
1775 else:
1776 id: Literal["scale_range"]
1777 kwargs: ScaleRangeKwargs = Field(default_factory=ScaleRangeKwargs.model_construct)
1780class ScaleMeanVarianceKwargs(KwargsNode):
1781 """key word arguments for [ScaleMeanVarianceKwargs][]"""
1783 reference_tensor: TensorId
1784 """ID of unprocessed input tensor to match."""
1786 axes: Annotated[Sequence[AxisId] | None, Field(examples=[("batch", "x", "y")])] = (
1787 None
1788 )
1789 """The subset of axes to normalize jointly, i.e. axes to reduce to compute mean/std.
1790 For example to normalize 'batch', 'x' and 'y' jointly in a tensor ('batch', 'channel', 'y', 'x')
1791 resulting in a tensor of equal shape normalized per channel, specify `axes=('batch', 'x', 'y')`.
1792 To normalize samples independently, leave out the 'batch' axis.
1793 Default: Scale all axes jointly."""
1795 eps: Annotated[float, Interval(gt=0, le=0.1)] = 1e-6
1796 """Epsilon for numeric stability:
1797 `out = (tensor - mean) / (std + eps) * (ref_std + eps) + ref_mean.`"""
1800class ScaleMeanVarianceDescr(NodeWithExplicitlySetFields):
1801 """Scale a tensor's data distribution to match another tensor's mean/std.
1802 `out = (tensor - mean) / (std + eps) * (ref_std + eps) + ref_mean.`
1803 """
1805 implemented_id: ClassVar[Literal["scale_mean_variance"]] = "scale_mean_variance"
1806 if TYPE_CHECKING:
1807 id: Literal["scale_mean_variance"] = "scale_mean_variance"
1808 else:
1809 id: Literal["scale_mean_variance"]
1810 kwargs: ScaleMeanVarianceKwargs
1813PreprocessingDescr = Annotated[
1814 Union[
1815 BinarizeDescr,
1816 ClipDescr,
1817 EnsureDtypeDescr,
1818 FixedZeroMeanUnitVarianceDescr,
1819 ScaleLinearDescr,
1820 ScaleRangeDescr,
1821 SigmoidDescr,
1822 SoftmaxDescr,
1823 ZeroMeanUnitVarianceDescr,
1824 ],
1825 Discriminator("id"),
1826]
1827PostprocessingDescr = Annotated[
1828 Union[
1829 BinarizeDescr,
1830 CellposeFlowDynamicsDescr,
1831 ClipDescr,
1832 CustomProcessingDescr,
1833 EnsureDtypeDescr,
1834 FixedZeroMeanUnitVarianceDescr,
1835 ScaleLinearDescr,
1836 ScaleMeanVarianceDescr,
1837 ScaleRangeDescr,
1838 SigmoidDescr,
1839 SoftmaxDescr,
1840 StardistPostprocessingDescr,
1841 ZeroMeanUnitVarianceDescr,
1842 ],
1843 Discriminator("id"),
1844]
1846IO_AxisT = TypeVar("IO_AxisT", InputAxis, OutputAxis)
1849class TensorDescrBase(Node, Generic[IO_AxisT]):
1850 id: TensorId
1851 """Tensor id. No duplicates are allowed."""
1853 description: Annotated[str, MaxLen(128)] = ""
1854 """free text description"""
1856 axes: NotEmpty[Sequence[IO_AxisT]]
1857 """tensor axes"""
1859 @property
1860 def shape(self):
1861 return tuple(a.size for a in self.axes)
1863 @field_validator("axes", mode="after", check_fields=False)
1864 @classmethod
1865 def _validate_axes(cls, axes: Sequence[AnyAxis]) -> Sequence[AnyAxis]:
1866 batch_axes = [a for a in axes if a.type == "batch"]
1867 if len(batch_axes) > 1:
1868 raise ValueError(
1869 f"Only one batch axis (per tensor) allowed, but got {batch_axes}"
1870 )
1872 seen_ids: set[AxisId] = set()
1873 duplicate_axes_ids: set[AxisId] = set()
1874 for a in axes:
1875 (duplicate_axes_ids if a.id in seen_ids else seen_ids).add(a.id)
1877 if duplicate_axes_ids:
1878 raise ValueError(f"Duplicate axis ids: {duplicate_axes_ids}")
1880 return axes
1882 test_tensor: FAIR[FileDescr_package | None] = None
1883 """An example tensor to use for testing.
1884 Using the model with the test input tensors is expected to yield the test output tensors.
1885 Each test tensor has be a an ndarray in the
1886 [numpy.lib file format](https://numpy.org/doc/stable/reference/generated/numpy.lib.format.html#module-numpy.lib.format).
1887 The file extension must be '.npy'."""
1889 sample_tensor: FAIR[FileDescr_package | None] = None
1890 """A sample tensor to illustrate a possible input/output for the model,
1891 The sample image primarily serves to inform a human user about an example use case
1892 and is typically stored as .hdf5, .png or .tiff.
1893 It has to be readable by the [imageio library](https://imageio.readthedocs.io/en/stable/formats/index.html#supported-formats)
1894 (numpy's `.npy` format is not supported).
1895 The image dimensionality has to match the number of axes specified in this tensor description.
1896 """
1898 @model_validator(mode="after")
1899 def _validate_sample_tensor(self) -> Self:
1900 if self.sample_tensor is None or not get_validation_context().perform_io_checks:
1901 return self
1903 reader = get_reader(self.sample_tensor.source, sha256=self.sample_tensor.sha256)
1904 tensor: NDArray[Any] = imread( # pyright: ignore[reportUnknownVariableType]
1905 reader.read(),
1906 extension=PurePosixPath(reader.original_file_name).suffix,
1907 )
1908 n_dims = len(tensor.squeeze().shape)
1909 n_dims_min = n_dims_max = len(self.axes)
1911 for a in self.axes:
1912 if isinstance(a, BatchAxis):
1913 n_dims_min -= 1
1914 elif isinstance(a.size, int):
1915 if a.size == 1:
1916 n_dims_min -= 1
1917 elif isinstance(a.size, (ParameterizedSize, DataDependentSize)):
1918 if a.size.min == 1:
1919 n_dims_min -= 1
1920 elif isinstance(a.size, SizeReference):
1921 if a.size.offset < 2:
1922 # size reference may result in singleton axis
1923 n_dims_min -= 1
1924 else:
1925 assert_never(a.size)
1927 n_dims_min = max(0, n_dims_min)
1928 if n_dims < n_dims_min or n_dims > n_dims_max:
1929 raise ValueError(
1930 f"Expected sample tensor to have {n_dims_min} to"
1931 + f" {n_dims_max} dimensions, but found {n_dims} (shape: {tensor.shape})."
1932 )
1934 return self
1936 data: TensorDataDescr | NotEmpty[Sequence[TensorDataDescr]] = (
1937 IntervalOrRatioDataDescr()
1938 )
1939 """Description of the tensor's data values, optionally per channel.
1940 If specified per channel, the data `type` needs to match across channels."""
1942 @property
1943 def dtype(
1944 self,
1945 ) -> Literal[
1946 "float32",
1947 "float64",
1948 "uint8",
1949 "int8",
1950 "uint16",
1951 "int16",
1952 "uint32",
1953 "int32",
1954 "uint64",
1955 "int64",
1956 "bool",
1957 ]:
1958 """dtype as specified under `data.type` or `data[i].type`"""
1959 if isinstance(self.data, collections.abc.Sequence):
1960 return self.data[0].type
1961 else:
1962 return self.data.type
1964 @field_validator("data", mode="after")
1965 @classmethod
1966 def _check_data_type_across_channels(
1967 cls, value: TensorDataDescr | NotEmpty[Sequence[TensorDataDescr]]
1968 ) -> TensorDataDescr | NotEmpty[Sequence[TensorDataDescr]]:
1969 if not isinstance(value, list):
1970 return value
1972 dtypes = {t.type for t in value}
1973 if len(dtypes) > 1:
1974 raise ValueError(
1975 "Tensor data descriptions per channel need to agree in their data"
1976 + f" `type`, but found {dtypes}."
1977 )
1979 return value
1981 @model_validator(mode="after")
1982 def _check_data_matches_channelaxis(self) -> Self:
1983 if not isinstance(self.data, (list, tuple)):
1984 return self
1986 for a in self.axes:
1987 if isinstance(a, ChannelAxis):
1988 size = a.size
1989 assert isinstance(size, int)
1990 break
1991 else:
1992 return self
1994 if len(self.data) != size:
1995 raise ValueError(
1996 f"Got tensor data descriptions for {len(self.data)} channels, but"
1997 + f" '{a.id}' axis has size {size}."
1998 )
2000 return self
2002 def get_axis_sizes_for_array(self, array: NDArray[Any]) -> dict[AxisId, int]:
2003 if len(array.shape) != len(self.axes):
2004 raise ValueError(
2005 f"Dimension mismatch: array shape {array.shape} (#{len(array.shape)})"
2006 + f" incompatible with {len(self.axes)} axes."
2007 )
2008 return {a.id: array.shape[i] for i, a in enumerate(self.axes)}
2011class ConstantPadding(Node):
2012 mode: Literal["constant"] = "constant"
2013 value: int | float = 0
2016class EdgePadding(Node):
2017 mode: Literal["edge"] = "edge"
2020class ReflectPadding(Node):
2021 mode: Literal["reflect"] = "reflect"
2024class SymmetricPadding(Node):
2025 mode: Literal["symmetric"] = "symmetric"
2028Padding = Union[ConstantPadding, EdgePadding, ReflectPadding, SymmetricPadding]
2031class ModelId(ResourceId):
2032 pass
2035class InputTensorDescr(TensorDescrBase[InputAxis]):
2036 id: TensorId = TensorId("input")
2037 """Input tensor id.
2038 No duplicates are allowed across all inputs and outputs."""
2040 output_of: ModelId | None = None
2041 """If this input tensor is the output of another model, specify the model id here.
2042 This model's input id must match the output id of the referenced model.
2043 """
2045 @model_validator(mode="after")
2046 def _validate_output_of(self) -> Self:
2047 if self.output_of is None:
2048 return self
2050 try:
2051 with get_validation_context().replace(perform_io_checks=False):
2052 opened_ref_model = open_bioimageio_yaml(self.output_of)
2053 format_version = opened_ref_model.content["format_version"]
2054 assert isinstance(format_version, str)
2055 if format_version.startswith("0.4"):
2056 ref_model = _ModelDescr04.model_validate(opened_ref_model.content)
2057 else:
2058 ref_model = ModelDescr.model_validate(opened_ref_model.content)
2059 except Exception as e:
2060 raise ValueError(
2061 f"Failed to load model '{self.output_of}' referenced under output_of: {e}"
2062 )
2064 try:
2065 ref_model_outputs = {
2066 t.id if isinstance(t, OutputTensorDescr) else TensorId(t.name)
2067 for t in ref_model.outputs
2068 }
2069 except Exception as e:
2070 raise ValueError(
2071 f"Failed to read output IDs of model '{self.output_of}' referenced under output_of: {e}"
2072 )
2074 if self.id not in ref_model_outputs:
2075 raise ValueError(
2076 f"Input tensor '{self.id}' is specified as output of model '{self.output_of}', "
2077 + f"but that model's outputs are {ref_model_outputs}."
2078 )
2079 return self
2081 optional: bool = False
2082 """indicates that this tensor may be `None`"""
2084 pad: Padding | None = None
2085 """Explicitly specify how to pad this input tensor.
2087 Use `axes[i].pad` to specify padding width.
2089 Note:
2090 Non-blockwise sample prediction only applies padding for axes with a `pad` specification.
2091 """
2093 preprocessing: list[PreprocessingDescr] = Field(
2094 default_factory=cast(Callable[[], List[PreprocessingDescr]], list)
2095 )
2096 """Description of how this input should be preprocessed.
2098 notes:
2099 - If preprocessing does not start with an 'ensure_dtype' entry, it is added
2100 to ensure an input tensor's data type matches the input tensor's data description.
2101 - If preprocessing does not end with an 'ensure_dtype' or 'binarize' entry, an
2102 'ensure_dtype' step is added to ensure preprocessing steps are not unintentionally
2103 changing the data type.
2104 """
2106 @model_validator(mode="after")
2107 def _validate_preprocessing_kwargs(self) -> Self:
2108 axes_ids = [a.id for a in self.axes]
2109 for p in self.preprocessing:
2110 kwargs_axes: Sequence[Any] | None = p.kwargs.get("axes")
2111 if kwargs_axes is None:
2112 continue
2114 if not isinstance(kwargs_axes, collections.abc.Sequence):
2115 raise ValueError(
2116 f"Expected `preprocessing.i.kwargs.axes` to be a sequence, but got {type(kwargs_axes)}"
2117 )
2119 if any(a not in axes_ids for a in kwargs_axes):
2120 raise ValueError(
2121 "`preprocessing.i.kwargs.axes` needs to be subset of axes ids"
2122 )
2124 if isinstance(self.data, (NominalOrOrdinalDataDescr, IntervalOrRatioDataDescr)):
2125 dtype = self.data.type
2126 else:
2127 dtype = self.data[0].type
2129 # ensure `preprocessing` begins with `EnsureDtypeDescr`
2130 if not self.preprocessing or not isinstance(
2131 self.preprocessing[0], EnsureDtypeDescr
2132 ):
2133 self.preprocessing.insert(
2134 0, EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype=dtype))
2135 )
2137 # ensure `preprocessing` ends with `EnsureDtypeDescr` or `BinarizeDescr`
2138 if not isinstance(self.preprocessing[-1], (EnsureDtypeDescr, BinarizeDescr)):
2139 self.preprocessing.append(
2140 EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype=dtype))
2141 )
2143 return self
2146def convert_axes(
2147 axes: str,
2148 *,
2149 shape: Sequence[int] | _ParameterizedInputShape_v0_4 | _ImplicitOutputShape_v0_4,
2150 tensor_type: Literal["input", "output"],
2151 halo: Sequence[int] | None,
2152 size_refs: Mapping[_TensorName_v0_4, Mapping[str, int]],
2153):
2154 ret: list[AnyAxis] = []
2155 for i, a in enumerate(axes):
2156 axis_type = _AXIS_TYPE_MAP.get(a, a)
2157 if axis_type == "batch":
2158 ret.append(BatchAxis())
2159 continue
2161 scale = 1.0
2162 if isinstance(shape, _ParameterizedInputShape_v0_4):
2163 if shape.step[i] == 0:
2164 size = shape.min[i]
2165 else:
2166 size = ParameterizedSize(min=shape.min[i], step=shape.step[i])
2167 elif isinstance(shape, _ImplicitOutputShape_v0_4):
2168 ref_t = str(shape.reference_tensor)
2169 if ref_t.count(".") == 1:
2170 t_id, orig_a_id = ref_t.split(".")
2171 else:
2172 t_id = ref_t
2173 orig_a_id = a
2175 a_id = _AXIS_ID_MAP.get(orig_a_id, a)
2176 if not (orig_scale := shape.scale[i]):
2177 # old way to insert a new axis dimension
2178 size = int(2 * shape.offset[i])
2179 else:
2180 scale = 1 / orig_scale
2181 if axis_type in ("channel", "index"):
2182 # these axes no longer have a scale
2183 offset_from_scale = orig_scale * size_refs.get(
2184 _TensorName_v0_4(t_id), {}
2185 ).get(orig_a_id, 0)
2186 else:
2187 offset_from_scale = 0
2188 size = SizeReference(
2189 tensor_id=TensorId(t_id),
2190 axis_id=AxisId(a_id),
2191 offset=int(offset_from_scale + 2 * shape.offset[i]),
2192 )
2193 else:
2194 size = shape[i]
2196 if axis_type == "time":
2197 if tensor_type == "input":
2198 ret.append(TimeInputAxis(size=size, scale=scale))
2199 else:
2200 assert not isinstance(size, ParameterizedSize)
2201 if halo is None:
2202 ret.append(TimeOutputAxis(size=size, scale=scale))
2203 else:
2204 assert not isinstance(size, int)
2205 ret.append(
2206 TimeOutputAxisWithHalo(size=size, scale=scale, halo=halo[i])
2207 )
2209 elif axis_type == "index":
2210 if tensor_type == "input":
2211 ret.append(IndexInputAxis(size=size))
2212 else:
2213 if isinstance(size, ParameterizedSize):
2214 size = DataDependentSize(min=size.min)
2216 ret.append(IndexOutputAxis(size=size))
2217 elif axis_type == "channel":
2218 assert not isinstance(size, ParameterizedSize)
2219 if isinstance(size, SizeReference):
2220 warnings.warn(
2221 "Conversion of channel size from an implicit output shape may be"
2222 + " wrong"
2223 )
2224 ret.append(
2225 ChannelAxis(
2226 channel_names=[f"channel{i}" for i in range(size.offset)]
2227 )
2228 )
2229 else:
2230 ret.append(
2231 ChannelAxis(channel_names=[f"channel{i}" for i in range(size)])
2232 )
2233 elif axis_type == "space":
2234 if tensor_type == "input":
2235 ret.append(SpaceInputAxis(id=AxisId(a), size=size, scale=scale))
2236 else:
2237 assert not isinstance(size, ParameterizedSize)
2238 if halo is None or halo[i] == 0:
2239 ret.append(SpaceOutputAxis(id=AxisId(a), size=size, scale=scale))
2240 elif isinstance(size, int):
2241 raise NotImplementedError(
2242 f"output axis with halo and fixed size (here {size}) not allowed"
2243 )
2244 else:
2245 ret.append(
2246 SpaceOutputAxisWithHalo(
2247 id=AxisId(a), size=size, scale=scale, halo=halo[i]
2248 )
2249 )
2251 return ret
2254def _axes_letters_to_ids(
2255 axes: str | None,
2256) -> list[AxisId] | None:
2257 if axes is None:
2258 return None
2260 return [AxisId(a) for a in axes]
2263def _get_complement_v04_axis(
2264 tensor_axes: Sequence[str], axes: Sequence[str] | None
2265) -> AxisId | None:
2266 if axes is None:
2267 return None
2269 non_complement_axes = set(axes) | {"b"}
2270 complement_axes = [a for a in tensor_axes if a not in non_complement_axes]
2271 if len(complement_axes) > 1:
2272 raise ValueError(
2273 f"Expected none or a single complement axis, but axes '{axes}' "
2274 + f"for tensor dims '{tensor_axes}' leave '{complement_axes}'."
2275 )
2277 return None if not complement_axes else AxisId(complement_axes[0])
2280def _convert_proc(
2281 p: _PreprocessingDescr_v0_4 | _PostprocessingDescr_v0_4,
2282 tensor_axes: Sequence[str],
2283) -> PreprocessingDescr | PostprocessingDescr:
2284 if isinstance(p, _BinarizeDescr_v0_4):
2285 return BinarizeDescr(kwargs=BinarizeKwargs(threshold=p.kwargs.threshold))
2286 elif isinstance(p, _ClipDescr_v0_4):
2287 return ClipDescr(kwargs=ClipKwargs(min=p.kwargs.min, max=p.kwargs.max))
2288 elif isinstance(p, _SigmoidDescr_v0_4):
2289 return SigmoidDescr()
2290 elif isinstance(p, _ScaleLinearDescr_v0_4):
2291 axes = _axes_letters_to_ids(p.kwargs.axes)
2292 if p.kwargs.axes is None:
2293 axis = None
2294 else:
2295 axis = _get_complement_v04_axis(tensor_axes, p.kwargs.axes)
2297 if axis is None:
2298 assert not isinstance(p.kwargs.gain, list)
2299 assert not isinstance(p.kwargs.offset, list)
2300 kwargs = ScaleLinearKwargs(gain=p.kwargs.gain, offset=p.kwargs.offset)
2301 else:
2302 kwargs = ScaleLinearAlongAxisKwargs(
2303 axis=axis, gain=p.kwargs.gain, offset=p.kwargs.offset
2304 )
2305 return ScaleLinearDescr(kwargs=kwargs)
2306 elif isinstance(p, _ScaleMeanVarianceDescr_v0_4):
2307 return ScaleMeanVarianceDescr(
2308 kwargs=ScaleMeanVarianceKwargs(
2309 axes=_axes_letters_to_ids(p.kwargs.axes),
2310 reference_tensor=TensorId(str(p.kwargs.reference_tensor)),
2311 eps=p.kwargs.eps,
2312 )
2313 )
2314 elif isinstance(p, _ZeroMeanUnitVarianceDescr_v0_4):
2315 if p.kwargs.mode == "fixed":
2316 mean = p.kwargs.mean
2317 std = p.kwargs.std
2318 assert mean is not None
2319 assert std is not None
2321 axis = _get_complement_v04_axis(tensor_axes, p.kwargs.axes)
2323 if axis is None:
2324 if isinstance(mean, list):
2325 raise ValueError("Expected single float value for mean, not <list>")
2326 if isinstance(std, list):
2327 raise ValueError("Expected single float value for std, not <list>")
2328 return FixedZeroMeanUnitVarianceDescr(
2329 kwargs=FixedZeroMeanUnitVarianceKwargs.model_construct(
2330 mean=mean,
2331 std=std,
2332 )
2333 )
2334 else:
2335 if not isinstance(mean, list):
2336 mean = [float(mean)]
2337 if not isinstance(std, list):
2338 std = [float(std)]
2340 return FixedZeroMeanUnitVarianceDescr(
2341 kwargs=FixedZeroMeanUnitVarianceAlongAxisKwargs(
2342 axis=axis, mean=mean, std=std
2343 )
2344 )
2346 else:
2347 axes = _axes_letters_to_ids(p.kwargs.axes) or []
2348 if p.kwargs.mode == "per_dataset":
2349 axes = [AxisId("batch")] + axes
2350 if not axes:
2351 axes = None
2352 return ZeroMeanUnitVarianceDescr(
2353 kwargs=ZeroMeanUnitVarianceKwargs(axes=axes, eps=p.kwargs.eps)
2354 )
2356 elif isinstance(p, _ScaleRangeDescr_v0_4):
2357 return ScaleRangeDescr(
2358 kwargs=ScaleRangeKwargs(
2359 axes=_axes_letters_to_ids(p.kwargs.axes),
2360 min_percentile=p.kwargs.min_percentile,
2361 max_percentile=p.kwargs.max_percentile,
2362 eps=p.kwargs.eps,
2363 )
2364 )
2365 else:
2366 assert_never(p)
2369class _InputTensorConv(
2370 Converter[
2371 _InputTensorDescr_v0_4,
2372 InputTensorDescr,
2373 FileSource,
2374 Optional[FileSource],
2375 Mapping[_TensorName_v0_4, Mapping[str, int]],
2376 ]
2377):
2378 def _convert(
2379 self,
2380 src: _InputTensorDescr_v0_4,
2381 tgt: type[InputTensorDescr | dict[str, Any]],
2382 test_tensor: FileSource,
2383 sample_tensor: FileSource | None,
2384 size_refs: Mapping[_TensorName_v0_4, Mapping[str, int]],
2385 ) -> InputTensorDescr | dict[str, Any]:
2386 axes: list[InputAxis] = convert_axes( # pyright: ignore[reportAssignmentType]
2387 src.axes,
2388 shape=src.shape,
2389 tensor_type="input",
2390 halo=None,
2391 size_refs=size_refs,
2392 )
2393 prep: list[PreprocessingDescr] = []
2394 for p in src.preprocessing:
2395 cp = _convert_proc(p, src.axes)
2396 assert not isinstance(
2397 cp,
2398 (
2399 CellposeFlowDynamicsDescr,
2400 CustomProcessingDescr,
2401 ScaleMeanVarianceDescr,
2402 StardistPostprocessingDescr,
2403 ),
2404 )
2405 prep.append(cp)
2407 prep.append(EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype="float32")))
2409 return tgt(
2410 axes=axes,
2411 id=TensorId(str(src.name)),
2412 test_tensor=FileDescr(source=test_tensor),
2413 sample_tensor=(
2414 None if sample_tensor is None else FileDescr(source=sample_tensor)
2415 ),
2416 data={"type": src.data_type}, # pyright: ignore[reportArgumentType]
2417 preprocessing=prep,
2418 )
2421_input_tensor_conv = _InputTensorConv(_InputTensorDescr_v0_4, InputTensorDescr)
2424class OutputTensorDescr(TensorDescrBase[OutputAxis]):
2425 id: TensorId = TensorId("output")
2426 """Output tensor id.
2427 No duplicates are allowed across all inputs and outputs."""
2429 postprocessing: list[PostprocessingDescr] = Field(
2430 default_factory=cast(Callable[[], List[PostprocessingDescr]], list)
2431 )
2432 """Description of how this output should be postprocessed.
2434 note: `postprocessing` always ends with an 'ensure_dtype' operation.
2435 If not given this is added to cast to this tensor's `data.type`.
2436 """
2438 @model_validator(mode="after")
2439 def _validate_postprocessing_kwargs(self) -> Self:
2440 axes_ids = [a.id for a in self.axes]
2441 for p in self.postprocessing:
2442 kwargs_axes = p.kwargs.get("axes")
2443 if kwargs_axes is None:
2444 continue
2446 if not isinstance(kwargs_axes, collections.abc.Sequence):
2447 raise ValueError(
2448 f"expected `axes` sequence, but got {type(kwargs_axes)}"
2449 )
2451 kwargs_axes_seq: Sequence[Any] = cast(Sequence[Any], kwargs_axes)
2452 if any(a not in axes_ids for a in kwargs_axes_seq):
2453 raise ValueError("`kwargs.axes` needs to be subset of axes ids")
2455 if isinstance(self.data, (NominalOrOrdinalDataDescr, IntervalOrRatioDataDescr)):
2456 dtype = self.data.type
2457 else:
2458 dtype = self.data[0].type
2460 # ensure `postprocessing` ends with `EnsureDtypeDescr` or `BinarizeDescr`
2461 if not self.postprocessing or not isinstance(
2462 self.postprocessing[-1], (EnsureDtypeDescr, BinarizeDescr)
2463 ):
2464 self.postprocessing.append(
2465 EnsureDtypeDescr(kwargs=EnsureDtypeKwargs(dtype=dtype))
2466 )
2467 return self
2470class _OutputTensorConv(
2471 Converter[
2472 _OutputTensorDescr_v0_4,
2473 OutputTensorDescr,
2474 FileSource,
2475 Optional[FileSource],
2476 Mapping[_TensorName_v0_4, Mapping[str, int]],
2477 ]
2478):
2479 def _convert(
2480 self,
2481 src: _OutputTensorDescr_v0_4,
2482 tgt: type[OutputTensorDescr | dict[str, Any]],
2483 test_tensor: FileSource,
2484 sample_tensor: FileSource | None,
2485 size_refs: Mapping[_TensorName_v0_4, Mapping[str, int]],
2486 ) -> OutputTensorDescr | dict[str, Any]:
2487 # TODO: split convert_axes into convert_output_axes and convert_input_axes
2488 axes: list[OutputAxis] = convert_axes( # pyright: ignore[reportAssignmentType]
2489 src.axes,
2490 shape=src.shape,
2491 tensor_type="output",
2492 halo=src.halo,
2493 size_refs=size_refs,
2494 )
2495 data_descr: dict[str, Any] = {"type": src.data_type}
2496 if data_descr["type"] == "bool":
2497 data_descr["values"] = [False, True]
2499 return tgt(
2500 axes=axes,
2501 id=TensorId(str(src.name)),
2502 test_tensor=FileDescr(source=test_tensor),
2503 sample_tensor=(
2504 None if sample_tensor is None else FileDescr(source=sample_tensor)
2505 ),
2506 data=data_descr, # pyright: ignore[reportArgumentType]
2507 postprocessing=[_convert_proc(p, src.axes) for p in src.postprocessing],
2508 )
2511_output_tensor_conv = _OutputTensorConv(_OutputTensorDescr_v0_4, OutputTensorDescr)
2514TensorDescr = Union[InputTensorDescr, OutputTensorDescr]
2517def get_halos(
2518 tensors: Mapping[TensorId, TensorDescr],
2519 /,
2520) -> dict[TensorId, dict[AxisId, tuple[int, int]]]:
2521 """Get all input and output halos from tensor descriptions.
2523 Note:
2524 - Input halos are to be padded
2525 - Output halos are to be cropped
2526 """
2527 halos: dict[TensorId, dict[AxisId, tuple[int, int]]] = {}
2528 for descr in tensors.values():
2529 if isinstance(descr, InputTensorDescr):
2530 continue
2531 for axis in descr.axes:
2532 if not isinstance(axis, WithHalo):
2533 continue
2535 ref_scale = next(
2536 a
2537 for a in tensors[axis.size.tensor_id].axes
2538 if a.id == axis.size.axis_id
2539 ).scale
2541 # set output halo (to be cropped)
2542 halos.setdefault(descr.id, {})[axis.id] = (axis.halo, axis.halo)
2543 # set input halo (to be padded)
2544 pad_width = int(axis.halo / axis.scale * ref_scale)
2545 halos.setdefault(axis.size.tensor_id, {})[axis.size.axis_id] = (
2546 pad_width,
2547 pad_width,
2548 )
2550 return halos
2553def validate_tensors(
2554 tensors: Mapping[TensorId, tuple[TensorDescr, NDArray[Any] | None]],
2555 tensor_origin: Literal[
2556 "source", "test_tensor"
2557 ] = "source", # for more precise error messages
2558 *,
2559 pad_inputs: bool | Literal["allow"] = True,
2560 crop_outputs: bool | Literal["allow"] = True,
2561):
2562 """Validate all inputs (and optionally output tensors) against their tensor descriptions.
2564 Args:
2565 tensors: Mapping of tensor id to a tuple of tensor description and optional numpy array.
2566 tensor_origin: String to use in error messages to indicate the origin of the tensors being validated.
2567 pad_inputs: Wether to apply/allow padding of inputs before shape comparison
2568 crop_outputs: Wether to apply/allow cropping of outputs before shape comparison.
2569 """
2570 all_tensor_axes: dict[TensorId, dict[AxisId, tuple[AnyAxis, int | None]]] = {}
2572 def e_msg_location(d: TensorDescr):
2573 return f"{'inputs' if isinstance(d, InputTensorDescr) else 'outputs'}[{d.id}]"
2575 for descr, array in tensors.values():
2576 if array is None:
2577 axis_sizes = {a.id: None for a in descr.axes}
2578 else:
2579 try:
2580 axis_sizes = descr.get_axis_sizes_for_array(array)
2581 except ValueError as e:
2582 raise ValueError(f"{e_msg_location(descr)} {e}")
2584 all_tensor_axes[descr.id] = {a.id: (a, axis_sizes[a.id]) for a in descr.axes}
2586 # get halos to be padded/cropped to validate against halo-adjusted sizes
2587 io_halos = get_halos({k: v[0] for k, v in tensors.items()})
2589 for descr, array in tensors.values():
2590 if array is None:
2591 continue
2593 if descr.dtype in ("float32", "float64"):
2594 invalid_test_tensor_dtype = array.dtype.name not in (
2595 "float32",
2596 "float64",
2597 "uint8",
2598 "int8",
2599 "uint16",
2600 "int16",
2601 "uint32",
2602 "int32",
2603 "uint64",
2604 "int64",
2605 )
2606 else:
2607 invalid_test_tensor_dtype = array.dtype.name != descr.dtype
2609 if invalid_test_tensor_dtype:
2610 raise ValueError(
2611 f"{tensor_origin} data type '{array.dtype.name}' does not"
2612 + f" match described {e_msg_location(descr)}.dtype '{descr.dtype}'"
2613 )
2615 if array.min() > -1e-4 and array.max() < 1e-4:
2616 raise ValueError(
2617 "Output values are too small for reliable testing."
2618 + f" Values <-1e5 or >=1e5 must be present in {tensor_origin}"
2619 )
2621 for a in descr.axes:
2622 actual_size = all_tensor_axes[descr.id][a.id][1]
2624 if actual_size is None:
2625 continue
2627 if a.size is None:
2628 continue
2630 # add padding width to actual tensor size
2631 total_axis_halo = sum(io_halos.get(descr.id, {}).get(a.id, (0, 0)))
2632 if isinstance(descr, InputTensorDescr):
2633 # pad input halos
2634 actual_size_with_halo = actual_size + total_axis_halo
2635 if pad_inputs is True:
2636 check_sizes = {actual_size_with_halo}
2637 size_hint = " (after padding input halo)"
2638 elif pad_inputs == "allow":
2639 check_sizes = {actual_size, actual_size_with_halo}
2640 size_hint = " (with or without padding input halo)"
2641 elif pad_inputs is False:
2642 check_sizes = {actual_size}
2643 size_hint = ""
2644 else:
2645 assert_never(pad_inputs)
2647 elif isinstance(descr, OutputTensorDescr):
2648 # crop output halos
2649 actual_size_with_halo = max(0, actual_size - total_axis_halo)
2650 if crop_outputs is True:
2651 check_sizes = {actual_size_with_halo}
2652 size_hint = " (after cropping output halo)"
2653 elif crop_outputs == "allow":
2654 check_sizes = {actual_size, actual_size_with_halo}
2655 size_hint = " (with or without cropping output halo)"
2656 elif crop_outputs is False:
2657 check_sizes = {actual_size}
2658 size_hint = ""
2659 else:
2660 assert_never(crop_outputs)
2661 else:
2662 assert_never(descr)
2664 del actual_size # make sure we explicitly use unchanged or halo-adjusted size from here on
2666 if isinstance(a.size, int):
2667 if a.size not in check_sizes:
2668 raise ValueError(
2669 f"{e_msg_location(descr)}.axes[{a.id}]: {tensor_origin} axis "
2670 + f"has incompatible size {check_sizes}{size_hint}, expected {a.size}"
2671 )
2672 elif isinstance(a.size, (ParameterizedSize, DataDependentSize)):
2673 _ = try_all_raise_last(
2674 (partial(a.size.validate_size, s) for s in check_sizes),
2675 f"{e_msg_location(descr)}.axes[{a.id}]: {tensor_origin} axis ",
2676 )
2677 elif isinstance(a.size, SizeReference):
2678 ref_tensor_axes = all_tensor_axes.get(a.size.tensor_id)
2679 if ref_tensor_axes is None:
2680 raise ValueError(
2681 f"{e_msg_location(descr)}.axes[{a.id}].size.tensor_id: Unknown tensor"
2682 + f" reference '{a.size.tensor_id}', available: {list(all_tensor_axes)}"
2683 )
2685 ref_axis, ref_size = ref_tensor_axes.get(a.size.axis_id, (None, None))
2686 if ref_axis is None or ref_size is None:
2687 raise ValueError(
2688 f"{e_msg_location(descr)}.axes[{a.id}].size.axis_id: Unknown tensor axis"
2689 + f" reference '{a.size.tensor_id}.{a.size.axis_id}, available: {list(ref_tensor_axes)}"
2690 )
2692 if a.unit != ref_axis.unit:
2693 raise ValueError(
2694 f"{e_msg_location(descr)}.axes[{a.id}].size: `SizeReference` requires"
2695 + " axis and reference axis to have the same `unit`, but"
2696 + f" {a.unit}!={ref_axis.unit}"
2697 )
2699 if (
2700 expected_size := (
2701 ref_size * ref_axis.scale / a.scale + a.size.offset
2702 )
2703 ) not in check_sizes:
2704 raise ValueError(
2705 f"{e_msg_location(descr)}.{tensor_origin}: axis '{a.id}' of size"
2706 + f" {check_sizes} invalid for referenced size {ref_size};"
2707 + f" expected {expected_size}"
2708 )
2709 else:
2710 assert_never(a.size)
2713FileDescr_dependencies = Annotated[
2714 FileDescr_package,
2715 WithSuffix((".yaml", ".yml"), case_sensitive=True),
2716 Field(examples=[{"source": "environment.yaml"}]),
2717]
2720class _ArchitectureCallableDescr(Node):
2721 callable: Annotated[Identifier, Field(examples=["MyNetworkClass", "get_my_model"])]
2722 """Identifier of the callable that returns a torch.nn.Module instance."""
2724 kwargs: dict[str, YamlValue] = Field(
2725 default_factory=cast(Callable[[], Dict[str, YamlValue]], dict)
2726 )
2727 """key word arguments for the `callable`"""
2730class ArchitectureFromFileDescr(_ArchitectureCallableDescr, FileDescr):
2731 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
2732 """Architecture source file"""
2734 @model_serializer(mode="wrap", when_used="unless-none")
2735 def _serialize(self, nxt: SerializerFunctionWrapHandler, info: SerializationInfo):
2736 return package_file_descr_serializer(self, nxt, info)
2739class ArchitectureFromLibraryDescr(_ArchitectureCallableDescr):
2740 import_from: str
2741 """Where to import the callable from, i.e. `from <import_from> import <callable>`"""
2744class _ArchFileConv(
2745 Converter[
2746 _CallableFromFile_v0_4,
2747 ArchitectureFromFileDescr,
2748 Optional[Sha256],
2749 Dict[str, Any],
2750 ]
2751):
2752 def _convert(
2753 self,
2754 src: _CallableFromFile_v0_4,
2755 tgt: type[ArchitectureFromFileDescr | dict[str, Any]],
2756 sha256: Sha256 | None,
2757 kwargs: dict[str, Any],
2758 ) -> ArchitectureFromFileDescr | dict[str, Any]:
2759 if src.startswith("http") and src.count(":") == 2:
2760 http, source, callable_ = src.split(":")
2761 source = f"{http}:{source}"
2762 elif not src.startswith("http") and src.count(":") == 1:
2763 source, callable_ = src.split(":")
2764 else:
2765 source = str(src)
2766 callable_ = str(src)
2767 return tgt(
2768 callable=Identifier(callable_),
2769 source=cast(FileSource, source),
2770 sha256=sha256,
2771 kwargs=kwargs,
2772 )
2775_arch_file_conv = _ArchFileConv(_CallableFromFile_v0_4, ArchitectureFromFileDescr)
2778class _ArchLibConv(
2779 Converter[
2780 _CallableFromDepencency_v0_4, ArchitectureFromLibraryDescr, Dict[str, Any]
2781 ]
2782):
2783 def _convert(
2784 self,
2785 src: _CallableFromDepencency_v0_4,
2786 tgt: type[ArchitectureFromLibraryDescr | dict[str, Any]],
2787 kwargs: dict[str, Any],
2788 ) -> ArchitectureFromLibraryDescr | dict[str, Any]:
2789 *mods, callable_ = src.split(".")
2790 import_from = ".".join(mods)
2791 return tgt(
2792 import_from=import_from, callable=Identifier(callable_), kwargs=kwargs
2793 )
2796_arch_lib_conv = _ArchLibConv(
2797 _CallableFromDepencency_v0_4, ArchitectureFromLibraryDescr
2798)
2801class WeightsEntryDescrBase(FileDescr):
2802 type: ClassVar[WeightsFormat]
2803 weights_format_name: ClassVar[str] # human readable
2805 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
2806 """Source of the weights file."""
2808 authors: list[Author] | None = None
2809 """Authors
2810 Either the person(s) that have trained this model resulting in the original weights file.
2811 (If this is the initial weights entry, i.e. it does not have a `parent`)
2812 Or the person(s) who have converted the weights to this weights format.
2813 (If this is a child weight, i.e. it has a `parent` field)
2814 """
2816 parent: Annotated[WeightsFormat | None, Field(examples=["pytorch_state_dict"])] = (
2817 None
2818 )
2819 """The source weights these weights were converted from.
2820 For example, if a model's weights were converted from the `pytorch_state_dict` format to `torchscript`,
2821 The `pytorch_state_dict` weights entry has no `parent` and is the parent of the `torchscript` weights.
2822 All weight entries except one (the initial set of weights resulting from training the model),
2823 need to have this field."""
2825 comment: str = ""
2826 """A comment about this weights entry, for example how these weights were created."""
2828 @model_validator(mode="after")
2829 def _validate(self) -> Self:
2830 if self.type == self.parent:
2831 raise ValueError("Weights entry can't be it's own parent.")
2833 return self
2835 @model_serializer(mode="wrap", when_used="unless-none")
2836 def _serialize(self, nxt: SerializerFunctionWrapHandler, info: SerializationInfo):
2837 return package_file_descr_serializer(self, nxt, info)
2840class KerasHdf5WeightsDescr(WeightsEntryDescrBase):
2841 type: ClassVar[WeightsFormat] = "keras_hdf5"
2842 weights_format_name: ClassVar[str] = "Keras HDF5"
2843 tensorflow_version: Version
2844 """TensorFlow version used to create these weights."""
2847class KerasV3WeightsDescr(WeightsEntryDescrBase):
2848 type: ClassVar[WeightsFormat] = "keras_v3"
2849 weights_format_name: ClassVar[str] = "Keras v3"
2850 keras_version: Annotated[Version, Ge(Version(3))]
2851 """Keras version used to create these weights."""
2852 backend: tuple[Literal["tensorflow", "jax", "torch"], Version]
2853 """Keras backend used to create these weights."""
2854 source: Annotated[
2855 FileSource,
2856 AfterValidator(wo_special_file_name),
2857 WithSuffix(".keras", case_sensitive=True),
2858 ]
2859 """Source of the .keras weights file."""
2862FileDescr_external_data = Annotated[
2863 FileDescr_package,
2864 WithSuffix(".data", case_sensitive=True),
2865 Field(examples=[{"source": "weights.onnx.data"}]),
2866]
2869class OnnxWeightsDescr(WeightsEntryDescrBase):
2870 type: ClassVar[WeightsFormat] = "onnx"
2871 weights_format_name: ClassVar[str] = "ONNX"
2872 opset_version: Annotated[int, Ge(7)]
2873 """ONNX opset version"""
2875 external_data: FileDescr_external_data | None = None
2876 """Source of the external ONNX data file holding the weights.
2877 (If present **source** holds the ONNX architecture without weights)."""
2879 @model_validator(mode="after")
2880 def _validate_external_data_unique_file_name(self) -> Self:
2881 if self.external_data is not None and (
2882 extract_file_name(self.source)
2883 == extract_file_name(self.external_data.source)
2884 ):
2885 raise ValueError(
2886 f"ONNX `external_data` file name '{extract_file_name(self.external_data.source)}'"
2887 + " must be different from ONNX `source` file name."
2888 )
2890 return self
2893class PytorchStateDictWeightsDescr(WeightsEntryDescrBase):
2894 type: ClassVar[WeightsFormat] = "pytorch_state_dict"
2895 weights_format_name: ClassVar[str] = "Pytorch State Dict"
2896 architecture: ArchitectureFromFileDescr | ArchitectureFromLibraryDescr
2897 pytorch_version: Version
2898 """Version of the PyTorch library used.
2899 If `architecture.depencencies` is specified it has to include pytorch and any version pinning has to be compatible.
2900 """
2901 dependencies: FileDescr_dependencies | None = None
2902 """Custom depencies beyond pytorch described in a Conda environment file.
2903 Allows to specify custom dependencies, see conda docs:
2904 - [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)
2905 - [Creating an environment file manually](https://conda.io/projects/conda/en/latest/user-guide/tasks/manage-environments.html#creating-an-environment-file-manually)
2907 The conda environment file should include pytorch and any version pinning has to be compatible with
2908 **pytorch_version**.
2909 """
2910 strict: bool = True
2911 """Whether to allow missing or unexpected keys or to be strict about the architecture matching the state dict weights."""
2914class TensorflowJsWeightsDescr(WeightsEntryDescrBase):
2915 type: ClassVar[WeightsFormat] = "tensorflow_js"
2916 weights_format_name: ClassVar[str] = "Tensorflow.js"
2917 tensorflow_version: Version
2918 """Version of the TensorFlow library used."""
2920 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
2921 """The multi-file weights.
2922 All required files/folders should be a zip archive."""
2925class TensorflowSavedModelBundleWeightsDescr(WeightsEntryDescrBase):
2926 type: ClassVar[WeightsFormat] = "tensorflow_saved_model_bundle"
2927 weights_format_name: ClassVar[str] = "Tensorflow Saved Model"
2928 tensorflow_version: Version
2929 """Version of the TensorFlow library used."""
2931 dependencies: FileDescr_dependencies | None = None
2932 """Custom dependencies beyond tensorflow.
2933 Should include tensorflow and any version pinning has to be compatible with **tensorflow_version**."""
2935 source: Annotated[FileSource, AfterValidator(wo_special_file_name)]
2936 """The multi-file weights.
2937 All required files/folders should be a zip archive."""
2940class TorchscriptWeightsDescr(WeightsEntryDescrBase):
2941 type: ClassVar[WeightsFormat] = "torchscript"
2942 weights_format_name: ClassVar[str] = "TorchScript"
2943 pytorch_version: Version
2944 """Version of the PyTorch library used."""
2947SpecificWeightsDescr = Union[
2948 KerasHdf5WeightsDescr,
2949 KerasV3WeightsDescr,
2950 OnnxWeightsDescr,
2951 PytorchStateDictWeightsDescr,
2952 TensorflowJsWeightsDescr,
2953 TensorflowSavedModelBundleWeightsDescr,
2954 TorchscriptWeightsDescr,
2955]
2958class WeightsDescr(Node):
2959 keras_hdf5: KerasHdf5WeightsDescr | None = None
2960 keras_v3: KerasV3WeightsDescr | None = None
2961 onnx: OnnxWeightsDescr | None = None
2962 pytorch_state_dict: PytorchStateDictWeightsDescr | None = None
2963 tensorflow_js: TensorflowJsWeightsDescr | None = None
2964 tensorflow_saved_model_bundle: TensorflowSavedModelBundleWeightsDescr | None = None
2965 torchscript: TorchscriptWeightsDescr | None = None
2967 @model_validator(mode="after")
2968 def check_entries(self) -> Self:
2969 entries = {wtype for wtype, entry in self if entry is not None}
2971 if not entries:
2972 raise ValueError("Missing weights entry")
2974 entries_wo_parent = {
2975 wtype
2976 for wtype, entry in self
2977 if entry is not None and hasattr(entry, "parent") and entry.parent is None
2978 }
2979 if len(entries_wo_parent) != 1:
2980 issue_warning(
2981 "Exactly one weights entry may not specify the `parent` field (got"
2982 + " {value}). That entry is considered the original set of model weights."
2983 + " Other weight formats are created through conversion of the orignal or"
2984 + " already converted weights. They have to reference the weights format"
2985 + " they were converted from as their `parent`.",
2986 value=len(entries_wo_parent),
2987 field="weights",
2988 )
2990 for wtype, entry in self:
2991 if entry is None:
2992 continue
2994 assert hasattr(entry, "type")
2995 assert hasattr(entry, "parent")
2996 assert wtype == entry.type
2997 if (
2998 entry.parent is not None and entry.parent not in entries
2999 ): # self reference checked for `parent` field
3000 raise ValueError(
3001 f"`weights.{wtype}.parent={entry.parent} not in specified weight"
3002 + f" formats: {entries}"
3003 )
3005 return self
3007 def __getitem__(
3008 self,
3009 key: WeightsFormat,
3010 ):
3011 if key == "keras_hdf5":
3012 ret = self.keras_hdf5
3013 elif key == "keras_v3":
3014 ret = self.keras_v3
3015 elif key == "onnx":
3016 ret = self.onnx
3017 elif key == "pytorch_state_dict":
3018 ret = self.pytorch_state_dict
3019 elif key == "tensorflow_js":
3020 ret = self.tensorflow_js
3021 elif key == "tensorflow_saved_model_bundle":
3022 ret = self.tensorflow_saved_model_bundle
3023 elif key == "torchscript":
3024 ret = self.torchscript
3025 else:
3026 raise KeyError(key)
3028 if ret is None:
3029 raise KeyError(key)
3031 return ret
3033 @overload
3034 def __setitem__(
3035 self, key: Literal["keras_hdf5"], value: KerasHdf5WeightsDescr | None
3036 ) -> None: ...
3037 @overload
3038 def __setitem__(
3039 self, key: Literal["keras_v3"], value: KerasV3WeightsDescr | None
3040 ) -> None: ...
3041 @overload
3042 def __setitem__(
3043 self, key: Literal["onnx"], value: OnnxWeightsDescr | None
3044 ) -> None: ...
3045 @overload
3046 def __setitem__(
3047 self,
3048 key: Literal["pytorch_state_dict"],
3049 value: PytorchStateDictWeightsDescr | None,
3050 ) -> None: ...
3051 @overload
3052 def __setitem__(
3053 self, key: Literal["tensorflow_js"], value: TensorflowJsWeightsDescr | None
3054 ) -> None: ...
3055 @overload
3056 def __setitem__(
3057 self,
3058 key: Literal["tensorflow_saved_model_bundle"],
3059 value: TensorflowSavedModelBundleWeightsDescr | None,
3060 ) -> None: ...
3061 @overload
3062 def __setitem__(
3063 self, key: Literal["torchscript"], value: TorchscriptWeightsDescr | None
3064 ) -> None: ...
3066 def __setitem__(
3067 self,
3068 key: WeightsFormat,
3069 value: SpecificWeightsDescr | None,
3070 ):
3071 if key == "keras_hdf5":
3072 if value is not None and not isinstance(value, KerasHdf5WeightsDescr):
3073 raise TypeError(
3074 f"Expected KerasHdf5WeightsDescr or None for key 'keras_hdf5', got {type(value)}"
3075 )
3076 self.keras_hdf5 = value
3077 elif key == "keras_v3":
3078 if value is not None and not isinstance(value, KerasV3WeightsDescr):
3079 raise TypeError(
3080 f"Expected KerasV3WeightsDescr or None for key 'keras_v3', got {type(value)}"
3081 )
3082 self.keras_v3 = value
3083 elif key == "onnx":
3084 if value is not None and not isinstance(value, OnnxWeightsDescr):
3085 raise TypeError(
3086 f"Expected OnnxWeightsDescr or None for key 'onnx', got {type(value)}"
3087 )
3088 self.onnx = value
3089 elif key == "pytorch_state_dict":
3090 if value is not None and not isinstance(
3091 value, PytorchStateDictWeightsDescr
3092 ):
3093 raise TypeError(
3094 f"Expected PytorchStateDictWeightsDescr or None for key 'pytorch_state_dict', got {type(value)}"
3095 )
3096 self.pytorch_state_dict = value
3097 elif key == "tensorflow_js":
3098 if value is not None and not isinstance(value, TensorflowJsWeightsDescr):
3099 raise TypeError(
3100 f"Expected TensorflowJsWeightsDescr or None for key 'tensorflow_js', got {type(value)}"
3101 )
3102 self.tensorflow_js = value
3103 elif key == "tensorflow_saved_model_bundle":
3104 if value is not None and not isinstance(
3105 value, TensorflowSavedModelBundleWeightsDescr
3106 ):
3107 raise TypeError(
3108 f"Expected TensorflowSavedModelBundleWeightsDescr or None for key 'tensorflow_saved_model_bundle', got {type(value)}"
3109 )
3110 self.tensorflow_saved_model_bundle = value
3111 elif key == "torchscript":
3112 if value is not None and not isinstance(value, TorchscriptWeightsDescr):
3113 raise TypeError(
3114 f"Expected TorchscriptWeightsDescr or None for key 'torchscript', got {type(value)}"
3115 )
3116 self.torchscript = value
3117 else:
3118 raise KeyError(key)
3120 @property
3121 def available_formats(self) -> dict[WeightsFormat, SpecificWeightsDescr]:
3122 return {
3123 **({} if self.keras_hdf5 is None else {"keras_hdf5": self.keras_hdf5}),
3124 **({} if self.keras_v3 is None else {"keras_v3": self.keras_v3}),
3125 **({} if self.onnx is None else {"onnx": self.onnx}),
3126 **(
3127 {}
3128 if self.pytorch_state_dict is None
3129 else {"pytorch_state_dict": self.pytorch_state_dict}
3130 ),
3131 **(
3132 {}
3133 if self.tensorflow_js is None
3134 else {"tensorflow_js": self.tensorflow_js}
3135 ),
3136 **(
3137 {}
3138 if self.tensorflow_saved_model_bundle is None
3139 else {
3140 "tensorflow_saved_model_bundle": self.tensorflow_saved_model_bundle
3141 }
3142 ),
3143 **({} if self.torchscript is None else {"torchscript": self.torchscript}),
3144 }
3146 @property
3147 def missing_formats(self) -> set[WeightsFormat]:
3148 return {
3149 wf for wf in get_args(WeightsFormat) if wf not in self.available_formats
3150 }
3153class LinkedModel(LinkedResourceBase):
3154 """Reference to a bioimage.io model."""
3156 id: ModelId
3157 """A valid model `id` from the bioimage.io collection."""
3160class _DataDepSize(NamedTuple):
3161 min: StrictInt
3162 max: StrictInt | None
3165class _AxisSizes(NamedTuple):
3166 """the lenghts of all axes of model inputs and outputs"""
3168 inputs: dict[tuple[TensorId, AxisId], int]
3169 outputs: dict[tuple[TensorId, AxisId], int | _DataDepSize]
3172class _TensorSizes(NamedTuple):
3173 """_AxisSizes as nested dicts"""
3175 inputs: dict[TensorId, dict[AxisId, int]]
3176 outputs: dict[TensorId, dict[AxisId, int | _DataDepSize]]
3179class ReproducibilityTolerance(Node, extra="allow"):
3180 """Describes what small numerical differences -- if any -- may be tolerated
3181 in the generated output when executing in different environments.
3183 A tensor element *output* is considered mismatched to the **test_tensor** if
3184 abs(*output* - **test_tensor**) > **absolute_tolerance** + **relative_tolerance** * abs(**test_tensor**).
3185 (Internally we call [numpy.testing.assert_allclose](https://numpy.org/doc/stable/reference/generated/numpy.testing.assert_allclose.html).)
3187 Motivation:
3188 For testing we can request the respective deep learning frameworks to be as
3189 reproducible as possible by setting seeds and chosing deterministic algorithms,
3190 but differences in operating systems, available hardware and installed drivers
3191 may still lead to numerical differences.
3192 """
3194 relative_tolerance: RelativeTolerance = 1e-3
3195 """Maximum relative tolerance of reproduced test tensor."""
3197 absolute_tolerance: AbsoluteTolerance = 1e-3
3198 """Maximum absolute tolerance of reproduced test tensor."""
3200 mismatched_elements_per_million: MismatchedElementsPerMillion = 100
3201 """Maximum number of mismatched elements/pixels per million to tolerate."""
3203 output_ids: Sequence[TensorId] = ()
3204 """Limits the output tensor IDs these reproducibility details apply to."""
3206 weights_formats: Sequence[WeightsFormat] = ()
3207 """Limits the weights formats these details apply to."""
3210class BiasRisksLimitations(Node, extra="allow"):
3211 """Known biases, risks, technical limitations, and recommendations for model use."""
3213 known_biases: str = dedent("""\
3214 In general bioimage models may suffer from biases caused by:
3216 - Imaging protocol dependencies
3217 - Use of a specific cell type
3218 - Species-specific training data limitations
3220 """)
3221 """Biases in training data or model behavior."""
3223 risks: str = dedent("""\
3224 Common risks in bioimage analysis include:
3226 - Erroneously assuming generalization to unseen experimental conditions
3227 - Trusting (overconfident) model outputs without validation
3228 - Misinterpretation of results
3230 """)
3231 """Potential risks in the context of bioimage analysis."""
3233 limitations: str | None = None
3234 """Technical limitations and failure modes."""
3236 recommendations: str = "Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model."
3237 """Mitigation strategies regarding `known_biases`, `risks`, and `limitations`, as well as applicable best practices.
3239 Consider:
3240 - How to use a validation dataset?
3241 - How to manually validate?
3242 - Feasibility of domain adaptation for different experimental setups?
3244 """
3246 def format_md(self) -> str:
3247 if self.limitations is None:
3248 limitations_header = ""
3249 else:
3250 limitations_header = "## Limitations\n\n"
3252 return f"""# Bias, Risks, and Limitations
3254{self.known_biases}
3256{self.risks}
3258{limitations_header}{self.limitations or ""}
3260## Recommendations
3262{self.recommendations}
3264"""
3267class TrainingDetails(Node, extra="allow"):
3268 training_preprocessing: str | None = None
3269 """Detailed image preprocessing steps during model training:
3271 Mention:
3272 - *Normalization methods*
3273 - *Augmentation strategies*
3274 - *Resizing/resampling procedures*
3275 - *Artifact handling*
3277 """
3279 training_epochs: float | None = None
3280 """Number of training epochs."""
3282 training_batch_size: float | None = None
3283 """Batch size used in training."""
3285 initial_learning_rate: float | None = None
3286 """Initial learning rate used in training."""
3288 learning_rate_schedule: str | None = None
3289 """Learning rate schedule used in training."""
3291 loss_function: str | None = None
3292 """Loss function used in training, e.g. nn.MSELoss."""
3294 loss_function_kwargs: dict[str, YamlValue] = Field(
3295 default_factory=cast(Callable[[], Dict[str, YamlValue]], dict)
3296 )
3297 """key word arguments for the `loss_function`"""
3299 optimizer: str | None = None
3300 """optimizer, e.g. torch.optim.Adam"""
3302 optimizer_kwargs: dict[str, YamlValue] = Field(
3303 default_factory=cast(Callable[[], Dict[str, YamlValue]], dict)
3304 )
3305 """key word arguments for the `optimizer`"""
3307 regularization: str | None = None
3308 """Regularization techniques used during training, e.g. drop-out or weight decay."""
3310 training_duration: float | None = None
3311 """Total training duration in hours."""
3314class Evaluation(Node, extra="allow"):
3315 model_id: ModelId | None = None
3316 """Model being evaluated."""
3318 dataset_id: DatasetId
3319 """Dataset used for evaluation."""
3321 dataset_source: HttpUrl
3322 """Source of the dataset."""
3324 dataset_role: Literal["train", "validation", "test", "independent", "unknown"]
3325 """Role of the dataset used for evaluation.
3327 - `train`: dataset was (part of) the training data
3328 - `validation`: dataset was (part of) the validation data used during training, e.g. used for model selection or hyperparameter tuning
3329 - `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
3330 - `independent`: dataset is entirely independent test data; not used during training or validation, and acquired from a different source/distribution than training data
3331 - `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.
3332 """
3334 sample_count: int
3335 """Number of evaluated samples."""
3337 evaluation_factors: list[Annotated[str, MaxLen(16)]]
3338 """(Abbreviations of) each evaluation factor.
3340 Evaluation factors are criteria along which model performance is evaluated, e.g. different image conditions
3341 like 'low SNR', 'high cell density', or different biological conditions like 'cell type A', 'cell type B'.
3342 An 'overall' factor may be included to summarize performance across all conditions.
3343 """
3345 evaluation_factors_long: list[str]
3346 """Descriptions (long form) of each evaluation factor."""
3348 metrics: list[Annotated[str, MaxLen(16)]]
3349 """(Abbreviations of) metrics used for evaluation."""
3351 metrics_long: list[str]
3352 """Description of each metric used."""
3354 @model_validator(mode="after")
3355 def _validate_list_lengths(self) -> Self:
3356 if len(self.evaluation_factors) != len(self.evaluation_factors_long):
3357 raise ValueError(
3358 "`evaluation_factors` and `evaluation_factors_long` must have the same length"
3359 )
3361 if len(self.metrics) != len(self.metrics_long):
3362 raise ValueError("`metrics` and `metrics_long` must have the same length")
3364 if len(self.results) != len(self.metrics):
3365 raise ValueError("`results` must have the same number of rows as `metrics`")
3367 for row in self.results:
3368 if len(row) != len(self.evaluation_factors):
3369 raise ValueError(
3370 "`results` must have the same number of columns (in every row) as `evaluation_factors`"
3371 )
3373 return self
3375 results: list[list[str | float | int]]
3376 """Results for each metric (rows; outer list) and each evaluation factor (columns; inner list)."""
3378 results_summary: str | None = None
3379 """Interpretation of results for general audience.
3381 Consider:
3382 - Overall model performance
3383 - Comparison to existing methods
3384 - Limitations and areas for improvement
3386"""
3388 def format_md(self):
3389 results_header = ["Metric"] + self.evaluation_factors
3390 results_table_cells = [results_header, ["---"] * len(results_header)] + [
3391 [metric] + [str(r) for r in row]
3392 for metric, row in zip(self.metrics, self.results)
3393 ]
3395 results_table = "".join(
3396 "| " + " | ".join(row) + " |\n" for row in results_table_cells
3397 )
3398 factors = "".join(
3399 f"\n - {ef}: {efl}"
3400 for ef, efl in zip(self.evaluation_factors, self.evaluation_factors_long)
3401 )
3402 metrics = "".join(
3403 f"\n - {em}: {eml}" for em, eml in zip(self.metrics, self.metrics_long)
3404 )
3406 return f"""## Testing Data, Factors & Metrics
3408Evaluation of {self.model_id or "this"} model on the {self.dataset_id} dataset (dataset role: {self.dataset_role}).
3410### Testing Data
3412- **Source:** [{self.dataset_id}]({self.dataset_source})
3413- **Size:** {self.sample_count} evaluated samples
3415### Factors
3416{factors}
3418### Metrics
3419{metrics}
3421## Results
3423### Quantitative Results
3425{results_table}
3427### Summary
3429{self.results_summary or "missing"}
3431"""
3434class EnvironmentalImpact(Node, extra="allow"):
3435 """Environmental considerations for model training and deployment.
3437 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).
3438 """
3440 hardware_type: str | None = None
3441 """GPU/CPU specifications"""
3443 hours_used: float | None = None
3444 """Total compute hours"""
3446 cloud_provider: str | None = None
3447 """If applicable"""
3449 compute_region: str | None = None
3450 """Geographic location"""
3452 co2_emitted: float | None = None
3453 """kg CO2 equivalent
3455 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).
3456 """
3458 def format_md(self):
3459 """Filled Markdown template section following [Hugging Face Model Card Template](https://huggingface.co/docs/hub/en/model-card-annotated)."""
3460 if self == self.__class__():
3461 return ""
3463 ret = "# Environmental Impact\n\n"
3464 if self.hardware_type is not None:
3465 ret += f"- **Hardware Type:** {self.hardware_type}\n"
3466 if self.hours_used is not None:
3467 ret += f"- **Hours used:** {self.hours_used}\n"
3468 if self.cloud_provider is not None:
3469 ret += f"- **Cloud Provider:** {self.cloud_provider}\n"
3470 if self.compute_region is not None:
3471 ret += f"- **Compute Region:** {self.compute_region}\n"
3472 if self.co2_emitted is not None:
3473 ret += f"- **Carbon Emitted:** {self.co2_emitted} kg CO2e\n"
3475 return ret + "\n"
3478class BioimageioConfig(Node, extra="allow"):
3479 reproducibility_tolerance: Sequence[ReproducibilityTolerance] = ()
3480 """Tolerances to allow when reproducing the model's test outputs
3481 from the model's test inputs.
3482 Only the first entry matching tensor id and weights format is considered.
3483 """
3485 funded_by: str | None = None
3486 """Funding agency, grant number if applicable"""
3488 architecture_type: Annotated[str, MaxLen(32)] | None = (
3489 None # TODO: add to differentiated tags
3490 )
3491 """Model architecture type, e.g., 3D U-Net, ResNet, transformer"""
3493 architecture_description: str | None = None
3494 """Text description of model architecture."""
3496 modality: str | None = None # TODO: add to differentiated tags
3497 """Input modality, e.g., fluorescence microscopy, electron microscopy"""
3499 target_structure: list[str] = Field( # TODO: add to differentiated tags
3500 default_factory=cast(Callable[[], List[str]], list)
3501 )
3502 """Biological structure(s) the model is designed to analyze, e.g., nuclei, mitochondria, cells"""
3504 task: str | None = None # TODO: add to differentiated tags
3505 """Bioimage-specific task type, e.g., segmentation, classification, detection, denoising"""
3507 new_version: ModelId | None = None
3508 """A new version of this model exists with a different model id."""
3510 out_of_scope_use: str | None = None
3511 """Describe how the model may be misused in bioimage analysis contexts and what users should **not** do with the model."""
3513 bias_risks_limitations: BiasRisksLimitations = Field(
3514 default_factory=BiasRisksLimitations.model_construct
3515 )
3516 """Description of known bias, risks, and technical limitations for in-scope model use."""
3518 model_parameter_count: int | None = None
3519 """Total number of model parameters."""
3521 training: TrainingDetails = Field(default_factory=TrainingDetails.model_construct)
3522 """Details on how the model was trained."""
3524 inference_time: str | None = None
3525 """Average inference time per image/tile. Specify hardware and image size. Multiple examples can be given."""
3527 memory_requirements_inference: str | None = None
3528 """GPU memory needed for inference. Multiple examples with different image size can be given."""
3530 memory_requirements_training: str | None = None
3531 """GPU memory needed for training. Multiple examples with different image/batch sizes can be given."""
3533 evaluations: list[Evaluation] = Field(
3534 default_factory=cast(Callable[[], List[Evaluation]], list)
3535 )
3536 """Quantitative model evaluations.
3538 Note:
3539 At the moment we recommend to include only a single test dataset
3540 (with evaluation factors that may mark subsets of the dataset)
3541 to avoid confusion and make the presentation of results cleaner.
3542 """
3544 environmental_impact: EnvironmentalImpact = Field(
3545 default_factory=EnvironmentalImpact.model_construct
3546 )
3547 """Environmental considerations for model training and deployment"""
3550class Config(Node, extra="allow"):
3551 bioimageio: BioimageioConfig = Field(
3552 default_factory=BioimageioConfig.model_construct
3553 )
3554 stardist: YamlValue = None
3557class ModelDescr(GenericModelDescrBase):
3558 """Specification of the fields used in a bioimage.io-compliant RDF to describe AI models with pretrained weights.
3559 These fields are typically stored in a YAML file which we call a model resource description file (model RDF).
3560 """
3562 implemented_format_version: ClassVar[Literal["0.5.13"]] = "0.5.13"
3563 if TYPE_CHECKING:
3564 format_version: Literal["0.5.13"] = "0.5.13"
3565 else:
3566 format_version: Literal["0.5.13"]
3567 """Version of the bioimage.io model description specification used.
3568 When creating a new model always use the latest micro/patch version described here.
3569 The `format_version` is important for any consumer software to understand how to parse the fields.
3570 """
3572 implemented_type: ClassVar[Literal["model"]] = "model"
3573 if TYPE_CHECKING:
3574 type: Literal["model"] = "model"
3575 else:
3576 type: Literal["model"]
3577 """Specialized resource type 'model'"""
3579 id: ModelId | None = None
3580 """bioimage.io-wide unique resource identifier
3581 assigned by bioimage.io; version **un**specific."""
3583 authors: FAIR[list[Author]] = Field(
3584 default_factory=cast(Callable[[], List[Author]], list)
3585 )
3586 """The authors are the creators of the model RDF and the primary points of contact."""
3588 documentation: FAIR[FileDescr_documentation | None] = None
3589 """Additional model documentation.
3590 The recommended documentation source file name is `README.md`. An `.md` suffix is mandatory.
3591 The documentation should include a '#[#] Validation' (sub)section
3592 with details on how to quantitatively validate the model on unseen data."""
3594 @field_validator("documentation", mode="after")
3595 @classmethod
3596 def _validate_documentation(cls, value: FileDescr | None) -> FileDescr | None:
3597 if not get_validation_context().perform_io_checks or value is None:
3598 return value
3600 doc_reader = get_reader(value)
3601 doc_content = doc_reader.read().decode(encoding="utf-8")
3602 if not re.search("#.*[vV]alidation", doc_content):
3603 issue_warning(
3604 "No '# Validation' (sub)section found in {value}.",
3605 value=value,
3606 field="documentation",
3607 )
3609 return value
3611 inputs: NotEmpty[Sequence[InputTensorDescr]]
3612 """Describes the input tensors expected by this model."""
3614 @field_validator("inputs", mode="after")
3615 @classmethod
3616 def _validate_input_axes(
3617 cls, inputs: Sequence[InputTensorDescr]
3618 ) -> Sequence[InputTensorDescr]:
3619 input_size_refs = cls._get_axes_with_independent_size(inputs)
3621 for i, ipt in enumerate(inputs):
3622 valid_independent_refs: dict[
3623 tuple[TensorId, AxisId],
3624 tuple[TensorDescr, AnyAxis, int | ParameterizedSize],
3625 ] = {
3626 **{
3627 (ipt.id, a.id): (ipt, a, a.size)
3628 for a in ipt.axes
3629 if not isinstance(a, BatchAxis)
3630 and isinstance(a.size, (int, ParameterizedSize))
3631 },
3632 **input_size_refs,
3633 }
3634 for a, ax in enumerate(ipt.axes):
3635 cls._validate_axis(
3636 "inputs",
3637 i=i,
3638 tensor_id=ipt.id,
3639 a=a,
3640 axis=ax,
3641 valid_independent_refs=valid_independent_refs,
3642 )
3643 return inputs
3645 @staticmethod
3646 def _validate_axis(
3647 field_name: str,
3648 i: int,
3649 tensor_id: TensorId,
3650 a: int,
3651 axis: AnyAxis,
3652 valid_independent_refs: dict[
3653 tuple[TensorId, AxisId],
3654 tuple[TensorDescr, AnyAxis, int | ParameterizedSize],
3655 ],
3656 ):
3657 if isinstance(axis, BatchAxis) or isinstance(
3658 axis.size, (int, ParameterizedSize, DataDependentSize)
3659 ):
3660 return
3661 elif not isinstance(axis.size, SizeReference):
3662 assert_never(axis.size)
3664 # validate axis.size SizeReference
3665 ref = (axis.size.tensor_id, axis.size.axis_id)
3666 if ref not in valid_independent_refs:
3667 raise ValueError(
3668 "Invalid tensor axis reference at"
3669 + f" {field_name}[{i}].axes[{a}].size: {axis.size}."
3670 )
3671 if ref == (tensor_id, axis.id):
3672 raise ValueError(
3673 "Self-referencing not allowed for"
3674 + f" {field_name}[{i}].axes[{a}].size: {axis.size}"
3675 )
3676 if axis.type == "channel":
3677 if valid_independent_refs[ref][1].type != "channel":
3678 raise ValueError(
3679 "A channel axis' size may only reference another fixed size"
3680 + " channel axis."
3681 )
3682 if isinstance(axis.channel_names, str) and "{i}" in axis.channel_names:
3683 ref_size = valid_independent_refs[ref][2]
3684 assert isinstance(ref_size, int), (
3685 "channel axis ref (another channel axis) has to specify fixed"
3686 + " size"
3687 )
3688 generated_channel_names = [
3689 axis.channel_names.format(i=i) for i in range(1, ref_size + 1)
3690 ]
3691 axis.channel_names = generated_channel_names
3693 if (ax_unit := getattr(axis, "unit", None)) != (
3694 ref_unit := getattr(valid_independent_refs[ref][1], "unit", None)
3695 ):
3696 raise ValueError(
3697 "The units of an axis and its reference axis need to match, but"
3698 + f" '{ax_unit}' != '{ref_unit}'."
3699 )
3700 ref_axis = valid_independent_refs[ref][1]
3701 if isinstance(ref_axis, BatchAxis):
3702 raise ValueError(
3703 f"Invalid reference axis '{ref_axis.id}' for {tensor_id}.{axis.id}"
3704 + " (a batch axis is not allowed as reference)."
3705 )
3707 if isinstance(axis, WithHalo):
3708 min_size = axis.size.get_size(axis, ref_axis, n=0)
3709 if (min_size - 2 * axis.halo) < 1:
3710 raise ValueError(
3711 f"axis {axis.id} with minimum size {min_size} is too small for halo"
3712 + f" {axis.halo}."
3713 )
3715 ref_halo = axis.halo * axis.scale / ref_axis.scale
3716 if ref_halo != int(ref_halo):
3717 raise ValueError(
3718 f"Inferred halo for {'.'.join(ref)} is not an integer ({ref_halo} ="
3719 + f" {tensor_id}.{axis.id}.halo {axis.halo}"
3720 + f" * {tensor_id}.{axis.id}.scale {axis.scale}"
3721 + f" / {'.'.join(ref)}.scale {ref_axis.scale})."
3722 )
3724 def validate_input_tensors(
3725 self,
3726 sources: Sequence[NDArray[Any]] | Mapping[TensorId, NDArray[Any] | None],
3727 *,
3728 pad_inputs: bool | Literal["allow"] = True,
3729 crop_outputs: bool | Literal["allow"] = True,
3730 ) -> Mapping[TensorId, NDArray[Any] | None]:
3731 """Check if the given input tensors match the model's input tensor descriptions.
3732 This includes checks of tensor shapes and dtypes, but not of the actual values.
3733 """
3734 if not isinstance(sources, collections.abc.Mapping):
3735 sources = {descr.id: tensor for descr, tensor in zip(self.inputs, sources)}
3737 tensors = {
3738 **{descr.id: (descr, sources.get(descr.id)) for descr in self.inputs},
3739 **{ # outputs are required for halo
3740 descr.id: (descr, None) for descr in self.outputs
3741 },
3742 }
3743 validate_tensors(tensors, pad_inputs=pad_inputs, crop_outputs=crop_outputs)
3745 return sources
3747 @model_validator(mode="after")
3748 def _validate_test_tensors(self) -> Self:
3749 if not get_validation_context().perform_io_checks:
3750 return self
3752 test_inputs = {
3753 descr.id: (
3754 descr,
3755 None if descr.test_tensor is None else load_array(descr.test_tensor),
3756 )
3757 for descr in self.inputs
3758 }
3759 test_outputs = {
3760 descr.id: (
3761 descr,
3762 None if descr.test_tensor is None else load_array(descr.test_tensor),
3763 )
3764 for descr in self.outputs
3765 }
3767 validate_tensors(
3768 {**test_inputs, **test_outputs},
3769 tensor_origin="test_tensor",
3770 pad_inputs="allow",
3771 crop_outputs="allow",
3772 )
3774 for rep_tol in self.config.bioimageio.reproducibility_tolerance:
3775 if not rep_tol.absolute_tolerance:
3776 continue
3778 if rep_tol.output_ids:
3779 out_arrays = {
3780 k: v[1] for k, v in test_outputs.items() if k in rep_tol.output_ids
3781 }
3782 else:
3783 out_arrays = {k: v[1] for k, v in test_outputs.items()}
3785 for out_id, array in out_arrays.items():
3786 if array is None:
3787 continue
3789 if rep_tol.absolute_tolerance > (max_test_value := array.max()) * 0.01:
3790 raise ValueError(
3791 "config.bioimageio.reproducibility_tolerance.absolute_tolerance="
3792 + f"{rep_tol.absolute_tolerance} > 0.01*{max_test_value}"
3793 + f" (1% of the maximum value of the test tensor '{out_id}')"
3794 )
3796 return self
3798 @model_validator(mode="after")
3799 def _validate_tensor_references_in_proc_kwargs(self, info: ValidationInfo) -> Self:
3800 ipt_refs = {t.id for t in self.inputs}
3801 missing_refs = [
3802 k["reference_tensor"]
3803 for k in [p.kwargs for ipt in self.inputs for p in ipt.preprocessing]
3804 + [p.kwargs for out in self.outputs for p in out.postprocessing]
3805 if "reference_tensor" in k
3806 and k["reference_tensor"] is not None
3807 and k["reference_tensor"] not in ipt_refs
3808 ]
3810 if missing_refs:
3811 raise ValueError(
3812 f"`reference_tensor`s {missing_refs} not found. Valid input tensor"
3813 + f" references are: {ipt_refs}."
3814 )
3816 return self
3818 name: Annotated[
3819 str,
3820 RestrictCharacters(string.ascii_letters + string.digits + "_+- ()"),
3821 MinLen(5),
3822 MaxLen(128),
3823 warn(MaxLen(64), "Name longer than 64 characters.", INFO),
3824 ]
3825 """A human-readable name of this model.
3826 It should be no longer than 64 characters
3827 and may only contain letter, number, underscore, minus, parentheses and spaces.
3828 We recommend to chose a name that refers to the model's task and image modality.
3829 """
3831 outputs: NotEmpty[Sequence[OutputTensorDescr]]
3832 """Describes the output tensors."""
3834 @field_validator("outputs", mode="after")
3835 @classmethod
3836 def _validate_tensor_ids(
3837 cls, outputs: Sequence[OutputTensorDescr], info: ValidationInfo
3838 ) -> Sequence[OutputTensorDescr]:
3839 tensor_ids = [
3840 t.id for t in info.data.get("inputs", []) + info.data.get("outputs", [])
3841 ]
3842 duplicate_tensor_ids: list[str] = []
3843 seen: set[str] = set()
3844 for t in tensor_ids:
3845 if t in seen:
3846 duplicate_tensor_ids.append(t)
3848 seen.add(t)
3850 if duplicate_tensor_ids:
3851 raise ValueError(f"Duplicate tensor ids: {duplicate_tensor_ids}")
3853 return outputs
3855 @staticmethod
3856 def _get_axes_with_parameterized_size(
3857 io: Sequence[InputTensorDescr] | Sequence[OutputTensorDescr],
3858 ):
3859 return {
3860 f"{t.id}.{a.id}": (t, a, a.size)
3861 for t in io
3862 for a in t.axes
3863 if not isinstance(a, BatchAxis) and isinstance(a.size, ParameterizedSize)
3864 }
3866 @staticmethod
3867 def _get_axes_with_independent_size(
3868 io: Sequence[InputTensorDescr] | Sequence[OutputTensorDescr],
3869 ):
3870 return {
3871 (t.id, a.id): (t, a, a.size)
3872 for t in io
3873 for a in t.axes
3874 if not isinstance(a, BatchAxis)
3875 and isinstance(a.size, (int, ParameterizedSize))
3876 }
3878 @field_validator("outputs", mode="after")
3879 @classmethod
3880 def _validate_output_axes(
3881 cls, outputs: list[OutputTensorDescr], info: ValidationInfo
3882 ) -> list[OutputTensorDescr]:
3883 input_size_refs = cls._get_axes_with_independent_size(
3884 info.data.get("inputs", [])
3885 )
3886 output_size_refs = cls._get_axes_with_independent_size(outputs)
3888 for i, out in enumerate(outputs):
3889 valid_independent_refs: dict[
3890 tuple[TensorId, AxisId],
3891 tuple[TensorDescr, AnyAxis, int | ParameterizedSize],
3892 ] = {
3893 **{
3894 (out.id, a.id): (out, a, a.size)
3895 for a in out.axes
3896 if not isinstance(a, BatchAxis)
3897 and isinstance(a.size, (int, ParameterizedSize))
3898 },
3899 **input_size_refs,
3900 **output_size_refs,
3901 }
3902 for a, ax in enumerate(out.axes):
3903 cls._validate_axis(
3904 "outputs",
3905 i,
3906 out.id,
3907 a,
3908 ax,
3909 valid_independent_refs=valid_independent_refs,
3910 )
3912 return outputs
3914 packaged_by: list[Author] = Field(
3915 default_factory=cast(Callable[[], List[Author]], list)
3916 )
3917 """The persons that have packaged and uploaded this model.
3918 Only required if those persons differ from the `authors`."""
3920 parent: LinkedModel | None = None
3921 """The model from which this model is derived, e.g. by fine-tuning the weights."""
3923 @model_validator(mode="after")
3924 def _validate_parent_is_not_self(self) -> Self:
3925 if self.parent is not None and self.parent.id == self.id:
3926 raise ValueError("A model description may not reference itself as parent.")
3928 return self
3930 run_mode: Annotated[
3931 RunMode | None,
3932 warn(None, "Run mode '{value}' has limited support across consumer softwares."),
3933 ] = None
3934 """Custom run mode for this model: for more complex prediction procedures like test time
3935 data augmentation that currently cannot be expressed in the specification.
3936 No standard run modes are defined yet."""
3938 timestamp: Datetime = Field(default_factory=Datetime.now)
3939 """Timestamp in [ISO 8601](#https://en.wikipedia.org/wiki/ISO_8601) format
3940 with a few restrictions listed [here](https://docs.python.org/3/library/datetime.html#datetime.datetime.fromisoformat).
3941 (In Python a datetime object is valid, too)."""
3943 training_data: Annotated[
3944 None | LinkedDataset | DatasetDescr | DatasetDescr02,
3945 Field(union_mode="left_to_right"),
3946 ] = None
3947 """The dataset used to train this model"""
3949 weights: Annotated[WeightsDescr, WrapSerializer(package_weights)]
3950 """The weights for this model.
3951 Weights can be given for different formats, but should otherwise be equivalent.
3952 The available weight formats determine which consumers can use this model."""
3954 config: Config = Field(default_factory=Config.model_construct)
3956 @model_validator(mode="after")
3957 def _add_default_cover(self) -> Self:
3958 if not get_validation_context().perform_io_checks or self.covers:
3959 return self
3961 try:
3962 generated_covers = generate_covers(
3963 [
3964 (t, load_array(t.test_tensor))
3965 for t in self.inputs
3966 if t.test_tensor is not None
3967 ],
3968 [
3969 (t, load_array(t.test_tensor))
3970 for t in self.outputs
3971 if t.test_tensor is not None
3972 ],
3973 )
3974 except Exception as e:
3975 issue_warning(
3976 "Failed to generate cover image(s): {e}",
3977 value=self.covers,
3978 msg_context={"e": e},
3979 field="covers",
3980 )
3981 else:
3982 self.covers.extend(generated_covers)
3984 return self
3986 def get_input_test_arrays(self) -> list[NDArray[Any]]:
3987 return self._get_test_arrays(self.inputs)
3989 def get_output_test_arrays(self) -> list[NDArray[Any]]:
3990 return self._get_test_arrays(self.outputs)
3992 @staticmethod
3993 def _get_test_arrays(
3994 io_descr: Sequence[InputTensorDescr] | Sequence[OutputTensorDescr],
3995 ):
3996 ts: list[FileDescr] = []
3997 for d in io_descr:
3998 if d.test_tensor is None:
3999 raise ValueError(
4000 f"Failed to get test arrays: description of '{d.id}' is missing a `test_tensor`."
4001 )
4002 ts.append(d.test_tensor)
4004 data = [load_array(t) for t in ts]
4005 assert all(isinstance(d, np.ndarray) for d in data)
4006 return data
4008 @staticmethod
4009 def get_batch_size(tensor_sizes: Mapping[TensorId, Mapping[AxisId, int]]) -> int:
4010 batch_size = 1
4011 tensor_with_batchsize: TensorId | None = None
4012 for tid in tensor_sizes:
4013 for aid, s in tensor_sizes[tid].items():
4014 if aid != BATCH_AXIS_ID or s == 1 or s == batch_size:
4015 continue
4017 if batch_size != 1:
4018 assert tensor_with_batchsize is not None
4019 raise ValueError(
4020 f"batch size mismatch for tensors '{tensor_with_batchsize}' ({batch_size}) and '{tid}' ({s})"
4021 )
4023 batch_size = s
4024 tensor_with_batchsize = tid
4026 return batch_size
4028 def get_output_tensor_sizes(
4029 self, input_sizes: Mapping[TensorId, Mapping[AxisId, int]]
4030 ) -> dict[TensorId, dict[AxisId, int | _DataDepSize]]:
4031 """Returns the tensor output sizes for given **input_sizes**.
4032 Only if **input_sizes** has a valid input shape, the tensor output size is exact.
4033 Otherwise it might be larger than the actual (valid) output"""
4034 batch_size = self.get_batch_size(input_sizes)
4035 ns = self.get_ns(input_sizes)
4037 tensor_sizes = self.get_tensor_sizes(ns, batch_size=batch_size)
4038 return tensor_sizes.outputs
4040 def get_ns(self, input_sizes: Mapping[TensorId, Mapping[AxisId, int]]):
4041 """get parameter `n` for each parameterized axis
4042 such that the valid input size is >= the given input size"""
4043 ret: dict[tuple[TensorId, AxisId], ParameterizedSize_N] = {}
4044 axes = {t.id: {a.id: a for a in t.axes} for t in self.inputs}
4045 for tid in input_sizes:
4046 for aid, s in input_sizes[tid].items():
4047 size_descr = axes[tid][aid].size
4048 if isinstance(size_descr, ParameterizedSize):
4049 ret[(tid, aid)] = size_descr.get_n(s)
4050 elif size_descr is None or isinstance(size_descr, (int, SizeReference)):
4051 pass
4052 else:
4053 assert_never(size_descr)
4055 return ret
4057 def get_tensor_sizes(
4058 self,
4059 ns: Mapping[tuple[TensorId, AxisId], ParameterizedSize_N],
4060 batch_size: int,
4061 max_input_shape: Mapping[TensorId, Mapping[AxisId, int]] | None = None,
4062 ) -> _TensorSizes:
4063 max_axis_sizes: dict[tuple[TensorId, AxisId], int] = {}
4064 for m, this_max_axis_sizes in (max_input_shape or {}).items():
4065 for a, s in this_max_axis_sizes.items():
4066 max_axis_sizes[(m, a)] = s
4068 axis_sizes = self.get_axis_sizes(
4069 ns, batch_size=batch_size, max_input_shape=max_axis_sizes
4070 )
4071 return _TensorSizes(
4072 {
4073 t: {
4074 aa: axis_sizes.inputs[(tt, aa)]
4075 for tt, aa in axis_sizes.inputs
4076 if tt == t
4077 }
4078 for t in {tt for tt, _ in axis_sizes.inputs}
4079 },
4080 {
4081 t: {
4082 aa: axis_sizes.outputs[(tt, aa)]
4083 for tt, aa in axis_sizes.outputs
4084 if tt == t
4085 }
4086 for t in {tt for tt, _ in axis_sizes.outputs}
4087 },
4088 )
4090 def get_axis_sizes(
4091 self,
4092 ns: Mapping[tuple[TensorId, AxisId], ParameterizedSize_N],
4093 batch_size: int | None = None,
4094 *,
4095 max_input_shape: Mapping[tuple[TensorId, AxisId], int] | None = None,
4096 ) -> _AxisSizes:
4097 """Determine input and output block shape for scale factors **ns**
4098 of parameterized input sizes.
4100 Args:
4101 ns: Scale factor `n` for each axis (keyed by (tensor_id, axis_id))
4102 that is parameterized as `size = min + n * step`.
4103 batch_size: The desired size of the batch dimension.
4104 If given **batch_size** overwrites any batch size present in
4105 **max_input_shape**. Default 1.
4106 max_input_shape: Limits the derived block shapes.
4107 Each axis for which the input size, parameterized by `n`, is larger
4108 than **max_input_shape** is set to the minimal value `n_min` for which
4109 this is still true.
4110 Use this for small input samples or large values of **ns**.
4111 Or simply whenever you know the full input shape.
4113 Returns:
4114 Resolved axis sizes for model inputs and outputs.
4115 """
4116 max_input_shape = max_input_shape or {}
4117 if batch_size is None:
4118 for (_t_id, a_id), s in max_input_shape.items():
4119 if a_id == BATCH_AXIS_ID:
4120 batch_size = s
4121 break
4122 else:
4123 batch_size = 1
4125 all_axes = {
4126 t.id: {a.id: a for a in t.axes} for t in chain(self.inputs, self.outputs)
4127 }
4129 inputs: dict[tuple[TensorId, AxisId], int] = {}
4130 outputs: dict[tuple[TensorId, AxisId], int | _DataDepSize] = {}
4132 def get_axis_size(a: InputAxis | OutputAxis):
4133 if isinstance(a, BatchAxis):
4134 if (t_descr.id, a.id) in ns:
4135 logger.warning(
4136 "Ignoring unexpected size increment factor (n) for batch axis"
4137 + " of tensor '{}'.",
4138 t_descr.id,
4139 )
4140 return batch_size
4141 elif isinstance(a.size, int):
4142 if (t_descr.id, a.id) in ns:
4143 logger.warning(
4144 "Ignoring unexpected size increment factor (n) for fixed size"
4145 + " axis '{}' of tensor '{}'.",
4146 a.id,
4147 t_descr.id,
4148 )
4149 return a.size
4150 elif isinstance(a.size, ParameterizedSize):
4151 if (t_descr.id, a.id) not in ns:
4152 raise ValueError(
4153 "Size increment factor (n) missing for parametrized axis"
4154 + f" '{a.id}' of tensor '{t_descr.id}'."
4155 )
4156 n = ns[(t_descr.id, a.id)]
4157 s_max = max_input_shape.get((t_descr.id, a.id))
4158 if s_max is not None:
4159 n = min(n, a.size.get_n(s_max))
4161 return a.size.get_size(n)
4163 elif isinstance(a.size, SizeReference):
4164 if (t_descr.id, a.id) in ns:
4165 logger.warning(
4166 "Ignoring unexpected size increment factor (n) for axis '{}'"
4167 + " of tensor '{}' with size reference.",
4168 a.id,
4169 t_descr.id,
4170 )
4171 assert not isinstance(a, BatchAxis)
4172 ref_axis = all_axes[a.size.tensor_id][a.size.axis_id]
4173 assert not isinstance(ref_axis, BatchAxis)
4174 ref_key = (a.size.tensor_id, a.size.axis_id)
4175 ref_size = inputs.get(ref_key, outputs.get(ref_key))
4176 assert ref_size is not None, ref_key
4177 assert not isinstance(ref_size, _DataDepSize), ref_key
4178 return a.size.get_size(
4179 axis=a,
4180 ref_axis=ref_axis,
4181 ref_size=ref_size,
4182 )
4183 elif isinstance(a.size, DataDependentSize):
4184 if (t_descr.id, a.id) in ns:
4185 logger.warning(
4186 "Ignoring unexpected increment factor (n) for data dependent"
4187 + " size axis '{}' of tensor '{}'.",
4188 a.id,
4189 t_descr.id,
4190 )
4191 return _DataDepSize(a.size.min, a.size.max)
4192 else:
4193 assert_never(a.size)
4195 # first resolve all , but the `SizeReference` input sizes
4196 for t_descr in self.inputs:
4197 for a in t_descr.axes:
4198 if not isinstance(a.size, SizeReference):
4199 s = get_axis_size(a)
4200 assert not isinstance(s, _DataDepSize)
4201 inputs[t_descr.id, a.id] = s
4203 # resolve all other input axis sizes
4204 for t_descr in self.inputs:
4205 for a in t_descr.axes:
4206 if isinstance(a.size, SizeReference):
4207 s = get_axis_size(a)
4208 assert not isinstance(s, _DataDepSize)
4209 inputs[t_descr.id, a.id] = s
4211 # resolve all output axis sizes
4212 for t_descr in self.outputs:
4213 for a in t_descr.axes:
4214 assert not isinstance(a.size, ParameterizedSize)
4215 s = get_axis_size(a)
4216 outputs[t_descr.id, a.id] = s
4218 return _AxisSizes(inputs=inputs, outputs=outputs)
4220 @model_validator(mode="before")
4221 @classmethod
4222 def _convert(cls, data: dict[str, Any]) -> dict[str, Any]:
4223 cls.convert_from_old_format_wo_validation(data)
4224 return data
4226 @classmethod
4227 def convert_from_old_format_wo_validation(cls, data: dict[str, Any]) -> None:
4228 """Convert metadata following an older format version to this classes' format
4229 without validating the result.
4230 """
4231 if (
4232 data.get("type") == "model"
4233 and isinstance(fv := data.get("format_version"), str)
4234 and fv.count(".") == 2
4235 ):
4236 fv_parts = fv.split(".")
4237 if any(not p.isdigit() for p in fv_parts):
4238 return
4240 fv_tuple = tuple(map(int, fv_parts))
4242 assert cls.implemented_format_version_tuple[0:2] == (0, 5)
4243 if fv_tuple[:2] in ((0, 3), (0, 4)):
4244 m04 = _ModelDescr_v0_4.load(data)
4245 if isinstance(m04, InvalidDescr):
4246 try:
4247 updated = _model_conv.convert_as_dict(
4248 m04 # pyright: ignore[reportArgumentType]
4249 )
4250 except Exception as e:
4251 logger.error(
4252 "Failed to convert from invalid model 0.4 description."
4253 + f"\nerror: {e}"
4254 + "\nProceeding with model 0.5 validation without conversion."
4255 )
4256 updated = None
4257 else:
4258 updated = _model_conv.convert_as_dict(m04)
4260 if updated is not None:
4261 data.clear()
4262 data.update(updated)
4264 elif fv_tuple[:2] == (0, 5):
4265 # bump patch version
4266 data["format_version"] = cls.implemented_format_version
4268 if fv_tuple[:2] in ((0, 3), (0, 4)) or (
4269 fv_tuple[:2] == (0, 5) and fv_tuple[2] < 11
4270 ):
4271 convert_plain_covers_and_docs_and_icon(data)
4274class _ModelConv(Converter[_ModelDescr_v0_4, ModelDescr]):
4275 def _convert(
4276 self, src: _ModelDescr_v0_4, tgt: type[ModelDescr | dict[str, Any]]
4277 ) -> ModelDescr | dict[str, Any]:
4278 name = "".join(
4279 c if c in string.ascii_letters + string.digits + "_+- ()" else " "
4280 for c in src.name
4281 )
4283 def conv_authors(auths: Sequence[_Author_v0_4] | None):
4284 conv = (
4285 _author_conv.convert if TYPE_CHECKING else _author_conv.convert_as_dict
4286 )
4287 return None if auths is None else [conv(a) for a in auths]
4289 if TYPE_CHECKING:
4290 arch_file_conv = _arch_file_conv.convert
4291 arch_lib_conv = _arch_lib_conv.convert
4292 else:
4293 arch_file_conv = _arch_file_conv.convert_as_dict
4294 arch_lib_conv = _arch_lib_conv.convert_as_dict
4296 input_size_refs = {
4297 ipt.name: {
4298 a: s
4299 for a, s in zip(
4300 ipt.axes,
4301 (
4302 ipt.shape.min
4303 if isinstance(ipt.shape, _ParameterizedInputShape_v0_4)
4304 else ipt.shape
4305 ),
4306 )
4307 }
4308 for ipt in src.inputs
4309 if ipt.shape
4310 }
4311 output_size_refs = {
4312 **{
4313 out.name: {a: s for a, s in zip(out.axes, out.shape)}
4314 for out in src.outputs
4315 if not isinstance(out.shape, _ImplicitOutputShape_v0_4)
4316 },
4317 **input_size_refs,
4318 }
4320 return tgt(
4321 attachments=(
4322 []
4323 if src.attachments is None
4324 else [FileDescr(source=f) for f in src.attachments.files]
4325 ),
4326 authors=[_author_conv.convert_as_dict(a) for a in src.authors], # pyright: ignore[reportArgumentType]
4327 cite=[{"text": c.text, "doi": c.doi, "url": c.url} for c in src.cite], # pyright: ignore[reportArgumentType]
4328 config=src.config, # pyright: ignore[reportArgumentType]
4329 covers=[{"source": c} for c in src.covers], # pyright: ignore[reportArgumentType]
4330 description=src.description,
4331 documentation={"source": src.documentation} if src.documentation else None, # pyright: ignore[reportArgumentType]
4332 format_version="0.5.13",
4333 git_repo=src.git_repo, # pyright: ignore[reportArgumentType]
4334 icon={"source": src.icon} if src.icon else None, # pyright: ignore[reportArgumentType]
4335 id=None if src.id is None else ModelId(src.id),
4336 id_emoji=src.id_emoji,
4337 license=src.license, # type: ignore
4338 links=src.links,
4339 maintainers=[_maintainer_conv.convert_as_dict(m) for m in src.maintainers], # pyright: ignore[reportArgumentType]
4340 name=name,
4341 tags=src.tags,
4342 type=src.type,
4343 uploader=src.uploader,
4344 version=src.version,
4345 inputs=[ # pyright: ignore[reportArgumentType]
4346 _input_tensor_conv.convert_as_dict(ipt, tt, st, input_size_refs)
4347 for ipt, tt, st in zip(
4348 src.inputs,
4349 src.test_inputs,
4350 src.sample_inputs or [None] * len(src.test_inputs),
4351 )
4352 ],
4353 outputs=[ # pyright: ignore[reportArgumentType]
4354 _output_tensor_conv.convert_as_dict(out, tt, st, output_size_refs)
4355 for out, tt, st in zip(
4356 src.outputs,
4357 src.test_outputs,
4358 src.sample_outputs or [None] * len(src.test_outputs),
4359 )
4360 ],
4361 parent=(
4362 None
4363 if src.parent is None
4364 else LinkedModel(
4365 id=ModelId(
4366 str(src.parent.id)
4367 + (
4368 ""
4369 if src.parent.version_number is None
4370 else f"/{src.parent.version_number}"
4371 )
4372 )
4373 )
4374 ),
4375 training_data=(
4376 None
4377 if src.training_data is None
4378 else (
4379 LinkedDataset(
4380 id=DatasetId(
4381 str(src.training_data.id)
4382 + (
4383 ""
4384 if src.training_data.version_number is None
4385 else f"/{src.training_data.version_number}"
4386 )
4387 )
4388 )
4389 if isinstance(src.training_data, LinkedDataset02)
4390 else src.training_data
4391 )
4392 ),
4393 packaged_by=[_author_conv.convert_as_dict(a) for a in src.packaged_by], # pyright: ignore[reportArgumentType]
4394 run_mode=src.run_mode,
4395 timestamp=src.timestamp,
4396 weights=(WeightsDescr if TYPE_CHECKING else dict)(
4397 keras_hdf5=(w := src.weights.keras_hdf5)
4398 and (KerasHdf5WeightsDescr if TYPE_CHECKING else dict)(
4399 authors=conv_authors(w.authors),
4400 source=w.source,
4401 tensorflow_version=w.tensorflow_version or Version("1.15"),
4402 parent=w.parent,
4403 ),
4404 onnx=(w := src.weights.onnx)
4405 and (OnnxWeightsDescr if TYPE_CHECKING else dict)(
4406 source=w.source,
4407 authors=conv_authors(w.authors),
4408 parent=w.parent,
4409 opset_version=w.opset_version or 15,
4410 ),
4411 pytorch_state_dict=(w := src.weights.pytorch_state_dict)
4412 and (PytorchStateDictWeightsDescr if TYPE_CHECKING else dict)(
4413 source=w.source,
4414 authors=conv_authors(w.authors),
4415 parent=w.parent,
4416 architecture=(
4417 arch_file_conv(
4418 w.architecture,
4419 w.architecture_sha256,
4420 w.kwargs,
4421 )
4422 if isinstance(w.architecture, _CallableFromFile_v0_4)
4423 else arch_lib_conv(w.architecture, w.kwargs)
4424 ),
4425 pytorch_version=w.pytorch_version or Version("1.10"),
4426 dependencies=(
4427 None
4428 if w.dependencies is None
4429 else (FileDescr if TYPE_CHECKING else dict)(
4430 source=cast(
4431 FileSource,
4432 str(deps := w.dependencies)[
4433 (
4434 len("conda:")
4435 if str(deps).startswith("conda:")
4436 else 0
4437 ) :
4438 ],
4439 )
4440 )
4441 ),
4442 ),
4443 tensorflow_js=(w := src.weights.tensorflow_js)
4444 and (TensorflowJsWeightsDescr if TYPE_CHECKING else dict)(
4445 source=w.source,
4446 authors=conv_authors(w.authors),
4447 parent=w.parent,
4448 tensorflow_version=w.tensorflow_version or Version("1.15"),
4449 ),
4450 tensorflow_saved_model_bundle=(
4451 w := src.weights.tensorflow_saved_model_bundle
4452 )
4453 and (TensorflowSavedModelBundleWeightsDescr if TYPE_CHECKING else dict)(
4454 authors=conv_authors(w.authors),
4455 parent=w.parent,
4456 source=w.source,
4457 tensorflow_version=w.tensorflow_version or Version("1.15"),
4458 dependencies=(
4459 None
4460 if w.dependencies is None
4461 else (FileDescr if TYPE_CHECKING else dict)(
4462 source=cast(
4463 FileSource,
4464 (
4465 str(w.dependencies)[len("conda:") :]
4466 if str(w.dependencies).startswith("conda:")
4467 else str(w.dependencies)
4468 ),
4469 )
4470 )
4471 ),
4472 ),
4473 torchscript=(w := src.weights.torchscript)
4474 and (TorchscriptWeightsDescr if TYPE_CHECKING else dict)(
4475 source=w.source,
4476 authors=conv_authors(w.authors),
4477 parent=w.parent,
4478 pytorch_version=w.pytorch_version or Version("1.10"),
4479 ),
4480 ),
4481 )
4484_model_conv = _ModelConv(_ModelDescr_v0_4, ModelDescr)
4487# create better cover images for 3d data and non-image outputs
4488def generate_covers(
4489 inputs: Sequence[tuple[InputTensorDescr, NDArray[Any]]],
4490 outputs: Sequence[tuple[OutputTensorDescr, NDArray[Any]]],
4491) -> list[FileDescr]:
4492 def squeeze(
4493 data: NDArray[Any], axes: Sequence[AnyAxis]
4494 ) -> tuple[NDArray[Any], list[AnyAxis]]:
4495 """apply numpy.ndarray.squeeze while keeping track of the axis descriptions remaining"""
4496 if data.ndim != len(axes):
4497 raise ValueError(
4498 f"tensor shape {data.shape} does not match described axes"
4499 + f" {[a.id for a in axes]}"
4500 )
4502 axes = [deepcopy(a) for a, s in zip(axes, data.shape) if s != 1]
4503 return data.squeeze(), axes
4505 def normalize(
4506 data: NDArray[Any], axis: tuple[int, ...] | None, eps: float = 1e-7
4507 ) -> NDArray[np.float32]:
4508 data = data.astype("float32")
4509 data -= data.min(axis=axis, keepdims=True)
4510 data /= data.max(axis=axis, keepdims=True) + eps
4511 return data
4513 def to_2d_image(data: NDArray[Any], axes: Sequence[AnyAxis]):
4514 original_shape = data.shape
4515 original_axes = list(axes)
4516 data, axes = squeeze(data, axes)
4518 # take slice fom any batch or index axis if needed
4519 # and convert the first channel axis and take a slice from any additional channel axes
4520 slices: tuple[slice, ...] = ()
4521 ndim = data.ndim
4522 ndim_need = 3 if any(isinstance(a, ChannelAxis) for a in axes) else 2
4523 has_c_axis = False
4524 for i, a in enumerate(axes):
4525 s = data.shape[i]
4526 assert s > 1
4527 if (
4528 isinstance(a, (BatchAxis, IndexInputAxis, IndexOutputAxis))
4529 and ndim > ndim_need
4530 ):
4531 data = data[slices + (slice(s // 2 - 1, s // 2),)]
4532 ndim -= 1
4533 elif isinstance(a, ChannelAxis):
4534 if has_c_axis:
4535 # second channel axis
4536 data = data[slices + (slice(0, 1),)]
4537 ndim -= 1
4538 else:
4539 has_c_axis = True
4540 if s == 2:
4541 # visualize two channels with cyan and magenta
4542 data = np.concatenate(
4543 [
4544 data[slices + (slice(1, 2),)],
4545 data[slices + (slice(0, 1),)],
4546 (
4547 data[slices + (slice(0, 1),)]
4548 + data[slices + (slice(1, 2),)]
4549 )
4550 / 2, # TODO: take maximum instead?
4551 ],
4552 axis=i,
4553 )
4554 elif data.shape[i] == 3:
4555 pass # visualize 3 channels as RGB
4556 else:
4557 # visualize first 3 channels as RGB
4558 data = data[slices + (slice(3),)]
4560 assert data.shape[i] == 3
4562 slices += (slice(None),)
4564 data, axes = squeeze(data, axes)
4565 assert len(axes) == ndim
4566 # take slice from z axis if needed
4567 slices = ()
4568 if ndim > ndim_need:
4569 for i, a in enumerate(axes):
4570 s = data.shape[i]
4571 if a.id == AxisId("z"):
4572 data = data[slices + (slice(s // 2 - 1, s // 2),)]
4573 data, axes = squeeze(data, axes)
4574 ndim -= 1
4575 break
4577 slices += (slice(None),)
4579 # take slice from any space or time axis
4580 slices = ()
4582 for i, a in enumerate(axes):
4583 if ndim <= ndim_need:
4584 break
4586 s = data.shape[i]
4587 assert s > 1
4588 if isinstance(
4589 a, (SpaceInputAxis, SpaceOutputAxis, TimeInputAxis, TimeOutputAxis)
4590 ):
4591 data = data[slices + (slice(s // 2 - 1, s // 2),)]
4592 ndim -= 1
4594 slices += (slice(None),)
4596 del slices
4597 data, axes = squeeze(data, axes)
4598 assert len(axes) == ndim
4600 if (has_c_axis and ndim != 3) or (not has_c_axis and ndim != 2):
4601 raise ValueError(
4602 f"Failed to construct cover image from shape {original_shape} with axes {[a.id for a in original_axes]}."
4603 )
4605 if not has_c_axis:
4606 assert ndim == 2
4607 data = np.repeat(data[:, :, None], 3, axis=2)
4608 axes.append(ChannelAxis(channel_names=list("RGB")))
4609 ndim += 1
4611 assert ndim == 3
4613 # transpose axis order such that longest axis comes first...
4614 axis_order: list[int] = [int(i) for i in np.argsort(list(data.shape))]
4615 axis_order.reverse()
4616 # ... and channel axis is last
4617 c = next(i for i in range(3) if isinstance(axes[i], ChannelAxis))
4618 axis_order.append(axis_order.pop(c))
4619 axes = [axes[ao] for ao in axis_order]
4620 data = data.transpose(axis_order)
4622 # h, w = data.shape[:2]
4623 # if h / w in (1.0 or 2.0):
4624 # pass
4625 # elif h / w < 2:
4626 # TODO: enforce 2:1 or 1:1 aspect ratio for generated cover images
4628 norm_along = (
4629 tuple(i for i, a in enumerate(axes) if a.type in ("space", "time")) or None
4630 )
4631 # normalize the data and map to 8 bit
4632 data = normalize(data, norm_along)
4633 data = (data * 255).astype("uint8")
4635 return data
4637 def create_diagonal_split_image(im0: NDArray[Any], im1: NDArray[Any]):
4638 assert im0.dtype == im1.dtype == np.uint8
4639 assert im0.shape == im1.shape
4640 assert im0.ndim == 3
4641 N, M, C = im0.shape
4642 assert C == 3
4643 out = np.ones((N, M, C), dtype="uint8")
4644 for c in range(C):
4645 outc = np.tril(im0[..., c])
4646 mask = outc == 0
4647 outc[mask] = np.triu(im1[..., c])[mask]
4648 out[..., c] = outc
4650 return out
4652 if not inputs:
4653 raise ValueError("Missing test input tensor for cover generation.")
4655 if not outputs:
4656 raise ValueError("Missing test output tensor for cover generation.")
4658 ipt_descr, ipt = inputs[0]
4659 out_descr, out = outputs[0]
4661 ipt_img = to_2d_image(ipt, ipt_descr.axes)
4662 out_img = to_2d_image(out, out_descr.axes)
4664 cover_folder = Path(mkdtemp())
4665 if ipt_img.shape == out_img.shape:
4666 covers = [cover_folder / "cover.png"]
4667 imwrite(covers[0], create_diagonal_split_image(ipt_img, out_img))
4668 else:
4669 covers = [cover_folder / "input.png", cover_folder / "output.png"]
4670 imwrite(covers[0], ipt_img)
4671 imwrite(covers[1], out_img)
4673 return [FileDescr(source=c) for c in covers]