add convenience functions for prediction

FynnBe · FynnBe · commit 0853f96690a2 · 2024-05-27T13:17:57.000+02:00
diff --git a/bioimageio/core/__init__.py b/bioimageio/core/__init__.py
@@ -27,6 +27,8 @@
 from .axis import AxisId as AxisId
 from .block_meta import BlockMeta as BlockMeta
 from .common import MemberId as MemberId
+from .prediction import predict as predict
+from .prediction import predict_many as predict_many
 from .sample import Sample as Sample
 from .tensor import Tensor as Tensor
 from .utils import VERSION
diff --git a/bioimageio/core/digest_spec.py b/bioimageio/core/digest_spec.py
@@ -2,6 +2,7 @@
 
 import importlib.util
 from itertools import chain
+from pathlib import Path
 from typing import (
     Any,
     Callable,
@@ -16,6 +17,7 @@
     Union,
 )
 
+import xarray as xr
 from numpy.typing import NDArray
 from typing_extensions import Unpack, assert_never
 
@@ -33,6 +35,7 @@
 from .axis import AxisId, AxisInfo, PerAxis
 from .block_meta import split_multiple_shapes_into_blocks
 from .common import Halo, MemberId, PerMember, SampleId, TotalNumberOfBlocks
+from .io import get_tensor
 from .sample import (
     LinearSampleAxisTransform,
     Sample,
@@ -334,7 +337,9 @@ def create_sample_for_model(
     *,
     stat: Optional[Stat] = None,
     sample_id: SampleId = None,
-    inputs: Optional[PerMember[NDArray[Any]]] = None,  # TODO: make non-optional
+    inputs: Optional[
+        PerMember[Union[Tensor, xr.DataArray, NDArray[Any], Path]]
+    ] = None,  # TODO: make non-optional
     **kwargs: NDArray[Any],  # TODO: deprecate in favor of `inputs`
 ) -> Sample:
     """Create a sample from a single set of input(s) for a specific bioimage.io model
@@ -359,7 +364,7 @@ def create_sample_for_model(
 
     return Sample(
         members={
-            m: Tensor.from_numpy(inputs[m], dims=get_axes_infos(ipt))
+            m: get_tensor(inputs[m], ipt)
             for m, ipt in model_inputs.items()
             if m in inputs
         },
diff --git a/bioimageio/core/io.py b/bioimageio/core/io.py
@@ -1,12 +1,15 @@
 from pathlib import Path
-from typing import Any, Dict, Optional, Sequence
+from typing import Any, Dict, Optional, Sequence, Union
 
 import imageio
+import numpy as np
+import xarray as xr
 from loguru import logger
 from numpy.typing import NDArray
+from typing_extensions import assert_never
 
-from bioimageio.spec.model import AnyModelDescr
-from bioimageio.spec.utils import load_array
+from bioimageio.spec.model import AnyModelDescr, v0_4, v0_5
+from bioimageio.spec.utils import load_array, save_array
 
 from .axis import Axis, AxisLike
 from .common import MemberId, PerMember, SampleId
@@ -26,6 +29,7 @@ def load_image(path: Path, is_volume: bool) -> NDArray[Any]:
 
 
 def load_tensor(path: Path, axes: Optional[Sequence[AxisLike]] = None) -> Tensor:
+    # TODO: load axis meta data
     array = load_image(
         path,
         is_volume=(
@@ -36,6 +40,50 @@ def load_tensor(path: Path, axes: Optional[Sequence[AxisLike]] = None) -> Tensor
     return Tensor.from_numpy(array, dims=axes)
 
 
+def get_tensor(
+    src: Union[Tensor, xr.DataArray, NDArray[Any], Path],
+    ipt: Union[v0_4.InputTensorDescr, v0_5.InputTensorDescr],
+):
+    """helper to cast/load various tensor sources"""
+
+    if isinstance(src, Tensor):
+        return src
+
+    if isinstance(src, xr.DataArray):
+        return Tensor.from_xarray(src)
+
+    if isinstance(src, np.ndarray):
+        return Tensor.from_numpy(src, dims=get_axes_infos(ipt))
+
+    if isinstance(src, Path):
+        return load_tensor(src, axes=get_axes_infos(ipt))
+
+    assert_never(src)
+
+
+def save_tensor(path: Path, tensor: Tensor) -> None:
+    # TODO: save axis meta data
+    data: NDArray[Any] = tensor.data.to_numpy()
+    if path.suffix == ".npy":
+        save_array(path, data)
+    else:
+        imageio.volwrite(path, data)
+
+
+def save_sample(path: Union[Path, str], sample: Sample) -> None:
+    """save a sample to path
+
+    `path` must contain `{member_id}` and may contain `{sample_id}`,
+    which are resolved with the `sample` object.
+    """
+    path = str(path).format(sample_id=sample.id)
+    if "{member_id}" not in path:
+        raise ValueError(f"missing `{{member_id}}` in path {path}")
+
+    for m, t in sample.members.items():
+        save_tensor(Path(path.format(member_id=m)), t)
+
+
 def load_sample_for_model(
     *,
     model: AnyModelDescr,
diff --git a/bioimageio/core/prediction.py b/bioimageio/core/prediction.py
@@ -4,4 +4,177 @@
 e..g load samples with core.io.load_sample_for_model()
 """
 
-# TODO: add convenience functions for predictions
+import collections
+from pathlib import Path
+from typing import (
+    Any,
+    Generator,
+    Hashable,
+    Iterable,
+    Iterator,
+    List,
+    Mapping,
+    Optional,
+    Tuple,
+    Union,
+)
+
+import xarray as xr
+from numpy.typing import NDArray
+from tqdm import tqdm
+
+from bioimageio.core.axis import AxisId
+from bioimageio.core.io import save_sample
+from bioimageio.spec import load_description
+from bioimageio.spec.common import PermissiveFileSource
+from bioimageio.spec.model import v0_4, v0_5
+
+from ._prediction_pipeline import PredictionPipeline, create_prediction_pipeline
+from .common import MemberId, PerMember
+from .digest_spec import create_sample_for_model
+from .sample import Sample
+from .tensor import Tensor
+
+
+def predict(
+    *,
+    model: Union[
+        PermissiveFileSource, v0_4.ModelDescr, v0_5.ModelDescr, PredictionPipeline
+    ],
+    inputs: PerMember[Union[Tensor, xr.DataArray, NDArray[Any], Path]],
+    sample_id: Hashable = "sample",
+    blocksize_parameter: Optional[
+        Union[
+            v0_5.ParameterizedSize.N,
+            Mapping[Tuple[MemberId, AxisId], v0_5.ParameterizedSize.N],
+        ]
+    ] = None,
+    skip_preprocessing: bool = False,
+    skip_postprocessing: bool = False,
+    save_output_path: Optional[Union[Path, str]] = None,
+) -> Sample:
+    """Run prediction for a single set of input(s) with a bioimage.io model
+
+    Args:
+        model: model to predict with.
+            May be given as RDF source, model description or prediction pipeline.
+        inputs: the named input(s) for this model as a dictionary
+        sample_id: the sample id.
+        blocksize_parameter: (optional) tile the input into blocks parametrized by
+            blocksize according to any parametrized axis sizes defined in the model RDF
+        skip_preprocessing: flag to skip the model's preprocessing
+        skip_postprocessing: flag to skip the model's postprocessing
+        save_output_path: A path with `{member_id}` `{sample_id}` in it
+            to save the output to.
+    """
+    if save_output_path is not None:
+        if "{member_id}" not in str(save_output_path):
+            raise ValueError(
+                f"Missing `{{member_id}}` in save_output_path={save_output_path}"
+            )
+
+    if isinstance(model, PredictionPipeline):
+        pp = model
+    else:
+        if not isinstance(model, (v0_4.ModelDescr, v0_5.ModelDescr)):
+            loaded = load_description(model)
+            if not isinstance(loaded, (v0_4.ModelDescr, v0_5.ModelDescr)):
+                raise ValueError(f"expected model description, but got {loaded}")
+            model = loaded
+
+        pp = create_prediction_pipeline(model)
+
+    sample = create_sample_for_model(
+        pp.model_description, inputs=inputs, sample_id=sample_id
+    )
+
+    if blocksize_parameter is None:
+        output = pp.predict_sample_without_blocking(
+            sample,
+            skip_preprocessing=skip_preprocessing,
+            skip_postprocessing=skip_postprocessing,
+        )
+    else:
+        output = pp.predict_sample_with_blocking(
+            sample,
+            skip_preprocessing=skip_preprocessing,
+            skip_postprocessing=skip_postprocessing,
+            ns=blocksize_parameter,
+        )
+    if save_output_path:
+        save_sample(save_output_path, output)
+
+    return output
+
+
+def predict_many(
+    *,
+    model: Union[
+        PermissiveFileSource, v0_4.ModelDescr, v0_5.ModelDescr, PredictionPipeline
+    ],
+    inputs: Iterable[PerMember[Union[Tensor, xr.DataArray, NDArray[Any], Path]]],
+    sample_id: str = "sample{i:03}",
+    blocksize_parameter: Optional[
+        Union[
+            v0_5.ParameterizedSize.N,
+            Mapping[Tuple[MemberId, AxisId], v0_5.ParameterizedSize.N],
+        ]
+    ] = None,
+    skip_preprocessing: bool = False,
+    skip_postprocessing: bool = False,
+    save_output_path: Optional[Union[Path, str]] = None,
+) -> Iterator[Sample]:
+    """Run prediction for a multiple sets of inputs with a bioimage.io model
+
+    Args:
+        model: model to predict with.
+            May be given as RDF source, model description or prediction pipeline.
+        inputs: An iterable of the named input(s) for this model as a dictionary.
+        sample_id: the sample id.
+            note: `{i}` will be formatted as the i-th sample.
+            If `{i}` (or `{i:`) is not present and `inputs` is an iterable `{i:03}` is appended.
+        blocksize_parameter: (optional) tile the input into blocks parametrized by
+            blocksize according to any parametrized axis sizes defined in the model RDF
+        skip_preprocessing: flag to skip the model's preprocessing
+        skip_postprocessing: flag to skip the model's postprocessing
+        save_output_path: A path with `{member_id}` `{sample_id}` in it
+            to save the output to.
+    """
+    if save_output_path is not None:
+        if "{member_id}" not in str(save_output_path):
+            raise ValueError(
+                f"Missing `{{member_id}}` in save_output_path={save_output_path}"
+            )
+
+        if not isinstance(inputs, collections.Mapping) and "{sample_id}" not in str(
+            save_output_path
+        ):
+            raise ValueError(
+                f"Missing `{{sample_id}}` in save_output_path={save_output_path}"
+            )
+
+    if isinstance(model, PredictionPipeline):
+        pp = model
+    else:
+        if not isinstance(model, (v0_4.ModelDescr, v0_5.ModelDescr)):
+            loaded = load_description(model)
+            if not isinstance(loaded, (v0_4.ModelDescr, v0_5.ModelDescr)):
+                raise ValueError(f"expected model description, but got {loaded}")
+            model = loaded
+
+        pp = create_prediction_pipeline(model)
+
+    if not isinstance(inputs, collections.Mapping):
+        sample_id = str(sample_id)
+        if "{i}" not in sample_id and "{i:" not in sample_id:
+            sample_id += "{i:03}"
+        for i, ipts in tqdm(enumerate(inputs)):
+            yield predict(
+                model=pp,
+                inputs=ipts,
+                sample_id=sample_id.format(i=i),
+                blocksize_parameter=blocksize_parameter,
+                skip_preprocessing=skip_preprocessing,
+                skip_postprocessing=skip_postprocessing,
+                save_output_path=save_output_path,
+            )
