Merge pull request #106 from bioimage-io/clean_pred_pipe

Simplify PredictionPipeline

Author: FynnBe

MANIFEST.in

@@ -0,0 +1 @@
+include bioimageio/core/VERSION

bioimageio/core/prediction.py

@@ -11,7 +11,7 @@
 import xarray as xr
 
 from bioimageio.core import load_resource_description
-from bioimageio.core.resource_io.nodes import Model
+from bioimageio.core.resource_io.nodes import InputTensor, Model, OutputTensor
 from bioimageio.core.prediction_pipeline import PredictionPipeline, create_prediction_pipeline
 from tqdm import tqdm
 
@@ -93,15 +93,19 @@ def pad(im, axes: Sequence[str], padding, pad_right=True) -> Tuple[np.ndarray, D
     return im, crop
 
 
-def load_image(in_path, axes):
+def load_image(in_path, axes: Sequence[str]) -> xr.DataArray:
     ext = os.path.splitext(in_path)[1]
     if ext == ".npy":
         im = np.load(in_path)
     else:
         is_volume = "z" in axes
         im = imageio.volread(in_path) if is_volume else imageio.imread(in_path)
         im = require_axes(im, axes)
-    return im
+    return xr.DataArray(im, dims=axes)
+
+
+def load_tensors(sources, tensor_specs: List[Union[InputTensor, OutputTensor]]) -> List[xr.DataArray]:
+    return [load_image(s, sspec.axes) for s, sspec in zip(sources, tensor_specs)]
 
 
 def _to_channel_last(image):
@@ -239,11 +243,14 @@ def load_tile(tile):
 #
 
 
-def predict(prediction_pipeline, inputs) -> List[xr.DataArray]:
+def predict(prediction_pipeline: PredictionPipeline, inputs) -> List[xr.DataArray]:
     if not isinstance(inputs, (tuple, list)):
         inputs = [inputs]
 
-    tagged_data = [xr.DataArray(ipt, dims=axes) for ipt, axes in zip(inputs, prediction_pipeline.input_axes)]
+    assert len(inputs) == len(prediction_pipeline.input_specs)
+    tagged_data = [
+        xr.DataArray(ipt, dims=ipt_spec.axes) for ipt, ipt_spec in zip(inputs, prediction_pipeline.input_specs)
+    ]
     return prediction_pipeline.forward(*tagged_data)
 
 
@@ -411,10 +418,15 @@ def predict_image(model_rdf, inputs, outputs, padding=None, tiling=None, weight_
 
     padding = parse_padding(padding, model)
     tiling = parse_tiling(tiling, model)
+
+    _predict_sample(prediction_pipeline, inputs, outputs, padding, tiling)
+
+
+def _predict_sample(prediction_pipeline, inputs, outputs, padding, tiling):
     if padding is not None and tiling is not None:
         raise ValueError("Only one of padding or tiling is supported")
 
-    input_data = [load_image(inp, axes) for inp, axes in zip(inputs, prediction_pipeline.input_axes)]
+    input_data = load_tensors(inputs, prediction_pipeline.input_specs)
     if padding is not None:
         result = predict_with_padding(prediction_pipeline, input_data, padding)
     elif tiling is not None:
@@ -439,6 +451,7 @@ def predict_images(
     verbose=False,
 ):
     """Predict multiple inputs with a bioimage.io model."""
+
     model = load_resource_description(model_rdf)
     assert isinstance(model, Model)
 
@@ -448,8 +461,6 @@ def predict_images(
 
     padding = parse_padding(padding, model)
     tiling = parse_tiling(tiling, model)
-    if padding is not None and tiling is not None:
-        raise ValueError("Only one of padding or tiling is supported")
 
     prog = zip(inputs, outputs)
     if verbose:
@@ -462,17 +473,7 @@ def predict_images(
         if not isinstance(outp, (tuple, list)):
             outp = [outp]
 
-        inp = [load_image(im, sp.axes) for im, sp in zip(inp, prediction_pipeline.input_specs)]
-        if padding is not None:
-            res = predict_with_padding(prediction_pipeline, inp, padding)
-        elif tiling is not None:
-            res = predict_with_tiling(prediction_pipeline, inp, tiling)
-        else:
-            res = predict(prediction_pipeline, inp)
-
-        assert isinstance(res, list)
-        for out, r in zip(outp, res):
-            save_image(out, r)
+        _predict_sample(prediction_pipeline, inp, outp, padding, tiling)
 
 
 def test_model(model_rdf: Union[URI, Path, str], weight_format=None, devices=None, decimal=4):

bioimageio/core/prediction_pipeline/_model_adapters/_model_adapter.py

@@ -1,7 +1,8 @@
 import abc
-from typing import Any, List, Optional, Type, Union
+from typing import List, Optional, Type
 
 import xarray as xr
+
 from bioimageio.core.resource_io import nodes
 
 #: Known weigh types in order of priority
@@ -18,7 +19,6 @@ class ModelAdapter(abc.ABC):
     def __init__(self, *, bioimageio_model: nodes.Model, devices=Optional[List[str]]):
         ...
 
-    # todo: separate preprocessing/actual forward/postprocessing
     @abc.abstractmethod
     def forward(self, *input_tensors: xr.DataArray) -> List[xr.DataArray]:
         """

bioimageio/core/prediction_pipeline/_prediction_pipeline.py

@@ -61,73 +61,20 @@ def output_specs(self) -> List[OutputTensor]:
         """
         ...
 
-    # todo: replace all uses of properties below with 'input_specs' and 'output_specs'
-    @property
-    @abc.abstractmethod
-    def input_axes(self) -> List[Tuple[str, ...]]:
-        """
-        Input axes excepted by this pipeline
-        Note: one character axes names
-        """
-        ...
-
-    @property
-    @abc.abstractmethod
-    def input_shape(self) -> List[List[Tuple[str, int]]]:
-        """
-        Named input dimensions
-        """
-        ...
-
-    @property
-    @abc.abstractmethod
-    def output_axes(self) -> List[Tuple[str, ...]]:
-        """
-        Output axes of this pipeline
-        Note: one character axes names
-        """
-        ...
-
-    @property
-    @abc.abstractmethod
-    def output_shape(self) -> List[Union[List[Tuple[str, float]], NamedImplicitOutputShape]]:
-        """
-        Named output dimensions. Either explicitly defined or implicitly in relation to an input
-        """
-        ...
-
-    @property
-    @abc.abstractmethod
-    def halo(self) -> List[List[Tuple[str, int]]]:
-        """
-        Size of output borders that have unreliable data due to artifacts (after application of postprocessing)
-        """
-        ...
-
 
 class _PredictionPipelineImpl(PredictionPipeline):
     def __init__(
         self,
         *,
         name: str,
         bioimageio_model: Model,
-        input_axes: Sequence[str],
-        input_shape: Sequence[List[Tuple[str, int]]],
-        output_axes: Sequence[str],
-        output_shape: Sequence[Union[List[Tuple[str, int]], NamedImplicitOutputShape]],
-        halo: Sequence[List[Tuple[str, int]]],
         preprocessing: Sequence[Transform],
         model: ModelAdapter,
         postprocessing: Sequence[Transform],
     ) -> None:
         self._name = name
         self._input_specs = bioimageio_model.inputs
         self._output_specs = bioimageio_model.outputs
-        self._input_axes = [tuple(axes) for axes in input_axes]
-        self._input_shape = input_shape
-        self._output_axes = [tuple(axes) for axes in output_axes]
-        self._output_shape = output_shape
-        self._halo = halo
         self._preprocessing = preprocessing
         self._model: ModelAdapter = model
         self._postprocessing = postprocessing
@@ -144,26 +91,6 @@ def input_specs(self):
     def output_specs(self):
         return self._output_specs
 
-    @property
-    def input_axes(self):
-        return self._input_axes
-
-    @property
-    def input_shape(self):
-        return self._input_shape
-
-    @property
-    def output_axes(self):
-        return self._output_axes
-
-    @property
-    def output_shape(self):
-        return self._output_shape
-
-    @property
-    def halo(self):
-        return self._halo
-
     def predict(self, *input_tensors: xr.DataArray) -> List[xr.DataArray]:
         """Predict input_tensor with the model without applying pre/postprocessing."""
         return self._model.forward(*input_tensors)
@@ -230,8 +157,6 @@ def create_prediction_pipeline(
         bioimageio_model=bioimageio_model, devices=devices, weight_format=weight_format
     )
 
-    input_axes: List[str] = []
-    named_input_shape: List[List[Tuple[str, int]]] = []
     preprocessing: List[Transform] = []
     for ipt in bioimageio_model.inputs:
         try:
@@ -241,42 +166,17 @@ def create_prediction_pipeline(
         except AttributeError:
             input_shape = ipt.shape
 
-        input_axes.append(ipt.axes)
-        named_input_shape.append(list(zip(ipt.axes, input_shape)))
         preprocessing_spec = [] if ipt.preprocessing is missing else ipt.preprocessing.copy()
         preprocessing.append(make_preprocessing(preprocessing_spec))
 
-    output_axes: List[str] = []
-    named_output_shape: List[Union[List[Tuple[str, int]], NamedImplicitOutputShape]] = []
-    named_halo: List[List[Tuple[str, int]]] = []
     postprocessing: List[Transform] = []
     for out in bioimageio_model.outputs:
-        output_axes.append(out.axes)
-        if isinstance(out.shape, list):  # explict output shape
-            named_output_shape.append(list(zip(out.axes, out.shape)))
-        elif isinstance(out.shape, ImplicitOutputShape):
-            named_output_shape.append(
-                NamedImplicitOutputShape(
-                    reference_input=out.shape.reference_tensor,
-                    scale=list(zip(out.axes, out.shape.scale)),
-                    offset=list(zip(out.axes, out.shape.offset)),
-                )
-            )
-        else:
-            raise TypeError(f"Unexpected type for output shape: {type(out.shape)}")
-
-        named_halo.append(list(zip(out.axes, out.halo or [0 for _ in out.axes])))
         postprocessing_spec = [] if out.postprocessing is missing else out.postprocessing.copy()
         postprocessing.append(make_postprocessing(postprocessing_spec))
 
     return _PredictionPipelineImpl(
         name=bioimageio_model.name,
         bioimageio_model=bioimageio_model,
-        input_axes=input_axes,
-        input_shape=named_input_shape,
-        output_axes=output_axes,
-        output_shape=named_output_shape,
-        halo=named_halo,
         preprocessing=preprocessing,
         model=model_adapter,
         postprocessing=postprocessing,

bioimageio/core/resource_io/nodes.py

@@ -1,7 +1,7 @@
 import pathlib
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Callable, Dict, List, Union
+from typing import Callable, Dict, List, Tuple, Union
 
 from marshmallow import missing
 from marshmallow.utils import _Missing
@@ -107,12 +107,22 @@ class Postprocessing(Node, model_raw_nodes.Postprocessing):
 
 @dataclass
 class InputTensor(Node, model_raw_nodes.InputTensor):
-    pass
+    axes: Tuple[str, ...] = missing
+
+    def __post_init__(self):
+        super().__post_init__()
+        # raw node has string with single letter axes names. Here we use a tuple of string here (like xr.DataArray).
+        self.axes = tuple(self.axes)
 
 
 @dataclass
 class OutputTensor(Node, model_raw_nodes.OutputTensor):
-    pass
+    axes: Tuple[str, ...] = missing
+
+    def __post_init__(self):
+        super().__post_init__()
+        # raw node has string with single letter axes names. Here we use a tuple of string here (like xr.DataArray).
+        self.axes = tuple(self.axes)
 
 
 @dataclass

setup.py

@@ -26,6 +26,7 @@
     ],
     packages=find_namespace_packages(exclude=["tests"]),  # Required
     install_requires=["bioimageio.spec", "imageio>=2.5", "numpy", "xarray"],
+    include_package_data=True,
     extras_require={
         "test": ["pytest", "tox"],
         "dev": ["pre-commit"],