Coverage for src/bioimageio/spec/model/v0_4.py: 91%
593 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-18 09:17 +0000
1from __future__ import annotations
3import collections.abc
4from typing import (
5 TYPE_CHECKING,
6 Any,
7 Callable,
8 ClassVar,
9 Dict,
10 List,
11 Literal,
12 Sequence,
13 Union,
14 cast,
15)
17import numpy as np
18from annotated_types import Ge, Interval, MaxLen, MinLen, MultipleOf
19from numpy.typing import NDArray
20from pydantic import (
21 AllowInfNan,
22 Discriminator,
23 Field,
24 RootModel,
25 SerializationInfo,
26 SerializerFunctionWrapHandler,
27 StringConstraints,
28 TypeAdapter,
29 ValidationInfo,
30 WrapSerializer,
31 field_validator,
32 model_validator,
33)
34from typing_extensions import Annotated, Self, assert_never, get_args
36from .._internal.common_nodes import (
37 KwargsNode,
38 Node,
39 NodeWithExplicitlySetFields,
40)
41from .._internal.constants import SHA256_HINT
42from .._internal.field_validation import validate_unique_entries
43from .._internal.field_warning import issue_warning, warn
44from .._internal.io import BioimageioYamlContent, WithSuffix
45from .._internal.io import FileDescr as FileDescr
46from .._internal.io_basics import Sha256 as Sha256
47from .._internal.io_packaging import FileSource_package, include_in_package
48from .._internal.io_utils import load_array
49from .._internal.packaging_context import packaging_context_var
50from .._internal.types import Datetime as Datetime
51from .._internal.types import FileSource, LowerCaseIdentifier
52from .._internal.types import Identifier as Identifier
53from .._internal.types import LicenseId as LicenseId
54from .._internal.types import NotEmpty as NotEmpty
55from .._internal.url import HttpUrl as HttpUrl
56from .._internal.validated_string_with_inner_node import ValidatedStringWithInnerNode
57from .._internal.validator_annotations import AfterValidator, RestrictCharacters
58from .._internal.version_type import Version as Version
59from .._internal.warning_levels import ALERT, INFO
60from ..dataset.v0_2 import VALID_COVER_IMAGE_EXTENSIONS as VALID_COVER_IMAGE_EXTENSIONS
61from ..dataset.v0_2 import DatasetDescr as DatasetDescr
62from ..dataset.v0_2 import LinkedDataset as LinkedDataset
63from ..generic.v0_2 import AttachmentsDescr as AttachmentsDescr
64from ..generic.v0_2 import Author as Author
65from ..generic.v0_2 import BadgeDescr as BadgeDescr
66from ..generic.v0_2 import CiteEntry as CiteEntry
67from ..generic.v0_2 import Doi as Doi
68from ..generic.v0_2 import GenericModelDescrBase
69from ..generic.v0_2 import LinkedResource as LinkedResource
70from ..generic.v0_2 import Maintainer as Maintainer
71from ..generic.v0_2 import OrcidId as OrcidId
72from ..generic.v0_2 import RelativeFilePath as RelativeFilePath
73from ..generic.v0_2 import ResourceId as ResourceId
74from ..generic.v0_2 import Uploader as Uploader
75from ._v0_4_converter import convert_from_older_format
78class ModelId(ResourceId):
79 pass
82AxesStr = Annotated[
83 str, RestrictCharacters("bitczyx"), AfterValidator(validate_unique_entries)
84]
85AxesInCZYX = Annotated[
86 str, RestrictCharacters("czyx"), AfterValidator(validate_unique_entries)
87]
89PostprocessingName = Literal[
90 "binarize",
91 "clip",
92 "scale_linear",
93 "sigmoid",
94 "zero_mean_unit_variance",
95 "scale_range",
96 "scale_mean_variance",
97]
98PreprocessingName = Literal[
99 "binarize",
100 "clip",
101 "scale_linear",
102 "sigmoid",
103 "zero_mean_unit_variance",
104 "scale_range",
105]
108class TensorName(LowerCaseIdentifier):
109 pass
112class CallableFromDepencencyNode(Node):
113 _submodule_adapter: ClassVar[TypeAdapter[Identifier]] = TypeAdapter(Identifier)
115 module_name: str
116 """The Python module that implements **callable_name**."""
118 @field_validator("module_name", mode="after")
119 def _check_submodules(cls, module_name: str) -> str:
120 for submod in module_name.split("."):
121 _ = cls._submodule_adapter.validate_python(submod)
123 return module_name
125 callable_name: Identifier
126 """The callable Python identifier implemented in module **module_name**."""
129class CallableFromDepencency(ValidatedStringWithInnerNode[CallableFromDepencencyNode]):
130 _inner_node_class = CallableFromDepencencyNode
131 root_model: ClassVar[type[RootModel[Any]]] = RootModel[
132 Annotated[
133 str,
134 StringConstraints(strip_whitespace=True, pattern=r"^.+\..+$"),
135 ]
136 ]
138 @classmethod
139 def _get_data(cls, valid_string_data: str):
140 *mods, callname = valid_string_data.split(".")
141 return {"module_name": ".".join(mods), "callable_name": callname}
143 @property
144 def module_name(self):
145 """The Python module that implements **callable_name**."""
146 return self._inner_node.module_name
148 @property
149 def callable_name(self):
150 """The callable Python identifier implemented in module **module_name**."""
151 return self._inner_node.callable_name
154class CallableFromFileNode(Node):
155 source_file: Annotated[
156 RelativeFilePath | HttpUrl,
157 Field(union_mode="left_to_right"),
158 include_in_package,
159 ]
160 """The Python source file that implements **callable_name**."""
161 callable_name: Identifier
162 """The callable Python identifier implemented in **source_file**."""
165class CallableFromFile(ValidatedStringWithInnerNode[CallableFromFileNode]):
166 _inner_node_class = CallableFromFileNode
167 root_model: ClassVar[type[RootModel[Any]]] = RootModel[
168 Annotated[
169 str,
170 StringConstraints(strip_whitespace=True, pattern=r"^.+:.+$"),
171 ]
172 ]
174 @classmethod
175 def _get_data(cls, valid_string_data: str):
176 *file_parts, callname = valid_string_data.split(":")
177 return {"source_file": ":".join(file_parts), "callable_name": callname}
179 @property
180 def source_file(self):
181 """The Python source file that implements **callable_name**."""
182 return self._inner_node.source_file
184 @property
185 def callable_name(self):
186 """The callable Python identifier implemented in **source_file**."""
187 return self._inner_node.callable_name
190CustomCallable = Annotated[
191 Union[CallableFromFile, CallableFromDepencency], Field(union_mode="left_to_right")
192]
195class DependenciesNode(Node):
196 manager: Annotated[NotEmpty[str], Field(examples=["conda", "maven", "pip"])]
197 """Dependency manager"""
199 file: FileSource_package
200 """Dependency file"""
203class Dependencies(ValidatedStringWithInnerNode[DependenciesNode]):
204 _inner_node_class = DependenciesNode
205 root_model: ClassVar[type[RootModel[Any]]] = RootModel[
206 Annotated[
207 str,
208 StringConstraints(strip_whitespace=True, pattern=r"^.+:.+$"),
209 ]
210 ]
212 @classmethod
213 def _get_data(cls, valid_string_data: str):
214 manager, *file_parts = valid_string_data.split(":")
215 return {"manager": manager, "file": ":".join(file_parts)}
217 @property
218 def manager(self):
219 """Dependency manager"""
220 return self._inner_node.manager
222 @property
223 def file(self):
224 """Dependency file"""
225 return self._inner_node.file
228WeightsFormat = Literal[
229 "keras_hdf5",
230 "onnx",
231 "pytorch_state_dict",
232 "tensorflow_js",
233 "tensorflow_saved_model_bundle",
234 "torchscript",
235]
238class WeightsEntryDescrBase(FileDescr):
239 type: ClassVar[WeightsFormat]
240 weights_format_name: ClassVar[str] # human readable
242 source: FileSource_package
243 """The weights file."""
245 attachments: Annotated[
246 AttachmentsDescr | None,
247 warn(None, "Weights entry depends on additional attachments.", ALERT),
248 ] = None
249 """Attachments that are specific to this weights entry."""
251 authors: list[Author] | None = None
252 """Authors
253 Either the person(s) that have trained this model resulting in the original weights file.
254 (If this is the initial weights entry, i.e. it does not have a `parent`)
255 Or the person(s) who have converted the weights to this weights format.
256 (If this is a child weight, i.e. it has a `parent` field)
257 """
259 dependencies: Annotated[
260 Dependencies | None,
261 warn(
262 None,
263 "Custom dependencies ({value}) specified. Avoid this whenever possible "
264 + "to allow execution in a wider range of software environments.",
265 ),
266 Field(
267 examples=[
268 "conda:environment.yaml",
269 "maven:./pom.xml",
270 "pip:./requirements.txt",
271 ]
272 ),
273 ] = None
274 """Dependency manager and dependency file, specified as `<dependency manager>:<relative file path>`."""
276 parent: Annotated[WeightsFormat | None, Field(examples=["pytorch_state_dict"])] = (
277 None
278 )
279 """The source weights these weights were converted from.
280 For example, if a model's weights were converted from the `pytorch_state_dict` format to `torchscript`,
281 The `pytorch_state_dict` weights entry has no `parent` and is the parent of the `torchscript` weights.
282 All weight entries except one (the initial set of weights resulting from training the model),
283 need to have this field."""
285 @model_validator(mode="after")
286 def check_parent_is_not_self(self) -> Self:
287 if self.type == self.parent:
288 raise ValueError("Weights entry can't be it's own parent.")
290 return self
293class KerasHdf5WeightsDescr(WeightsEntryDescrBase):
294 type: ClassVar[WeightsFormat] = "keras_hdf5"
295 weights_format_name: ClassVar[str] = "Keras HDF5"
296 tensorflow_version: Version | None = None
297 """TensorFlow version used to create these weights"""
299 @field_validator("tensorflow_version", mode="after")
300 @classmethod
301 def _tfv(cls, value: Any):
302 if value is None:
303 issue_warning(
304 "missing. Please specify the TensorFlow version"
305 + " these weights were created with.",
306 value=value,
307 severity=ALERT,
308 field="tensorflow_version",
309 )
310 return value
313class OnnxWeightsDescr(WeightsEntryDescrBase):
314 type: ClassVar[WeightsFormat] = "onnx"
315 weights_format_name: ClassVar[str] = "ONNX"
316 opset_version: Annotated[int, Ge(7)] | None = None
317 """ONNX opset version"""
319 @field_validator("opset_version", mode="after")
320 @classmethod
321 def _ov(cls, value: Any):
322 if value is None:
323 issue_warning(
324 "Missing ONNX opset version (aka ONNX opset number). "
325 + "Please specify the ONNX opset version these weights were created"
326 + " with.",
327 value=value,
328 severity=ALERT,
329 field="opset_version",
330 )
331 return value
334class PytorchStateDictWeightsDescr(WeightsEntryDescrBase):
335 type: ClassVar[WeightsFormat] = "pytorch_state_dict"
336 weights_format_name: ClassVar[str] = "Pytorch State Dict"
337 architecture: CustomCallable = Field(
338 examples=["my_function.py:MyNetworkClass", "my_module.submodule.get_my_model"]
339 )
340 """callable returning a torch.nn.Module instance.
341 Local implementation: `<relative path to file>:<identifier of implementation within the file>`.
342 Implementation in a dependency: `<dependency-package>.<[dependency-module]>.<identifier>`."""
344 architecture_sha256: Annotated[
345 Sha256 | None,
346 Field(
347 description=(
348 "The SHA256 of the architecture source file, if the architecture is not"
349 " defined in a module listed in `dependencies`\n"
350 )
351 + SHA256_HINT,
352 ),
353 ] = None
354 """The SHA256 of the architecture source file,
355 if the architecture is not defined in a module listed in `dependencies`"""
357 @model_validator(mode="after")
358 def check_architecture_sha256(self) -> Self:
359 if isinstance(self.architecture, CallableFromFile):
360 if self.architecture_sha256 is None:
361 raise ValueError(
362 "Missing required `architecture_sha256` for `architecture` with"
363 + " source file."
364 )
365 elif self.architecture_sha256 is not None:
366 raise ValueError(
367 "Got `architecture_sha256` for architecture that does not have a source"
368 + " file."
369 )
371 return self
373 kwargs: dict[str, Any] = Field(
374 default_factory=cast(Callable[[], Dict[str, Any]], dict)
375 )
376 """key word arguments for the `architecture` callable"""
378 pytorch_version: Version | None = None
379 """Version of the PyTorch library used.
380 If `depencencies` is specified it should include pytorch and the verison has to match.
381 (`dependencies` overrules `pytorch_version`)"""
383 @field_validator("pytorch_version", mode="after")
384 @classmethod
385 def _ptv(cls, value: Any):
386 if value is None:
387 issue_warning(
388 "missing. Please specify the PyTorch version these"
389 + " PyTorch state dict weights were created with.",
390 value=value,
391 severity=ALERT,
392 field="pytorch_version",
393 )
394 return value
397class TorchscriptWeightsDescr(WeightsEntryDescrBase):
398 type: ClassVar[WeightsFormat] = "torchscript"
399 weights_format_name: ClassVar[str] = "TorchScript"
400 pytorch_version: Version | None = None
401 """Version of the PyTorch library used."""
403 @field_validator("pytorch_version", mode="after")
404 @classmethod
405 def _ptv(cls, value: Any):
406 if value is None:
407 issue_warning(
408 "missing. Please specify the PyTorch version these"
409 + " Torchscript weights were created with.",
410 value=value,
411 severity=ALERT,
412 field="pytorch_version",
413 )
414 return value
417class TensorflowJsWeightsDescr(WeightsEntryDescrBase):
418 type: ClassVar[WeightsFormat] = "tensorflow_js"
419 weights_format_name: ClassVar[str] = "Tensorflow.js"
420 tensorflow_version: Version | None = None
421 """Version of the TensorFlow library used."""
423 @field_validator("tensorflow_version", mode="after")
424 @classmethod
425 def _tfv(cls, value: Any):
426 if value is None:
427 issue_warning(
428 "missing. Please specify the TensorFlow version"
429 + " these TensorflowJs weights were created with.",
430 value=value,
431 severity=ALERT,
432 field="tensorflow_version",
433 )
434 return value
436 source: FileSource_package
437 """The multi-file weights.
438 All required files/folders should be a zip archive."""
441class TensorflowSavedModelBundleWeightsDescr(WeightsEntryDescrBase):
442 type: ClassVar[WeightsFormat] = "tensorflow_saved_model_bundle"
443 weights_format_name: ClassVar[str] = "Tensorflow Saved Model"
444 tensorflow_version: Version | None = None
445 """Version of the TensorFlow library used."""
447 @field_validator("tensorflow_version", mode="after")
448 @classmethod
449 def _tfv(cls, value: Any):
450 if value is None:
451 issue_warning(
452 "missing. Please specify the TensorFlow version"
453 + " these Tensorflow saved model bundle weights were created with.",
454 value=value,
455 severity=ALERT,
456 field="tensorflow_version",
457 )
458 return value
461class WeightsDescr(Node):
462 keras_hdf5: KerasHdf5WeightsDescr | None = None
463 onnx: OnnxWeightsDescr | None = None
464 pytorch_state_dict: PytorchStateDictWeightsDescr | None = None
465 tensorflow_js: TensorflowJsWeightsDescr | None = None
466 tensorflow_saved_model_bundle: TensorflowSavedModelBundleWeightsDescr | None = None
467 torchscript: TorchscriptWeightsDescr | None = None
469 @model_validator(mode="after")
470 def check_one_entry(self) -> Self:
471 if all(
472 entry is None
473 for entry in [
474 self.keras_hdf5,
475 self.onnx,
476 self.pytorch_state_dict,
477 self.tensorflow_js,
478 self.tensorflow_saved_model_bundle,
479 self.torchscript,
480 ]
481 ):
482 raise ValueError("Missing weights entry")
484 return self
486 def __getitem__(
487 self,
488 key: WeightsFormat,
489 ):
490 if key == "keras_hdf5":
491 ret = self.keras_hdf5
492 elif key == "onnx":
493 ret = self.onnx
494 elif key == "pytorch_state_dict":
495 ret = self.pytorch_state_dict
496 elif key == "tensorflow_js":
497 ret = self.tensorflow_js
498 elif key == "tensorflow_saved_model_bundle":
499 ret = self.tensorflow_saved_model_bundle
500 elif key == "torchscript":
501 ret = self.torchscript
502 else:
503 raise KeyError(key)
505 if ret is None:
506 raise KeyError(key)
508 return ret
510 @property
511 def available_formats(self):
512 return {
513 **({} if self.keras_hdf5 is None else {"keras_hdf5": self.keras_hdf5}),
514 **({} if self.onnx is None else {"onnx": self.onnx}),
515 **(
516 {}
517 if self.pytorch_state_dict is None
518 else {"pytorch_state_dict": self.pytorch_state_dict}
519 ),
520 **(
521 {}
522 if self.tensorflow_js is None
523 else {"tensorflow_js": self.tensorflow_js}
524 ),
525 **(
526 {}
527 if self.tensorflow_saved_model_bundle is None
528 else {
529 "tensorflow_saved_model_bundle": self.tensorflow_saved_model_bundle
530 }
531 ),
532 **({} if self.torchscript is None else {"torchscript": self.torchscript}),
533 }
535 @property
536 def missing_formats(self):
537 return {
538 wf for wf in get_args(WeightsFormat) if wf not in self.available_formats
539 }
542class ParameterizedInputShape(Node):
543 """A sequence of valid shapes given by `shape_k = min + k * step for k in {0, 1, ...}`."""
545 min: NotEmpty[list[int]]
546 """The minimum input shape"""
548 step: NotEmpty[list[int]]
549 """The minimum shape change"""
551 def __len__(self) -> int:
552 return len(self.min)
554 @model_validator(mode="after")
555 def matching_lengths(self) -> Self:
556 if len(self.min) != len(self.step):
557 raise ValueError("`min` and `step` required to have the same length")
559 return self
562class ImplicitOutputShape(Node):
563 """Output tensor shape depending on an input tensor shape.
564 `shape(output_tensor) = shape(input_tensor) * scale + 2 * offset`"""
566 reference_tensor: TensorName
567 """Name of the reference tensor."""
569 scale: NotEmpty[list[float | None]]
570 """output_pix/input_pix for each dimension.
571 'null' values indicate new dimensions, whose length is defined by 2*`offset`"""
573 offset: NotEmpty[list[int | Annotated[float, MultipleOf(0.5)]]]
574 """Position of origin wrt to input."""
576 def __len__(self) -> int:
577 return len(self.scale)
579 @model_validator(mode="after")
580 def matching_lengths(self) -> Self:
581 if len(self.scale) != len(self.offset):
582 raise ValueError(
583 f"scale {self.scale} has to have same length as offset {self.offset}!"
584 )
585 # if we have an expanded dimension, make sure that it's offet is not zero
586 for sc, off in zip(self.scale, self.offset):
587 if sc is None and not off:
588 raise ValueError("`offset` must not be zero if `scale` is none/zero")
590 return self
593class TensorDescrBase(Node):
594 name: TensorName
595 """Tensor name. No duplicates are allowed."""
597 description: str = ""
599 axes: AxesStr
600 """Axes identifying characters. Same length and order as the axes in `shape`.
601 | axis | description |
602 | --- | --- |
603 | b | batch (groups multiple samples) |
604 | i | instance/index/element |
605 | t | time |
606 | c | channel |
607 | z | spatial dimension z |
608 | y | spatial dimension y |
609 | x | spatial dimension x |
610 """
612 data_range: (
613 tuple[Annotated[float, AllowInfNan(True)], Annotated[float, AllowInfNan(True)]]
614 | None
615 ) = None
616 """Tuple `(minimum, maximum)` specifying the allowed range of the data in this tensor.
617 If not specified, the full data range that can be expressed in `data_type` is allowed."""
620class BinarizeKwargs(KwargsNode):
621 """key word arguments for `BinarizeDescr`"""
623 threshold: float
624 """The fixed threshold"""
627class BinarizeDescr(NodeWithExplicitlySetFields):
628 """BinarizeDescr the tensor with a fixed `BinarizeKwargs.threshold`.
629 Values above the threshold will be set to one, values below the threshold to zero.
630 """
632 implemented_name: ClassVar[Literal["binarize"]] = "binarize"
633 if TYPE_CHECKING:
634 name: Literal["binarize"] = "binarize"
635 else:
636 name: Literal["binarize"]
638 kwargs: BinarizeKwargs
641class ClipKwargs(KwargsNode):
642 """key word arguments for `ClipDescr`"""
644 min: float
645 """minimum value for clipping"""
646 max: float
647 """maximum value for clipping"""
650class ClipDescr(NodeWithExplicitlySetFields):
651 """Clip tensor values to a range.
653 Set tensor values below `ClipKwargs.min` to `ClipKwargs.min`
654 and above `ClipKwargs.max` to `ClipKwargs.max`.
655 """
657 implemented_name: ClassVar[Literal["clip"]] = "clip"
658 if TYPE_CHECKING:
659 name: Literal["clip"] = "clip"
660 else:
661 name: Literal["clip"]
663 kwargs: ClipKwargs
666class ScaleLinearKwargs(KwargsNode):
667 """key word arguments for `ScaleLinearDescr`"""
669 axes: Annotated[AxesInCZYX | None, Field(examples=["xy"])] = None
670 """The subset of axes to scale jointly.
671 For example xy to scale the two image axes for 2d data jointly."""
673 gain: float | list[float] = 1.0
674 """multiplicative factor"""
676 offset: float | list[float] = 0.0
677 """additive term"""
679 @model_validator(mode="after")
680 def either_gain_or_offset(self) -> Self:
681 if (
682 self.gain == 1.0
683 or isinstance(self.gain, list)
684 and all(g == 1.0 for g in self.gain)
685 ) and (
686 self.offset == 0.0
687 or isinstance(self.offset, list)
688 and all(off == 0.0 for off in self.offset)
689 ):
690 raise ValueError(
691 "Redunt linear scaling not allowd. Set `gain` != 1.0 and/or `offset` !="
692 + " 0.0."
693 )
695 return self
698class ScaleLinearDescr(NodeWithExplicitlySetFields):
699 """Fixed linear scaling."""
701 implemented_name: ClassVar[Literal["scale_linear"]] = "scale_linear"
702 if TYPE_CHECKING:
703 name: Literal["scale_linear"] = "scale_linear"
704 else:
705 name: Literal["scale_linear"]
707 kwargs: ScaleLinearKwargs
710class SigmoidDescr(NodeWithExplicitlySetFields):
711 """The logistic sigmoid funciton, a.k.a. expit function."""
713 implemented_name: ClassVar[Literal["sigmoid"]] = "sigmoid"
714 if TYPE_CHECKING:
715 name: Literal["sigmoid"] = "sigmoid"
716 else:
717 name: Literal["sigmoid"]
719 @property
720 def kwargs(self) -> KwargsNode:
721 """empty kwargs"""
722 return KwargsNode()
725class ZeroMeanUnitVarianceKwargs(KwargsNode):
726 """key word arguments for `ZeroMeanUnitVarianceDescr`"""
728 mode: Literal["fixed", "per_dataset", "per_sample"] = "fixed"
729 """Mode for computing mean and variance.
730 | mode | description |
731 | ----------- | ------------------------------------ |
732 | fixed | Fixed values for mean and variance |
733 | per_dataset | Compute for the entire dataset |
734 | per_sample | Compute for each sample individually |
735 """
736 axes: Annotated[AxesInCZYX, Field(examples=["xy"])]
737 """The subset of axes to normalize jointly.
738 For example `xy` to normalize the two image axes for 2d data jointly."""
740 mean: Annotated[
741 float | NotEmpty[list[float]] | None, Field(examples=[(1.1, 2.2, 3.3)])
742 ] = None
743 """The mean value(s) to use for `mode: fixed`.
744 For example `[1.1, 2.2, 3.3]` in the case of a 3 channel image with `axes: xy`."""
745 # todo: check if means match input axes (for mode 'fixed')
747 std: Annotated[
748 float | NotEmpty[list[float]] | None, Field(examples=[(0.1, 0.2, 0.3)])
749 ] = None
750 """The standard deviation values to use for `mode: fixed`. Analogous to mean."""
752 eps: Annotated[float, Interval(gt=0, le=0.1)] = 1e-6
753 """epsilon for numeric stability: `out = (tensor - mean) / (std + eps)`."""
755 @model_validator(mode="after")
756 def mean_and_std_match_mode(self) -> Self:
757 if self.mode == "fixed" and (self.mean is None or self.std is None):
758 raise ValueError("`mean` and `std` are required for `mode: fixed`.")
759 elif self.mode != "fixed" and (self.mean is not None or self.std is not None):
760 raise ValueError(f"`mean` and `std` not allowed for `mode: {self.mode}`")
762 return self
765class ZeroMeanUnitVarianceDescr(NodeWithExplicitlySetFields):
766 """Subtract mean and divide by variance."""
768 implemented_name: ClassVar[Literal["zero_mean_unit_variance"]] = (
769 "zero_mean_unit_variance"
770 )
771 if TYPE_CHECKING:
772 name: Literal["zero_mean_unit_variance"] = "zero_mean_unit_variance"
773 else:
774 name: Literal["zero_mean_unit_variance"]
776 kwargs: ZeroMeanUnitVarianceKwargs
779class ScaleRangeKwargs(KwargsNode):
780 """key word arguments for `ScaleRangeDescr`
782 For `min_percentile`=0.0 (the default) and `max_percentile`=100 (the default)
783 this processing step normalizes data to the [0, 1] intervall.
784 For other percentiles the normalized values will partially be outside the [0, 1]
785 intervall. Use `ScaleRange` followed by `ClipDescr` if you want to limit the
786 normalized values to a range.
787 """
789 mode: Literal["per_dataset", "per_sample"]
790 """Mode for computing percentiles.
791 | mode | description |
792 | ----------- | ------------------------------------ |
793 | per_dataset | compute for the entire dataset |
794 | per_sample | compute for each sample individually |
795 """
796 axes: Annotated[AxesInCZYX, Field(examples=["xy"])]
797 """The subset of axes to normalize jointly.
798 For example xy to normalize the two image axes for 2d data jointly."""
800 min_percentile: Annotated[int | float, Interval(ge=0, lt=100)] = 0.0
801 """The lower percentile used to determine the value to align with zero."""
803 max_percentile: Annotated[int | float, Interval(gt=1, le=100)] = 100.0
804 """The upper percentile used to determine the value to align with one.
805 Has to be bigger than `min_percentile`.
806 The range is 1 to 100 instead of 0 to 100 to avoid mistakenly
807 accepting percentiles specified in the range 0.0 to 1.0."""
809 @model_validator(mode="after")
810 def min_smaller_max(self, info: ValidationInfo) -> Self:
811 if self.min_percentile >= self.max_percentile:
812 raise ValueError(
813 f"min_percentile {self.min_percentile} >= max_percentile"
814 + f" {self.max_percentile}"
815 )
817 return self
819 eps: Annotated[float, Interval(gt=0, le=0.1)] = 1e-6
820 """Epsilon for numeric stability.
821 `out = (tensor - v_lower) / (v_upper - v_lower + eps)`;
822 with `v_lower,v_upper` values at the respective percentiles."""
824 reference_tensor: TensorName | None = None
825 """Tensor name to compute the percentiles from. Default: The tensor itself.
826 For any tensor in `inputs` only input tensor references are allowed.
827 For a tensor in `outputs` only input tensor refereences are allowed if `mode: per_dataset`"""
830class ScaleRangeDescr(NodeWithExplicitlySetFields):
831 """Scale with percentiles."""
833 implemented_name: ClassVar[Literal["scale_range"]] = "scale_range"
834 if TYPE_CHECKING:
835 name: Literal["scale_range"] = "scale_range"
836 else:
837 name: Literal["scale_range"]
839 kwargs: ScaleRangeKwargs
842class ScaleMeanVarianceKwargs(KwargsNode):
843 """key word arguments for `ScaleMeanVarianceDescr`"""
845 mode: Literal["per_dataset", "per_sample"]
846 """Mode for computing mean and variance.
847 | mode | description |
848 | ----------- | ------------------------------------ |
849 | per_dataset | Compute for the entire dataset |
850 | per_sample | Compute for each sample individually |
851 """
853 reference_tensor: TensorName
854 """Name of tensor to match."""
856 axes: Annotated[AxesInCZYX | None, Field(examples=["xy"])] = None
857 """The subset of axes to scale jointly.
858 For example xy to normalize the two image axes for 2d data jointly.
859 Default: scale all non-batch axes jointly."""
861 eps: Annotated[float, Interval(gt=0, le=0.1)] = 1e-6
862 """Epsilon for numeric stability:
863 "`out = (tensor - mean) / (std + eps) * (ref_std + eps) + ref_mean."""
866class ScaleMeanVarianceDescr(NodeWithExplicitlySetFields):
867 """Scale the tensor s.t. its mean and variance match a reference tensor."""
869 implemented_name: ClassVar[Literal["scale_mean_variance"]] = "scale_mean_variance"
870 if TYPE_CHECKING:
871 name: Literal["scale_mean_variance"] = "scale_mean_variance"
872 else:
873 name: Literal["scale_mean_variance"]
875 kwargs: ScaleMeanVarianceKwargs
878PreprocessingDescr = Annotated[
879 Union[
880 BinarizeDescr,
881 ClipDescr,
882 ScaleLinearDescr,
883 SigmoidDescr,
884 ZeroMeanUnitVarianceDescr,
885 ScaleRangeDescr,
886 ],
887 Discriminator("name"),
888]
889PostprocessingDescr = Annotated[
890 Union[
891 BinarizeDescr,
892 ClipDescr,
893 ScaleLinearDescr,
894 SigmoidDescr,
895 ZeroMeanUnitVarianceDescr,
896 ScaleRangeDescr,
897 ScaleMeanVarianceDescr,
898 ],
899 Discriminator("name"),
900]
903class InputTensorDescr(TensorDescrBase):
904 data_type: Literal["float32", "uint8", "uint16"]
905 """For now an input tensor is expected to be given as `float32`.
906 The data flow in bioimage.io models is explained
907 [in this diagram.](https://docs.google.com/drawings/d/1FTw8-Rn6a6nXdkZ_SkMumtcjvur9mtIhRqLwnKqZNHM/edit)."""
909 shape: Annotated[
910 Sequence[int] | ParameterizedInputShape,
911 Field(
912 examples=[(1, 512, 512, 1), {"min": (1, 64, 64, 1), "step": (0, 32, 32, 0)}]
913 ),
914 ]
915 """Specification of input tensor shape."""
917 preprocessing: list[PreprocessingDescr] = Field(
918 default_factory=cast( # TODO: (py>3.8) use list[PreprocessingDesr]
919 Callable[[], List[PreprocessingDescr]], list
920 )
921 )
922 """Description of how this input should be preprocessed."""
924 @model_validator(mode="after")
925 def zero_batch_step_and_one_batch_size(self) -> Self:
926 bidx = self.axes.find("b")
927 if bidx == -1:
928 return self
930 if isinstance(self.shape, ParameterizedInputShape):
931 step = self.shape.step
932 shape = self.shape.min
933 if step[bidx] != 0:
934 raise ValueError(
935 "Input shape step has to be zero in the batch dimension (the batch"
936 + " dimension can always be increased, but `step` should specify how"
937 + " to increase the minimal shape to find the largest single batch"
938 + " shape)"
939 )
940 else:
941 shape = self.shape
943 if shape[bidx] != 1:
944 raise ValueError("Input shape has to be 1 in the batch dimension b.")
946 return self
948 @model_validator(mode="after")
949 def validate_preprocessing_kwargs(self) -> Self:
950 for p in self.preprocessing:
951 kwargs_axes = p.kwargs.get("axes")
952 if isinstance(kwargs_axes, str) and any(
953 a not in self.axes for a in kwargs_axes
954 ):
955 raise ValueError("`kwargs.axes` needs to be subset of `axes`")
957 return self
960class OutputTensorDescr(TensorDescrBase):
961 data_type: Literal[
962 "float32",
963 "float64",
964 "uint8",
965 "int8",
966 "uint16",
967 "int16",
968 "uint32",
969 "int32",
970 "uint64",
971 "int64",
972 "bool",
973 ]
974 """Data type.
975 The data flow in bioimage.io models is explained
976 [in this diagram.](https://docs.google.com/drawings/d/1FTw8-Rn6a6nXdkZ_SkMumtcjvur9mtIhRqLwnKqZNHM/edit)."""
978 shape: Sequence[int] | ImplicitOutputShape
979 """Output tensor shape."""
981 halo: Sequence[int] | None = None
982 """The `halo` that should be cropped from the output tensor to avoid boundary effects.
983 The `halo` is to be cropped from both sides, i.e. `shape_after_crop = shape - 2 * halo`.
984 To document a `halo` that is already cropped by the model `shape.offset` has to be used instead."""
986 postprocessing: list[PostprocessingDescr] = Field(
987 default_factory=cast(Callable[[], List[PostprocessingDescr]], list)
988 )
989 """Description of how this output should be postprocessed."""
991 @model_validator(mode="after")
992 def matching_halo_length(self) -> Self:
993 if self.halo and len(self.halo) != len(self.shape):
994 raise ValueError(
995 f"halo {self.halo} has to have same length as shape {self.shape}!"
996 )
998 return self
1000 @model_validator(mode="after")
1001 def validate_postprocessing_kwargs(self) -> Self:
1002 for p in self.postprocessing:
1003 kwargs_axes = p.kwargs.get("axes", "")
1004 if not isinstance(kwargs_axes, str):
1005 raise ValueError(f"Expected {kwargs_axes} to be a string")
1007 if any(a not in self.axes for a in kwargs_axes):
1008 raise ValueError("`kwargs.axes` needs to be subset of axes")
1010 return self
1013KnownRunMode = Literal["deepimagej"]
1016class RunMode(Node):
1017 name: Annotated[
1018 KnownRunMode | str, warn(KnownRunMode, "Unknown run mode '{value}'.")
1019 ]
1020 """Run mode name"""
1022 kwargs: dict[str, Any] = Field(
1023 default_factory=cast(Callable[[], Dict[str, Any]], dict)
1024 )
1025 """Run mode specific key word arguments"""
1028class LinkedModel(Node):
1029 """Reference to a bioimage.io model."""
1031 id: Annotated[ModelId, Field(examples=["affable-shark", "ambitious-sloth"])]
1032 """A valid model `id` from the bioimage.io collection."""
1034 version_number: int | None = None
1035 """version number (n-th published version, not the semantic version) of linked model"""
1038def package_weights(
1039 value: Node, # Union[v0_4.WeightsDescr, v0_5.WeightsDescr]
1040 handler: SerializerFunctionWrapHandler,
1041 info: SerializationInfo,
1042):
1043 ctxt = packaging_context_var.get()
1044 if ctxt is not None and ctxt.weights_priority_order is not None:
1045 for wf in ctxt.weights_priority_order:
1046 w = getattr(value, wf, None)
1047 if w is not None:
1048 break
1049 else:
1050 raise ValueError(
1051 "None of the weight formats in `weights_priority_order`"
1052 + f" ({ctxt.weights_priority_order}) is present in the given model."
1053 )
1055 assert isinstance(w, Node), type(w)
1056 # construct WeightsDescr with new single weight format entry
1057 new_w = w.model_construct(**{k: v for k, v in w if k != "parent"})
1058 value = value.model_construct(None, **{wf: new_w})
1060 return handler(
1061 value,
1062 info, # pyright: ignore[reportArgumentType] # taken from pydantic docs
1063 )
1066class ModelDescr(GenericModelDescrBase):
1067 """Specification of the fields used in a bioimage.io-compliant RDF that describes AI models with pretrained weights.
1069 These fields are typically stored in a YAML file which we call a model resource description file (model RDF).
1070 """
1072 implemented_format_version: ClassVar[Literal["0.4.10"]] = "0.4.10"
1073 if TYPE_CHECKING:
1074 format_version: Literal["0.4.10"] = "0.4.10"
1075 else:
1076 format_version: Literal["0.4.10"]
1077 """Version of the bioimage.io model description specification used.
1078 When creating a new model always use the latest micro/patch version described here.
1079 The `format_version` is important for any consumer software to understand how to parse the fields.
1080 """
1082 implemented_type: ClassVar[Literal["model"]] = "model"
1083 if TYPE_CHECKING:
1084 type: Literal["model"] = "model"
1085 else:
1086 type: Literal["model"]
1087 """Specialized resource type 'model'"""
1089 id: ModelId | None = None
1090 """bioimage.io-wide unique resource identifier
1091 assigned by bioimage.io; version **un**specific."""
1093 authors: NotEmpty[ # pyright: ignore[reportGeneralTypeIssues] # make mandatory
1094 list[Author]
1095 ]
1096 """The authors are the creators of the model RDF and the primary points of contact."""
1098 documentation: Annotated[
1099 FileSource_package,
1100 Field(
1101 examples=[
1102 "https://raw.githubusercontent.com/bioimage-io/spec-bioimage-io/main/example_descriptions/models/unet2d_nuclei_broad/README.md",
1103 "README.md",
1104 ],
1105 ),
1106 ]
1107 """URL or relative path to a markdown file with additional documentation.
1108 The recommended documentation file name is `README.md`. An `.md` suffix is mandatory.
1109 The documentation should include a '[#[#]]# Validation' (sub)section
1110 with details on how to quantitatively validate the model on unseen data."""
1112 inputs: NotEmpty[list[InputTensorDescr]]
1113 """Describes the input tensors expected by this model."""
1115 license: Annotated[
1116 LicenseId | str,
1117 warn(LicenseId, "Unknown license id '{value}'."),
1118 Field(examples=["CC0-1.0", "MIT", "BSD-2-Clause"]),
1119 ]
1120 """A [SPDX license identifier](https://spdx.org/licenses/).
1121 We do notsupport custom license beyond the SPDX license list, if you need that please
1122 [open a GitHub issue](https://github.com/bioimage-io/spec-bioimage-io/issues/new/choose
1123 ) to discuss your intentions with the community."""
1125 name: Annotated[
1126 str,
1127 MinLen(1),
1128 warn(MinLen(5), "Name shorter than 5 characters.", INFO),
1129 warn(MaxLen(64), "Name longer than 64 characters.", INFO),
1130 ]
1131 """A human-readable name of this model.
1132 It should be no longer than 64 characters and only contain letter, number, underscore, minus or space characters."""
1134 outputs: NotEmpty[list[OutputTensorDescr]]
1135 """Describes the output tensors."""
1137 @field_validator("inputs", "outputs")
1138 @classmethod
1139 def unique_tensor_descr_names(
1140 cls, value: Sequence[InputTensorDescr | OutputTensorDescr]
1141 ) -> Sequence[InputTensorDescr | OutputTensorDescr]:
1142 unique_names = {str(v.name) for v in value}
1143 if len(unique_names) != len(value):
1144 raise ValueError("Duplicate tensor descriptor names")
1146 return value
1148 @model_validator(mode="after")
1149 def unique_io_names(self) -> Self:
1150 unique_names = {str(ss.name) for s in (self.inputs, self.outputs) for ss in s}
1151 if len(unique_names) != (len(self.inputs) + len(self.outputs)):
1152 raise ValueError("Duplicate tensor descriptor names across inputs/outputs")
1154 return self
1156 @model_validator(mode="after")
1157 def minimum_shape2valid_output(self) -> Self:
1158 tensors_by_name: dict[TensorName, InputTensorDescr | OutputTensorDescr] = {
1159 t.name: t for t in self.inputs + self.outputs
1160 }
1162 for out in self.outputs:
1163 if isinstance(out.shape, ImplicitOutputShape):
1164 ndim_ref = len(tensors_by_name[out.shape.reference_tensor].shape)
1165 ndim_out_ref = len(
1166 [scale for scale in out.shape.scale if scale is not None]
1167 )
1168 if ndim_ref != ndim_out_ref:
1169 expanded_dim_note = (
1170 " Note that expanded dimensions (`scale`: null) are not"
1171 + f" counted for {out.name}'sdimensionality here."
1172 if None in out.shape.scale
1173 else ""
1174 )
1175 raise ValueError(
1176 f"Referenced tensor '{out.shape.reference_tensor}' with"
1177 + f" {ndim_ref} dimensions does not match output tensor"
1178 + f" '{out.name}' with"
1179 + f" {ndim_out_ref} dimensions.{expanded_dim_note}"
1180 )
1182 min_out_shape = self._get_min_shape(out, tensors_by_name)
1183 if out.halo:
1184 halo = out.halo
1185 halo_msg = f" for halo {out.halo}"
1186 else:
1187 halo = [0] * len(min_out_shape)
1188 halo_msg = ""
1190 if any(s - 2 * h < 1 for s, h in zip(min_out_shape, halo)):
1191 raise ValueError(
1192 f"Minimal shape {min_out_shape} of output {out.name} is too"
1193 + f" small{halo_msg}."
1194 )
1196 return self
1198 @classmethod
1199 def _get_min_shape(
1200 cls,
1201 t: InputTensorDescr | OutputTensorDescr,
1202 tensors_by_name: dict[TensorName, InputTensorDescr | OutputTensorDescr],
1203 ) -> Sequence[int]:
1204 """output with subtracted halo has to result in meaningful output even for the minimal input
1205 see https://github.com/bioimage-io/spec-bioimage-io/issues/392
1206 """
1207 if isinstance(t.shape, collections.abc.Sequence):
1208 return t.shape
1209 elif isinstance(t.shape, ParameterizedInputShape):
1210 return t.shape.min
1211 elif isinstance(t.shape, ImplicitOutputShape):
1212 pass
1213 else:
1214 assert_never(t.shape)
1216 ref_shape = cls._get_min_shape(
1217 tensors_by_name[t.shape.reference_tensor], tensors_by_name
1218 )
1220 if None not in t.shape.scale:
1221 scale: Sequence[float, ...] = t.shape.scale # type: ignore
1222 else:
1223 expanded_dims = [idx for idx, sc in enumerate(t.shape.scale) if sc is None]
1224 new_ref_shape: list[int] = []
1225 for idx in range(len(t.shape.scale)):
1226 ref_idx = idx - sum(int(exp < idx) for exp in expanded_dims)
1227 new_ref_shape.append(1 if idx in expanded_dims else ref_shape[ref_idx])
1229 ref_shape = new_ref_shape
1230 assert len(ref_shape) == len(t.shape.scale)
1231 scale = [0.0 if sc is None else sc for sc in t.shape.scale]
1233 offset = t.shape.offset
1234 assert len(offset) == len(scale)
1235 return [int(rs * s + 2 * off) for rs, s, off in zip(ref_shape, scale, offset)]
1237 @model_validator(mode="after")
1238 def validate_tensor_references_in_inputs(self) -> Self:
1239 for t in self.inputs:
1240 for proc in t.preprocessing:
1241 if "reference_tensor" not in proc.kwargs:
1242 continue
1244 ref_tensor = proc.kwargs["reference_tensor"]
1245 if ref_tensor is not None and str(ref_tensor) not in {
1246 str(t.name) for t in self.inputs
1247 }:
1248 raise ValueError(f"'{ref_tensor}' not found in inputs")
1250 if ref_tensor == t.name:
1251 raise ValueError(
1252 f"invalid self reference for preprocessing of tensor {t.name}"
1253 )
1255 return self
1257 @model_validator(mode="after")
1258 def validate_tensor_references_in_outputs(self) -> Self:
1259 for t in self.outputs:
1260 for proc in t.postprocessing:
1261 if "reference_tensor" not in proc.kwargs:
1262 continue
1263 ref_tensor = proc.kwargs["reference_tensor"]
1264 if ref_tensor is not None and str(ref_tensor) not in {
1265 str(t.name) for t in self.inputs
1266 }:
1267 raise ValueError(f"{ref_tensor} not found in inputs")
1269 return self
1271 packaged_by: list[Author] = Field(
1272 default_factory=cast(Callable[[], List[Author]], list)
1273 )
1274 """The persons that have packaged and uploaded this model.
1275 Only required if those persons differ from the `authors`."""
1277 parent: LinkedModel | None = None
1278 """The model from which this model is derived, e.g. by fine-tuning the weights."""
1280 @field_validator("parent", mode="before")
1281 @classmethod
1282 def ignore_url_parent(cls, parent: Any):
1283 if isinstance(parent, dict):
1284 return None
1286 else:
1287 return parent
1289 run_mode: RunMode | None = None
1290 """Custom run mode for this model: for more complex prediction procedures like test time
1291 data augmentation that currently cannot be expressed in the specification.
1292 No standard run modes are defined yet."""
1294 sample_inputs: list[FileSource_package] = Field(
1295 default_factory=cast(Callable[[], List[FileSource]], list)
1296 )
1297 """URLs/relative paths to sample inputs to illustrate possible inputs for the model,
1298 for example stored as PNG or TIFF images.
1299 The sample files primarily serve to inform a human user about an example use case"""
1301 sample_outputs: list[FileSource_package] = Field(
1302 default_factory=cast(Callable[[], List[FileSource]], list)
1303 )
1304 """URLs/relative paths to sample outputs corresponding to the `sample_inputs`."""
1306 test_inputs: NotEmpty[
1307 list[Annotated[FileSource_package, WithSuffix(".npy", case_sensitive=True)]]
1308 ]
1309 """Test input tensors compatible with the `inputs` description for a **single test case**.
1310 This means if your model has more than one input, you should provide one URL/relative path for each input.
1311 Each test input should be a file with an ndarray in
1312 [numpy.lib file format](https://numpy.org/doc/stable/reference/generated/numpy.lib.format.html#module-numpy.lib.format).
1313 The extension must be '.npy'."""
1315 test_outputs: NotEmpty[
1316 list[Annotated[FileSource_package, WithSuffix(".npy", case_sensitive=True)]]
1317 ]
1318 """Analog to `test_inputs`."""
1320 timestamp: Datetime
1321 """Timestamp in [ISO 8601](#https://en.wikipedia.org/wiki/ISO_8601) format
1322 with a few restrictions listed [here](https://docs.python.org/3/library/datetime.html#datetime.datetime.fromisoformat)."""
1324 training_data: LinkedDataset | DatasetDescr | None = None
1325 """The dataset used to train this model"""
1327 weights: Annotated[WeightsDescr, WrapSerializer(package_weights)]
1328 """The weights for this model.
1329 Weights can be given for different formats, but should otherwise be equivalent.
1330 The available weight formats determine which consumers can use this model."""
1332 @model_validator(mode="before")
1333 @classmethod
1334 def _convert_from_older_format(
1335 cls, data: BioimageioYamlContent, /
1336 ) -> BioimageioYamlContent:
1337 convert_from_older_format(data)
1338 return data
1340 def get_input_test_arrays(self) -> list[NDArray[Any]]:
1341 data = [load_array(ipt) for ipt in self.test_inputs]
1342 assert all(isinstance(d, np.ndarray) for d in data)
1343 return data
1345 def get_output_test_arrays(self) -> list[NDArray[Any]]:
1346 data = [load_array(out) for out in self.test_outputs]
1347 assert all(isinstance(d, np.ndarray) for d in data)
1348 return data