41@dataclass
42class Axis:
43    id: AxisId
44    type: Literal["batch", "channel", "index", "space", "time"]
45
46    @classmethod
47    def create(cls, axis: AxisLike) -> Axis:
48        if isinstance(axis, cls):
49            return axis
50        elif isinstance(axis, Axis):
51            return Axis(id=axis.id, type=axis.type)
52        elif isinstance(axis, str):
53            return Axis(id=AxisId(axis), type=_get_axis_type(axis))
54        elif isinstance(axis, v0_5.AxisBase):
55            return Axis(id=AxisId(axis.id), type=axis.type)
56        else:
57            assert_never(axis)

199class AxisId(LowerCaseIdentifier):
200    root_model: ClassVar[Type[RootModel[Any]]] = RootModel[
201        Annotated[LowerCaseIdentifierAnno, MaxLen(16)]
202    ]

 46@dataclass(frozen=True)
 47class BlockMeta:
 48    """Block meta data of a sample member (a tensor in a sample)
 49
 50    Figure for illustration:
 51    The first 2d block (dashed) of a sample member (**bold**).
 52    The inner slice (thin) is expanded by a halo in both dimensions on both sides.
 53    The outer slice reaches from the sample member origin (0, 0) to the right halo point.
 54
 55    ```terminal
 56    ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─  ─ ─ ─ ─ ─ ─ ─ ┐
 57    ╷ halo(left)                         ╷
 58    ╷                                    ╷
 59    ╷  (0, 0)┏━━━━━━━━━━━━━━━━━┯━━━━━━━━━┯━━━➔
 60    ╷        ┃                 │         ╷  sample member
 61    ╷        ┃      inner      │         ╷
 62    ╷        ┃   (and outer)   │  outer  ╷
 63    ╷        ┃      slice      │  slice  ╷
 64    ╷        ┃                 │         ╷
 65    ╷        ┣─────────────────┘         ╷
 66    ╷        ┃   outer slice             ╷
 67    ╷        ┃               halo(right) ╷
 68    └ ─ ─ ─ ─┃─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─┘
 69             ⬇
 70    ```
 71
 72    note:
 73    - Inner and outer slices are specified in sample member coordinates.
 74    - The outer_slice of a block at the sample edge may overlap by more than the
 75        halo with the neighboring block (the inner slices will not overlap though).
 76
 77    """
 78
 79    sample_shape: PerAxis[int]
 80    """the axis sizes of the whole (unblocked) sample"""
 81
 82    inner_slice: PerAxis[SliceInfo]
 83    """inner region (without halo) wrt the sample"""
 84
 85    halo: PerAxis[Halo]
 86    """halo enlarging the inner region to the block's sizes"""
 87
 88    block_index: BlockIndex
 89    """the i-th block of the sample"""
 90
 91    blocks_in_sample: TotalNumberOfBlocks
 92    """total number of blocks in the sample"""
 93
 94    @cached_property
 95    def shape(self) -> PerAxis[int]:
 96        """axis lengths of the block"""
 97        return Frozen(
 98            {
 99                a: s.stop - s.start + (sum(self.halo[a]) if a in self.halo else 0)
100                for a, s in self.inner_slice.items()
101            }
102        )
103
104    @cached_property
105    def padding(self) -> PerAxis[PadWidth]:
106        """padding to realize the halo at the sample edge
107        where we cannot simply enlarge the inner slice"""
108        return Frozen(
109            {
110                a: PadWidth(
111                    (
112                        self.halo[a].left
113                        - (self.inner_slice[a].start - self.outer_slice[a].start)
114                        if a in self.halo
115                        else 0
116                    ),
117                    (
118                        self.halo[a].right
119                        - (self.outer_slice[a].stop - self.inner_slice[a].stop)
120                        if a in self.halo
121                        else 0
122                    ),
123                )
124                for a in self.inner_slice
125            }
126        )
127
128    @cached_property
129    def outer_slice(self) -> PerAxis[SliceInfo]:
130        """slice of the outer block (without padding) wrt the sample"""
131        return Frozen(
132            {
133                a: SliceInfo(
134                    max(
135                        0,
136                        min(
137                            self.inner_slice[a].start
138                            - (self.halo[a].left if a in self.halo else 0),
139                            self.sample_shape[a]
140                            - self.inner_shape[a]
141                            - (self.halo[a].left if a in self.halo else 0),
142                        ),
143                    ),
144                    min(
145                        self.sample_shape[a],
146                        self.inner_slice[a].stop
147                        + (self.halo[a].right if a in self.halo else 0),
148                    ),
149                )
150                for a in self.inner_slice
151            }
152        )
153
154    @cached_property
155    def inner_shape(self) -> PerAxis[int]:
156        """axis lengths of the inner region (without halo)"""
157        return Frozen({a: s.stop - s.start for a, s in self.inner_slice.items()})
158
159    @cached_property
160    def local_slice(self) -> PerAxis[SliceInfo]:
161        """inner slice wrt the block, **not** the sample"""
162        return Frozen(
163            {
164                a: SliceInfo(
165                    self.halo[a].left,
166                    self.halo[a].left + self.inner_shape[a],
167                )
168                for a in self.inner_slice
169            }
170        )
171
172    @property
173    def dims(self) -> Collection[AxisId]:
174        return set(self.inner_shape)
175
176    @property
177    def tagged_shape(self) -> PerAxis[int]:
178        """alias for shape"""
179        return self.shape
180
181    @property
182    def inner_slice_wo_overlap(self):
183        """subslice of the inner slice, such that all `inner_slice_wo_overlap` can be
184        stiched together trivially to form the original sample.
185
186        This can also be used to calculate statistics
187        without overrepresenting block edge regions."""
188        # TODO: update inner_slice_wo_overlap when adding block overlap
189        return self.inner_slice
190
191    def __post_init__(self):
192        # freeze mutable inputs
193        if not isinstance(self.sample_shape, Frozen):
194            object.__setattr__(self, "sample_shape", Frozen(self.sample_shape))
195
196        if not isinstance(self.inner_slice, Frozen):
197            object.__setattr__(self, "inner_slice", Frozen(self.inner_slice))
198
199        if not isinstance(self.halo, Frozen):
200            object.__setattr__(self, "halo", Frozen(self.halo))
201
202        assert all(
203            a in self.sample_shape for a in self.inner_slice
204        ), "block has axes not present in sample"
205
206        assert all(
207            a in self.inner_slice for a in self.halo
208        ), "halo has axes not present in block"
209
210        if any(s > self.sample_shape[a] for a, s in self.shape.items()):
211            logger.warning(
212                "block {} larger than sample {}", self.shape, self.sample_shape
213            )
214
215    def get_transformed(
216        self, new_axes: PerAxis[Union[LinearAxisTransform, int]]
217    ) -> Self:
218        return self.__class__(
219            sample_shape={
220                a: (
221                    trf
222                    if isinstance(trf, int)
223                    else trf.compute(self.sample_shape[trf.axis])
224                )
225                for a, trf in new_axes.items()
226            },
227            inner_slice={
228                a: (
229                    SliceInfo(0, trf)
230                    if isinstance(trf, int)
231                    else SliceInfo(
232                        trf.compute(self.inner_slice[trf.axis].start),
233                        trf.compute(self.inner_slice[trf.axis].stop),
234                    )
235                )
236                for a, trf in new_axes.items()
237            },
238            halo={
239                a: (
240                    Halo(0, 0)
241                    if isinstance(trf, int)
242                    else Halo(self.halo[trf.axis].left, self.halo[trf.axis].right)
243                )
244                for a, trf in new_axes.items()
245            },
246            block_index=self.block_index,
247            blocks_in_sample=self.blocks_in_sample,
248        )

130def build_description(
131    content: BioimageioYamlContent,
132    /,
133    *,
134    context: Optional[ValidationContext] = None,
135    format_version: Union[FormatVersionPlaceholder, str] = DISCOVER,
136) -> Union[ResourceDescr, InvalidDescr]:
137    """build a bioimage.io resource description from an RDF's content.
138
139    Use `load_description` if you want to build a resource description from an rdf.yaml
140    or bioimage.io zip-package.
141
142    Args:
143        content: loaded rdf.yaml file (loaded with YAML, not bioimageio.spec)
144        context: validation context to use during validation
145        format_version: (optional) use this argument to load the resource and
146                        convert its metadata to a higher format_version
147
148    Returns:
149        An object holding all metadata of the bioimage.io resource
150
151    """
152
153    return build_description_impl(
154        content,
155        context=context,
156        format_version=format_version,
157        get_rd_class=_get_rd_class,
158    )

560def compute_dataset_measures(
561    measures: Iterable[DatasetMeasure], dataset: Iterable[Sample]
562) -> Dict[DatasetMeasure, MeasureValue]:
563    """compute all dataset `measures` for the given `dataset`"""
564    sample_calculators, calculators = get_measure_calculators(measures)
565    assert not sample_calculators
566
567    ret: Dict[DatasetMeasure, MeasureValue] = {}
568
569    for sample in dataset:
570        for calc in calculators:
571            calc.update(sample)
572
573    for calc in calculators:
574        ret.update(calc.finalize().items())
575
576    return ret

319def create_prediction_pipeline(
320    bioimageio_model: AnyModelDescr,
321    *,
322    devices: Optional[Sequence[str]] = None,
323    weight_format: Optional[WeightsFormat] = None,
324    weights_format: Optional[WeightsFormat] = None,
325    dataset_for_initial_statistics: Iterable[Union[Sample, Sequence[Tensor]]] = tuple(),
326    keep_updating_initial_dataset_statistics: bool = False,
327    fixed_dataset_statistics: Mapping[DatasetMeasure, MeasureValue] = MappingProxyType(
328        {}
329    ),
330    model_adapter: Optional[ModelAdapter] = None,
331    ns: Union[
332        v0_5.ParameterizedSize_N,
333        Mapping[Tuple[MemberId, AxisId], v0_5.ParameterizedSize_N],
334    ] = 10,
335    **deprecated_kwargs: Any,
336) -> PredictionPipeline:
337    """
338    Creates prediction pipeline which includes:
339    * computation of input statistics
340    * preprocessing
341    * model prediction
342    * computation of output statistics
343    * postprocessing
344    """
345    weights_format = weight_format or weights_format
346    del weight_format
347    if deprecated_kwargs:
348        warnings.warn(
349            f"deprecated create_prediction_pipeline kwargs: {set(deprecated_kwargs)}"
350        )
351
352    model_adapter = model_adapter or create_model_adapter(
353        model_description=bioimageio_model,
354        devices=devices,
355        weight_format_priority_order=weights_format and (weights_format,),
356    )
357
358    input_ids = get_member_ids(bioimageio_model.inputs)
359
360    def dataset():
361        common_stat: Stat = {}
362        for i, x in enumerate(dataset_for_initial_statistics):
363            if isinstance(x, Sample):
364                yield x
365            else:
366                yield Sample(members=dict(zip(input_ids, x)), stat=common_stat, id=i)
367
368    preprocessing, postprocessing = setup_pre_and_postprocessing(
369        bioimageio_model,
370        dataset(),
371        keep_updating_initial_dataset_stats=keep_updating_initial_dataset_statistics,
372        fixed_dataset_stats=fixed_dataset_statistics,
373    )
374
375    return PredictionPipeline(
376        name=bioimageio_model.name,
377        model_description=bioimageio_model,
378        model_adapter=model_adapter,
379        preprocessing=preprocessing,
380        postprocessing=postprocessing,
381        default_ns=ns,
382    )

65def dump_description(
66    rd: Union[ResourceDescr, InvalidDescr], exclude_unset: bool = True
67) -> BioimageioYamlContent:
68    """Converts a resource to a dictionary containing only simple types that can directly be serialzed to YAML."""
69    return rd.model_dump(mode="json", exclude_unset=exclude_unset)

36def enable_determinism(mode: Literal["seed_only", "full"]):
37    """Seed and configure ML frameworks for maximum reproducibility.
38    May degrade performance. Only recommended for testing reproducibility!
39
40    Seed any random generators and (if **mode**=="full") request ML frameworks to use
41    deterministic algorithms.
42    Notes:
43        - **mode** == "full"  might degrade performance and throw exceptions.
44        - Subsequent inference calls might still differ. Call before each function
45          (sequence) that is expected to be reproducible.
46        - Degraded performance: Use for testing reproducibility only!
47        - Recipes:
48            - [PyTorch](https://pytorch.org/docs/stable/notes/randomness.html)
49            - [Keras](https://keras.io/examples/keras_recipes/reproducibility_recipes/)
50            - [NumPy](https://numpy.org/doc/2.0/reference/random/generated/numpy.random.seed.html)
51    """
52    try:
53        try:
54            import numpy.random
55        except ImportError:
56            pass
57        else:
58            numpy.random.seed(0)
59    except Exception as e:
60        logger.debug(str(e))
61
62    try:
63        try:
64            import torch
65        except ImportError:
66            pass
67        else:
68            _ = torch.manual_seed(0)
69            torch.use_deterministic_algorithms(mode == "full")
70    except Exception as e:
71        logger.debug(str(e))
72
73    try:
74        try:
75            import keras
76        except ImportError:
77            pass
78        else:
79            keras.utils.set_random_seed(0)
80    except Exception as e:
81        logger.debug(str(e))
82
83    try:
84        try:
85            import tensorflow as tf  # pyright: ignore[reportMissingImports]
86        except ImportError:
87            pass
88        else:
89            tf.random.seed(0)
90            if mode == "full":
91                tf.config.experimental.enable_op_determinism()
92            # TODO: find possibility to switch it off again??
93    except Exception as e:
94        logger.debug(str(e))

 98def load_dataset_description(
 99    source: Union[PermissiveFileSource, ZipFile],
100    /,
101    *,
102    format_version: Union[Literal["discover"], Literal["latest"], str] = DISCOVER,
103    perform_io_checks: bool = settings.perform_io_checks,
104    known_files: Optional[Dict[str, Sha256]] = None,
105) -> AnyDatasetDescr:
106    """same as `load_description`, but addtionally ensures that the loaded
107    description is valid and of type 'dataset'.
108    """
109    rd = load_description(
110        source,
111        format_version=format_version,
112        perform_io_checks=perform_io_checks,
113        known_files=known_files,
114    )
115    return ensure_description_is_dataset(rd)

147def load_description_and_test(
148    source: Union[ResourceDescr, PermissiveFileSource, BioimageioYamlContent],
149    *,
150    format_version: Union[Literal["discover", "latest"], str] = "discover",
151    weight_format: Optional[WeightsFormat] = None,
152    devices: Optional[Sequence[str]] = None,
153    absolute_tolerance: float = 1.5e-4,
154    relative_tolerance: float = 1e-4,
155    decimal: Optional[int] = None,
156    determinism: Literal["seed_only", "full"] = "seed_only",
157    expected_type: Optional[str] = None,
158) -> Union[ResourceDescr, InvalidDescr]:
159    """Test RDF dynamically, e.g. model inference of test inputs"""
160    if (
161        isinstance(source, ResourceDescrBase)
162        and format_version != "discover"
163        and source.format_version != format_version
164    ):
165        warnings.warn(
166            f"deserializing source to ensure we validate and test using format {format_version}"
167        )
168        source = dump_description(source)
169
170    if isinstance(source, ResourceDescrBase):
171        rd = source
172    elif isinstance(source, dict):
173        rd = build_description(source, format_version=format_version)
174    else:
175        rd = load_description(source, format_version=format_version)
176
177    rd.validation_summary.env.add(
178        InstalledPackage(name="bioimageio.core", version=VERSION)
179    )
180
181    if expected_type is not None:
182        _test_expected_resource_type(rd, expected_type)
183
184    if isinstance(rd, (v0_4.ModelDescr, v0_5.ModelDescr)):
185        if weight_format is None:
186            weight_formats: List[WeightsFormat] = [
187                w for w, we in rd.weights if we is not None
188            ]  # pyright: ignore[reportAssignmentType]
189        else:
190            weight_formats = [weight_format]
191
192        if decimal is None:
193            atol = absolute_tolerance
194            rtol = relative_tolerance
195        else:
196            warnings.warn(
197                "The argument `decimal` has been deprecated in favour of"
198                + " `relative_tolerance` and `absolute_tolerance`, with different"
199                + " validation logic, using `numpy.testing.assert_allclose, see"
200                + " 'https://numpy.org/doc/stable/reference/generated/"
201                + " numpy.testing.assert_allclose.html'. Passing a value for `decimal`"
202                + " will cause validation to revert to the old behaviour."
203            )
204            atol = 1.5 * 10 ** (-decimal)
205            rtol = 0
206
207        enable_determinism(determinism)
208        for w in weight_formats:
209            _test_model_inference(rd, w, devices, atol, rtol)
210            if not isinstance(rd, v0_4.ModelDescr):
211                _test_model_inference_parametrized(rd, w, devices)
212
213    # TODO: add execution of jupyter notebooks
214    # TODO: add more tests
215
216    return rd

137def load_description_and_validate_format_only(
138    source: Union[PermissiveFileSource, ZipFile],
139    /,
140    *,
141    format_version: Union[Literal["discover"], Literal["latest"], str] = DISCOVER,
142    perform_io_checks: bool = settings.perform_io_checks,
143    known_files: Optional[Dict[str, Sha256]] = None,
144) -> ValidationSummary:
145    """load a bioimage.io resource description
146
147    Args:
148        source: Path or URL to an rdf.yaml or a bioimage.io package
149                (zip-file with rdf.yaml in it).
150        format_version: (optional) Use this argument to load the resource and
151                        convert its metadata to a higher format_version.
152        perform_io_checks: Wether or not to perform validation that requires file io,
153                           e.g. downloading a remote files. The existence of local
154                           absolute file paths is still being checked.
155        known_files: Allows to bypass download and hashing of referenced files
156                     (even if perform_io_checks is True).
157
158    Returns:
159        Validation summary of the bioimage.io resource found at `source`.
160
161    """
162    rd = load_description(
163        source,
164        format_version=format_version,
165        perform_io_checks=perform_io_checks,
166        known_files=known_files,
167    )
168    assert rd.validation_summary is not None
169    return rd.validation_summary