WIP axis, tensor, tile, sample

bioimageio/core/__init__.py

@@ -24,6 +24,10 @@
 from ._resource_tests import test_model as test_model
 from ._settings import settings as settings
 from .utils import VERSION
+from .tensor import Tensor as Tensor
+from .tile import Tile as Tile
+from .sample import Sample as Sample
+
 
 __version__ = VERSION
 

bioimageio/core/_prediction_pipeline.py

@@ -2,13 +2,13 @@
 from types import MappingProxyType
 from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Union
 
-from bioimageio.core.common import Tensor, TensorId
 from bioimageio.core.model_adapters import ModelAdapter, create_model_adapter
 from bioimageio.core.model_adapters import get_weight_formats as get_weight_formats
 from bioimageio.core.proc_ops import Processing
 from bioimageio.core.proc_setup import setup_pre_and_postprocessing
 from bioimageio.core.sample import Sample
 from bioimageio.core.stat_measures import DatasetMeasure, MeasureValue
+from bioimageio.core.Tensor import Tensor, TensorId
 from bioimageio.spec.model import AnyModelDescr, v0_4
 from bioimageio.spec.model.v0_5 import WeightsFormat
 

bioimageio/core/_resource_tests.py

@@ -5,7 +5,7 @@
 import numpy as np
 
 from bioimageio.core._prediction_pipeline import create_prediction_pipeline
-from bioimageio.core.common import AxisId, BatchSize
+from bioimageio.core.axis import AxisId, BatchSize
 from bioimageio.core.utils import VERSION, get_test_inputs, get_test_outputs
 from bioimageio.core.utils.tiling import resize_to
 from bioimageio.spec import (

bioimageio/core/axis.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal, Mapping, Optional, TypeVar, Union
+
+from typing_extensions import assert_never
+
+from bioimageio.spec.model import v0_5
+
+
+def _get_axis_type(a: Literal["b", "t", "i", "c", "x", "y", "z"]):
+    if a == "b":
+        return "batch"
+    elif a == "t":
+        return "time"
+    elif a == "i":
+        return "index"
+    elif a == "c":
+        return "channel"
+    elif a in ("x", "y", "z"):
+        return "space"
+    else:
+        assert_never(a)
+
+
+S = TypeVar("S", bound=str)
+
+
+def _get_axis_id(a: Union[Literal["b", "t", "i", "c"], S]):
+    if a == "b":
+        return AxisId("batch")
+    elif a == "t":
+        return AxisId("time")
+    elif a == "i":
+        return AxisId("index")
+    elif a == "c":
+        return AxisId("channel")
+    else:
+        return AxisId(a)
+
+
+AxisId = v0_5.AxisId
+
+T = TypeVar("T")
+PerAxis = Mapping[AxisId, T]
+
+BatchSize = int
+
+AxisLetter = Literal["b", "i", "t", "c", "z", "y", "x"]
+AxisLike = Union[AxisLetter, v0_5.AnyAxis, "Axis"]
+
+
+@dataclass
+class Axis:
+    id: AxisId
+    type: Literal["batch", "channel", "index", "space", "time"]
+
+    @classmethod
+    def create(cls, axis: AxisLike) -> Axis:
+        if isinstance(axis, cls):
+            return axis
+        elif isinstance(axis, Axis):
+            return Axis(id=axis.id, type=axis.type)
+        elif isinstance(axis, str):
+            return Axis(id=_get_axis_id(axis), type=_get_axis_type(axis))
+        elif isinstance(axis, v0_5.AxisBase):
+            return Axis(id=AxisId(axis.id), type=axis.type)
+        else:
+            assert_never(axis)
+
+
+@dataclass
+class AxisInfo(Axis):
+    maybe_singleton: bool
+
+    @classmethod
+    def create(cls, axis: AxisLike, maybe_singleton: Optional[bool] = None) -> AxisInfo:
+        if isinstance(axis, AxisInfo):
+            return axis
+
+        axis_base = super().create(axis)
+        if maybe_singleton is None:
+            if isinstance(axis, Axis):
+                maybe_singleton = False
+            elif isinstance(axis, str):
+                maybe_singleton = axis == "b"
+            else:
+                if axis.size is None:
+                    maybe_singleton = True
+                elif isinstance(axis.size, int):
+                    maybe_singleton = axis.size == 1
+                elif isinstance(axis.size, v0_5.SizeReference):
+                    maybe_singleton = (
+                        False  # TODO: check if singleton is ok for a `SizeReference`
+                    )
+                elif isinstance(
+                    axis.size, (v0_5.ParameterizedSize, v0_5.DataDependentSize)
+                ):
+                    try:
+                        maybe_size_one = axis.size.validate_size(
+                            1
+                        )  # TODO: refactor validate_size() to have boolean func here
+                    except ValueError:
+                        maybe_singleton = False
+                    else:
+                        maybe_singleton = maybe_size_one == 1
+                else:
+                    assert_never(axis.size)
+
+        return AxisInfo(
+            id=axis_base.id, type=axis_base.type, maybe_singleton=maybe_singleton
+        )

bioimageio/core/common.py

@@ -1,51 +1,64 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
-from typing import (
-    TYPE_CHECKING,
-    Dict,
-    Literal,
-    Mapping,
-    NamedTuple,
-    Protocol,
-    Tuple,
-    Union,
-)
+from typing import Literal, NamedTuple, Tuple, TypeVar, Union
 
-import xarray as xr
+from typing_extensions import Self, assert_never
 
-from bioimageio.spec.model import v0_5
+DTypeStr = Literal[
+    "bool",
+    "float32",
+    "float64",
+    "int8",
+    "int16",
+    "int32",
+    "int64",
+    "uint8",
+    "uint16",
+    "uint32",
+    "uint64",
+]
 
-if TYPE_CHECKING:
-    from bioimageio.core.stat_measures import Measure, MeasureValue
 
-TensorId = v0_5.TensorId
-AxisId = v0_5.AxisId
+LeftRight_T = TypeVar("LeftRight_T", bound="LeftRight")
+LeftRightLike = Union[int, Tuple[int, int], LeftRight_T]
 
 
-@dataclass
-class Axis:
-    id: AxisId
-    type: Literal["batch", "channel", "index", "space", "time"]
+class LeftRight(NamedTuple):
+    left: int
+    right: int
 
+    @classmethod
+    def create(cls, like: LeftRightLike[Self]) -> Self:
+        if isinstance(like, cls):
+            return like
+        elif isinstance(like, tuple):
+            return cls(*like)
+        elif isinstance(like, int):
+            return cls(like, like)
+        else:
+            assert_never(like)
 
-class AxisLike(Protocol):
-    id: str
-    type: Literal["batch", "channel", "index", "space", "time"]
 
+class Halo(LeftRight):
+    pass
 
-BatchSize = int
-Tensor = xr.DataArray
 
-Data = Dict[TensorId, Tensor]
-Stat = Dict["Measure", "MeasureValue"]
+HaloLike = LeftRightLike[Halo]
 
 
-class LeftRight(NamedTuple):
-    left: int
-    right: int
+class PadWidth(LeftRight):
+    pass
+
+
+PadWidthLike = LeftRightLike[PadWidth]
+PadMode = Literal["edge", "reflect", "symmetric"]
+PadWhere = Literal["before", "center", "after"]
 
 
 class SliceInfo(NamedTuple):
     start: int
     stop: int
+
+
+TileNumber = int
+TotalNumberOfTiles = int

bioimageio/core/io.py

@@ -0,0 +1,30 @@
+from pathlib import Path
+from typing import Optional, Sequence, Union
+
+import imageio
+
+from bioimageio.core.axis import Axis, AxisLike
+from bioimageio.spec.model import v0_5
+from bioimageio.spec.model.v0_4 import InputTensorDescr as InputTensorDescr04
+from bioimageio.spec.model.v0_4 import OutputTensorDescr as OutputTensorDescr04
+from bioimageio.spec.utils import load_array
+
+from .tensor import Tensor, TensorId
+
+
+def load_tensor(
+    path: Path, axes: Optional[Sequence[AxisLike]] = None, id: Optional[TensorId] = None
+) -> Tensor:
+
+    ext = path.suffix
+    if ext == ".npy":
+        array = load_array(path)
+    else:
+        is_volume = (
+            True
+            if axes is None
+            else sum(Axis.create(a).type != "channel" for a in axes) > 2
+        )
+        array = imageio.volread(path) if is_volume else imageio.imread(path)
+
+    return Tensor.from_numpy(array, axes, id=TensorId(path.stem) if id is None else id)

bioimageio/core/model_adapters/_keras_model_adapter.py

@@ -4,7 +4,7 @@
 from loguru import logger
 from numpy.typing import NDArray
 
-from bioimageio.core.common import Tensor
+from bioimageio.core.Tensor import Tensor
 from bioimageio.spec._internal.io_utils import download
 from bioimageio.spec.model import v0_4, v0_5
 from bioimageio.spec.model.v0_5 import Version

bioimageio/core/model_adapters/_model_adapter.py

@@ -2,7 +2,7 @@
 from abc import ABC, abstractmethod
 from typing import List, Optional, Sequence, Tuple, Union, final
 
-from bioimageio.core.common import Tensor
+from bioimageio.core.Tensor import Tensor
 from bioimageio.spec.model import v0_4, v0_5
 
 WeightsFormat = Union[v0_4.WeightsFormat, v0_5.WeightsFormat]

bioimageio/core/model_adapters/_onnx_model_adapter.py

@@ -3,7 +3,7 @@
 
 from numpy.typing import NDArray
 
-from bioimageio.core.common import Tensor
+from bioimageio.core.Tensor import Tensor
 from bioimageio.spec.model import v0_4, v0_5
 
 from ._model_adapter import ModelAdapter

bioimageio/core/model_adapters/_pytorch_model_adapter.py

@@ -2,7 +2,7 @@
 import warnings
 from typing import Any, List, Optional, Sequence, Tuple, Union
 
-from bioimageio.core.common import Tensor
+from bioimageio.core.Tensor import Tensor
 from bioimageio.core.utils import import_callable
 from bioimageio.spec.model import v0_4, v0_5
 from bioimageio.spec.utils import download

bioimageio/core/model_adapters/_tensorflow_model_adapter.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 
-from bioimageio.core.common import Tensor
+from bioimageio.core.Tensor import Tensor
 from bioimageio.spec.common import FileSource
 from bioimageio.spec.model import v0_4, v0_5
 from bioimageio.spec.utils import download

bioimageio/core/model_adapters/_torchscript_model_adapter.py

@@ -5,7 +5,7 @@
 import numpy as np
 from numpy.typing import NDArray
 
-from bioimageio.core.common import Tensor
+from bioimageio.core.Tensor import Tensor
 from bioimageio.spec.model import v0_4, v0_5
 from bioimageio.spec.utils import download
 

bioimageio/core/proc_ops.py

@@ -18,11 +18,8 @@
 from typing_extensions import Self, assert_never
 
 from bioimageio.core._op_base import Operator
-from bioimageio.core.common import (
+from bioimageio.core.axis import (
     AxisId,
-    Stat,
-    Tensor,
-    TensorId,
 )
 from bioimageio.core.sample import Sample
 from bioimageio.core.stat_calculators import StatsCalculator
@@ -37,8 +34,10 @@
     SampleMean,
     SamplePercentile,
     SampleStd,
+    Stat,
     StdMeasure,
 )
+from bioimageio.core.Tensor import Tensor, TensorId
 from bioimageio.spec.model import v0_4, v0_5