18def add_weights(
 19    model_descr: ModelDescr,
 20    *,
 21    output_path: DirectoryPath,
 22    source_format: Optional[WeightsFormat] = None,
 23    target_format: Optional[WeightsFormat] = None,
 24    verbose: bool = False,
 25) -> Optional[ModelDescr]:
 26    """Convert model weights to other formats and add them to the model description
 27
 28    Args:
 29        output_path: Path to save updated model package to.
 30        source_format: convert from a specific weights format.
 31                       Default: choose automatically from any available.
 32        target_format: convert to a specific weights format.
 33                       Default: attempt to convert to any missing format.
 34        devices: Devices that may be used during conversion.
 35        verbose: log more (error) output
 36
 37    Returns:
 38        - An updated model description if any converted weights were added.
 39        - `None` if no conversion was possible.
 40    """
 41    if not isinstance(model_descr, ModelDescr):
 42        if model_descr.type == "model" and not isinstance(model_descr, InvalidDescr):
 43            raise TypeError(
 44                f"Model format {model_descr.format} is not supported, please update"
 45                + f" model to format {ModelDescr.implemented_format_version} first."
 46            )
 47
 48        raise TypeError(type(model_descr))
 49
 50    # save model to local folder
 51    output_path = save_bioimageio_package_as_folder(
 52        model_descr, output_path=output_path
 53    )
 54    # reload from local folder to make sure we do not edit the given model
 55    _model_descr = load_model_description(output_path, perform_io_checks=False)
 56    assert isinstance(_model_descr, ModelDescr)
 57    model_descr = _model_descr
 58    del _model_descr
 59
 60    if source_format is None:
 61        available = set(model_descr.weights.available_formats)
 62    else:
 63        available = {source_format}
 64
 65    if target_format is None:
 66        missing = set(model_descr.weights.missing_formats)
 67    else:
 68        missing = {target_format}
 69
 70    originally_missing = set(missing)
 71
 72    if "pytorch_state_dict" in available and "torchscript" in missing:
 73        logger.info(
 74            "Attempting to convert 'pytorch_state_dict' weights to 'torchscript'."
 75        )
 76        from .pytorch_to_torchscript import convert
 77
 78        try:
 79            torchscript_weights_path = output_path / "weights_torchscript.pt"
 80            model_descr.weights.torchscript = convert(
 81                model_descr,
 82                output_path=torchscript_weights_path,
 83                use_tracing=False,
 84            )
 85        except Exception as e:
 86            if verbose:
 87                traceback.print_exception(e)
 88
 89            logger.error(e)
 90        else:
 91            available.add("torchscript")
 92            missing.discard("torchscript")
 93
 94    if "pytorch_state_dict" in available and "torchscript" in missing:
 95        logger.info(
 96            "Attempting to convert 'pytorch_state_dict' weights to 'torchscript' by tracing."
 97        )
 98        from .pytorch_to_torchscript import convert
 99
100        try:
101            torchscript_weights_path = output_path / "weights_torchscript_traced.pt"
102
103            model_descr.weights.torchscript = convert(
104                model_descr,
105                output_path=torchscript_weights_path,
106                use_tracing=True,
107            )
108        except Exception as e:
109            if verbose:
110                traceback.print_exception(e)
111
112            logger.error(e)
113        else:
114            available.add("torchscript")
115            missing.discard("torchscript")
116
117    if "torchscript" in available and "onnx" in missing:
118        logger.info("Attempting to convert 'torchscript' weights to 'onnx'.")
119        from .torchscript_to_onnx import convert
120
121        try:
122            onnx_weights_path = output_path / "weights.onnx"
123            model_descr.weights.onnx = convert(
124                model_descr,
125                output_path=onnx_weights_path,
126            )
127        except Exception as e:
128            if verbose:
129                traceback.print_exception(e)
130
131            logger.error(e)
132        else:
133            available.add("onnx")
134            missing.discard("onnx")
135
136    if "pytorch_state_dict" in available and "onnx" in missing:
137        logger.info("Attempting to convert 'pytorch_state_dict' weights to 'onnx'.")
138        from .pytorch_to_onnx import convert
139
140        try:
141            onnx_weights_path = output_path / "weights.onnx"
142
143            model_descr.weights.onnx = convert(
144                model_descr,
145                output_path=onnx_weights_path,
146                verbose=verbose,
147            )
148        except Exception as e:
149            if verbose:
150                traceback.print_exception(e)
151
152            logger.error(e)
153        else:
154            available.add("onnx")
155            missing.discard("onnx")
156
157    if missing:
158        logger.warning(
159            f"Converting from any of the available weights formats {available} to any"
160            + f" of {missing} failed or is not yet implemented. Please create an issue"
161            + " at https://github.com/bioimage-io/core-bioimage-io-python/issues/new/choose"
162            + " if you would like bioimageio.core to support a particular conversion."
163        )
164
165    if originally_missing == missing:
166        logger.warning("failed to add any converted weights")
167        return None
168    else:
169        logger.info("added weights formats {}", originally_missing - missing)
170        # resave model with updated rdf.yaml
171        _ = save_bioimageio_package_as_folder(model_descr, output_path=output_path)
172        tested_model_descr = load_description_and_test(model_descr)
173        assert isinstance(tested_model_descr, ModelDescr)
174        return tested_model_descr

49@dataclass
50class Axis:
51    id: AxisId
52    type: Literal["batch", "channel", "index", "space", "time"]
53
54    def __post_init__(self):
55        if self.type == "batch":
56            self.id = AxisId("batch")
57        elif self.type == "channel":
58            self.id = AxisId("channel")
59
60    @classmethod
61    def create(cls, axis: AxisLike) -> Axis:
62        if isinstance(axis, cls):
63            return axis
64        elif isinstance(axis, Axis):
65            return Axis(id=axis.id, type=axis.type)
66        elif isinstance(axis, v0_5.AxisBase):
67            return Axis(id=AxisId(axis.id), type=axis.type)
68        elif isinstance(axis, str):
69            return Axis(id=AxisId(axis), type=_guess_axis_type(axis))
70        else:
71            assert_never(axis)

242class AxisId(LowerCaseIdentifier):
243    root_model: ClassVar[Type[RootModel[Any]]] = RootModel[
244        Annotated[
245            LowerCaseIdentifierAnno,
246            MaxLen(16),
247            AfterValidator(_normalize_axis_id),
248        ]
249    ]

 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        )

173def build_description(
174    content: BioimageioYamlContentView,
175    /,
176    *,
177    context: Optional[ValidationContext] = None,
178    format_version: Union[FormatVersionPlaceholder, str] = DISCOVER,
179) -> Union[ResourceDescr, InvalidDescr]:
180    """build a bioimage.io resource description from an RDF's content.
181
182    Use `load_description` if you want to build a resource description from an rdf.yaml
183    or bioimage.io zip-package.
184
185    Args:
186        content: loaded rdf.yaml file (loaded with YAML, not bioimageio.spec)
187        context: validation context to use during validation
188        format_version: (optional) use this argument to load the resource and
189                        convert its metadata to a higher format_version
190
191    Returns:
192        An object holding all metadata of the bioimage.io resource
193
194    """
195
196    return build_description_impl(
197        content,
198        context=context,
199        format_version=format_version,
200        get_rd_class=_get_rd_class,
201    )

578def compute_dataset_measures(
579    measures: Iterable[DatasetMeasure], dataset: Iterable[Sample]
580) -> Dict[DatasetMeasure, MeasureValue]:
581    """compute all dataset `measures` for the given `dataset`"""
582    sample_calculators, calculators = get_measure_calculators(measures)
583    assert not sample_calculators
584
585    ret: Dict[DatasetMeasure, MeasureValue] = {}
586
587    for sample in dataset:
588        for calc in calculators:
589            calc.update(sample)
590
591    for calc in calculators:
592        ret.update(calc.finalize().items())
593
594    return ret

 72    @final
 73    @classmethod
 74    def create(
 75        cls,
 76        model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr],
 77        *,
 78        devices: Optional[Sequence[str]] = None,
 79        weight_format_priority_order: Optional[Sequence[SupportedWeightsFormat]] = None,
 80    ):
 81        """
 82        Creates model adapter based on the passed spec
 83        Note: All specific adapters should happen inside this function to prevent different framework
 84        initializations interfering with each other
 85        """
 86        if not isinstance(model_description, (v0_4.ModelDescr, v0_5.ModelDescr)):
 87            raise TypeError(
 88                f"expected v0_4.ModelDescr or v0_5.ModelDescr, but got {type(model_description)}"
 89            )
 90
 91        weights = model_description.weights
 92        errors: List[Exception] = []
 93        weight_format_priority_order = (
 94            DEFAULT_WEIGHT_FORMAT_PRIORITY_ORDER
 95            if weight_format_priority_order is None
 96            else weight_format_priority_order
 97        )
 98        # limit weight formats to the ones present
 99        weight_format_priority_order_present: Sequence[SupportedWeightsFormat] = [
100            w for w in weight_format_priority_order if getattr(weights, w) is not None
101        ]
102        if not weight_format_priority_order_present:
103            raise ValueError(
104                f"None of the specified weight formats ({weight_format_priority_order}) is present ({weight_format_priority_order_present})"
105            )
106
107        for wf in weight_format_priority_order_present:
108            if wf == "pytorch_state_dict":
109                assert weights.pytorch_state_dict is not None
110                try:
111                    from .pytorch_backend import PytorchModelAdapter
112
113                    return PytorchModelAdapter(
114                        model_description=model_description, devices=devices
115                    )
116                except Exception as e:
117                    errors.append(e)
118            elif wf == "tensorflow_saved_model_bundle":
119                assert weights.tensorflow_saved_model_bundle is not None
120                try:
121                    from .tensorflow_backend import create_tf_model_adapter
122
123                    return create_tf_model_adapter(
124                        model_description=model_description, devices=devices
125                    )
126                except Exception as e:
127                    errors.append(e)
128            elif wf == "onnx":
129                assert weights.onnx is not None
130                try:
131                    from .onnx_backend import ONNXModelAdapter
132
133                    return ONNXModelAdapter(
134                        model_description=model_description, devices=devices
135                    )
136                except Exception as e:
137                    errors.append(e)
138            elif wf == "torchscript":
139                assert weights.torchscript is not None
140                try:
141                    from .torchscript_backend import TorchscriptModelAdapter
142
143                    return TorchscriptModelAdapter(
144                        model_description=model_description, devices=devices
145                    )
146                except Exception as e:
147                    errors.append(e)
148            elif wf == "keras_hdf5":
149                assert weights.keras_hdf5 is not None
150                # keras can either be installed as a separate package or used as part of tensorflow
151                # we try to first import the keras model adapter using the separate package and,
152                # if it is not available, try to load the one using tf
153                try:
154                    try:
155                        from .keras_backend import KerasModelAdapter
156                    except Exception:
157                        from .tensorflow_backend import KerasModelAdapter
158
159                    return KerasModelAdapter(
160                        model_description=model_description, devices=devices
161                    )
162                except Exception as e:
163                    errors.append(e)
164            else:
165                assert_never(wf)
166
167        assert errors
168        if len(weight_format_priority_order) == 1:
169            assert len(errors) == 1
170            raise errors[0]
171
172        else:
173            msg = (
174                "None of the weight format specific model adapters could be created"
175                + " in this environment."
176            )
177            if sys.version_info[:2] >= (3, 11):
178                raise ExceptionGroup(msg, errors)
179            else:
180                raise ValueError(msg) from Exception(errors)

317def create_prediction_pipeline(
318    bioimageio_model: AnyModelDescr,
319    *,
320    devices: Optional[Sequence[str]] = None,
321    weight_format: Optional[SupportedWeightsFormat] = None,
322    weights_format: Optional[SupportedWeightsFormat] = None,
323    dataset_for_initial_statistics: Iterable[Union[Sample, Sequence[Tensor]]] = tuple(),
324    keep_updating_initial_dataset_statistics: bool = False,
325    fixed_dataset_statistics: Mapping[DatasetMeasure, MeasureValue] = MappingProxyType(
326        {}
327    ),
328    model_adapter: Optional[ModelAdapter] = None,
329    ns: Optional[BlocksizeParameter] = None,
330    default_blocksize_parameter: BlocksizeParameter = 10,
331    **deprecated_kwargs: Any,
332) -> PredictionPipeline:
333    """
334    Creates prediction pipeline which includes:
335    * computation of input statistics
336    * preprocessing
337    * model prediction
338    * computation of output statistics
339    * postprocessing
340
341    Args:
342        bioimageio_model: A bioimageio model description.
343        devices: (optional)
344        weight_format: deprecated in favor of **weights_format**
345        weights_format: (optional) Use a specific **weights_format** rather than
346            choosing one automatically.
347            A corresponding `bioimageio.core.model_adapters.ModelAdapter` will be
348            created to run inference with the **bioimageio_model**.
349        dataset_for_initial_statistics: (optional) If preprocessing steps require input
350            dataset statistics, **dataset_for_initial_statistics** allows you to
351            specifcy a dataset from which these statistics are computed.
352        keep_updating_initial_dataset_statistics: (optional) Set to `True` if you want
353            to update dataset statistics with each processed sample.
354        fixed_dataset_statistics: (optional) Allows you to specify a mapping of
355            `DatasetMeasure`s to precomputed `MeasureValue`s.
356        model_adapter: (optional) Allows you to use a custom **model_adapter** instead
357            of creating one according to the present/selected **weights_format**.
358        ns: deprecated in favor of **default_blocksize_parameter**
359        default_blocksize_parameter: Allows to control the default block size for
360            blockwise predictions, see `BlocksizeParameter`.
361
362    """
363    weights_format = weight_format or weights_format
364    del weight_format
365    default_blocksize_parameter = ns or default_blocksize_parameter
366    del ns
367    if deprecated_kwargs:
368        warnings.warn(
369            f"deprecated create_prediction_pipeline kwargs: {set(deprecated_kwargs)}"
370        )
371
372    model_adapter = model_adapter or create_model_adapter(
373        model_description=bioimageio_model,
374        devices=devices,
375        weight_format_priority_order=weights_format and (weights_format,),
376    )
377
378    input_ids = get_member_ids(bioimageio_model.inputs)
379
380    def dataset():
381        common_stat: Stat = {}
382        for i, x in enumerate(dataset_for_initial_statistics):
383            if isinstance(x, Sample):
384                yield x
385            else:
386                yield Sample(members=dict(zip(input_ids, x)), stat=common_stat, id=i)
387
388    preprocessing, postprocessing = setup_pre_and_postprocessing(
389        bioimageio_model,
390        dataset(),
391        keep_updating_initial_dataset_stats=keep_updating_initial_dataset_statistics,
392        fixed_dataset_stats=fixed_dataset_statistics,
393    )
394
395    return PredictionPipeline(
396        name=bioimageio_model.name,
397        model_description=bioimageio_model,
398        model_adapter=model_adapter,
399        preprocessing=preprocessing,
400        postprocessing=postprocessing,
401        default_blocksize_parameter=default_blocksize_parameter,
402    )