Merge pull request #284 from bioimage-io/raw_pred

allow to create prediction pipeline from raw_nodes.Model

Author: FynnBe

bioimageio/core/prediction_pipeline/_model_adapters/_model_adapter.py

@@ -1,12 +1,15 @@
 import abc
-from typing import List, Optional, Sequence, Type
+from typing import List, Optional, Sequence, Type, Union
 
 import xarray as xr
 
+from bioimageio.core import load_resource_description
 from bioimageio.core.resource_io import nodes
 
 #: Known weight formats in order of priority
 #: First match wins
+from bioimageio.spec.model import raw_nodes
+
 _WEIGHT_FORMATS = ["pytorch_state_dict", "tensorflow_saved_model_bundle", "torchscript", "onnx", "keras_hdf5"]
 
 
@@ -15,11 +18,24 @@ class ModelAdapter(abc.ABC):
     Represents model *without* any preprocessing and postprocessing
     """
 
-    def __init__(self, *, bioimageio_model: nodes.Model, devices: Optional[Sequence[str]] = None):
-        self.bioimageio_model = bioimageio_model
+    def __init__(
+        self, *, bioimageio_model: Union[nodes.Model, raw_nodes.Model], devices: Optional[Sequence[str]] = None
+    ):
+        self.bioimageio_model = self._prepare_model(bioimageio_model)
         self.default_devices = devices
         self.loaded = False
 
+    @staticmethod
+    def _prepare_model(bioimageio_model):
+        """the (raw) model node is prepared (here: loaded as non-raw model node) for the model adapter to be ready
+        for operation.
+        Note: To write a model adapter that uses the raw model node one can overwrite this method.
+        """
+        if isinstance(bioimageio_model, nodes.Model):
+            return bioimageio_model
+        else:
+            return load_resource_description(bioimageio_model)
+
     def __enter__(self):
         """load on entering context"""
         assert not self.loaded
@@ -93,7 +109,10 @@ def get_weight_formats() -> List[str]:
 
 
 def create_model_adapter(
-    *, bioimageio_model: nodes.Model, devices=Optional[Sequence[str]], weight_format: Optional[str] = None
+    *,
+    bioimageio_model: Union[nodes.Model, raw_nodes.Model],
+    devices=Optional[Sequence[str]],
+    weight_format: Optional[str] = None,
 ) -> ModelAdapter:
     """
     Creates model adapter based on the passed spec
@@ -105,7 +124,7 @@ def create_model_adapter(
 
     if weight_format is not None:
         if weight_format not in weight_formats:
-            raise ValueError(f"Weight format {weight_format} is not in supported formats {_WEIGHT_FORMATS}")
+            raise ValueError(f"Weight format {weight_format} is not in supported formats {weight_formats}")
         weight_formats = [weight_format]
 
     for weight in weight_formats:

bioimageio/core/prediction_pipeline/_prediction_pipeline.py

@@ -1,16 +1,18 @@
 import abc
 from dataclasses import dataclass
-from typing import Iterable, List, Optional, Sequence, Tuple
+from typing import Iterable, List, Optional, Sequence, Tuple, Union
 
 import xarray as xr
 from marshmallow import missing
 
 from bioimageio.core.resource_io import nodes
-from bioimageio.core.resource_io.nodes import InputTensor, Model, OutputTensor
+from bioimageio.spec.model import raw_nodes
 from ._combined_processing import CombinedProcessing
 from ._model_adapters import ModelAdapter, create_model_adapter
 from ._stat_state import StatsState
 from ._utils import ComputedMeasures, Sample, TensorName
+from .. import load_resource_description
+from ..resource_io.utils import resolve_raw_node
 
 
 @dataclass
@@ -54,15 +56,15 @@ def name(self) -> str:
 
     @property
     @abc.abstractmethod
-    def input_specs(self) -> List[InputTensor]:
+    def input_specs(self) -> List[nodes.InputTensor]:
         """
         specs of inputs
         """
         ...
 
     @property
     @abc.abstractmethod
-    def output_specs(self) -> List[OutputTensor]:
+    def output_specs(self) -> List[nodes.OutputTensor]:
         """
         specs of outputs
         """
@@ -88,7 +90,7 @@ def __init__(
         self,
         *,
         name: str,
-        bioimageio_model: Model,
+        bioimageio_model: Union[nodes.Model, raw_nodes.Model],
         preprocessing: CombinedProcessing,
         postprocessing: CombinedProcessing,
         ipt_stats: StatsState,
@@ -99,8 +101,14 @@ def __init__(
             raise NotImplementedError(f"Not yet implemented inference for run mode '{bioimageio_model.run_mode.name}'")
 
         self._name = name
-        self._input_specs = bioimageio_model.inputs
-        self._output_specs = bioimageio_model.outputs
+        if isinstance(bioimageio_model, nodes.Model):
+            self._input_specs = bioimageio_model.inputs
+            self._output_specs = bioimageio_model.outputs
+        else:
+            assert isinstance(bioimageio_model, raw_nodes.Model)
+            self._input_specs = [resolve_raw_node(s, nodes) for s in bioimageio_model.inputs]
+            self._output_specs = [resolve_raw_node(s, nodes) for s in bioimageio_model.outputs]
+
         self._preprocessing = preprocessing
         self._postprocessing = postprocessing
         self._ipt_stats = ipt_stats
@@ -176,7 +184,7 @@ def unload(self):
 
 
 def create_prediction_pipeline(
-    bioimageio_model: nodes.Model,
+    bioimageio_model: Union[nodes.Model, raw_nodes.Model],
     *,
     devices: Optional[Sequence[str]] = None,
     weight_format: Optional[str] = None,
@@ -196,8 +204,16 @@ def create_prediction_pipeline(
     model_adapter: ModelAdapter = model_adapter or create_model_adapter(
         bioimageio_model=bioimageio_model, devices=devices, weight_format=weight_format
     )
+    if isinstance(bioimageio_model, nodes.Model):
+        ipts = bioimageio_model.inputs
+        outs = bioimageio_model.outputs
+
+    else:
+        assert isinstance(bioimageio_model, raw_nodes.Model)
+        ipts = [resolve_raw_node(s, nodes) for s in bioimageio_model.inputs]
+        outs = [resolve_raw_node(s, nodes) for s in bioimageio_model.outputs]
 
-    preprocessing = CombinedProcessing(bioimageio_model.inputs)
+    preprocessing = CombinedProcessing(ipts)
 
     def sample_dataset():
         for tensors in dataset_for_initial_statistics:
@@ -209,7 +225,7 @@ def sample_dataset():
         update_dataset_stats_after_n_samples=update_dataset_stats_after_n_samples,
         update_dataset_stats_for_n_samples=update_dataset_stats_for_n_samples,
     )
-    postprocessing = CombinedProcessing(bioimageio_model.outputs)
+    postprocessing = CombinedProcessing(outs)
     out_stats = StatsState(
         postprocessing.required_measures,
         dataset=tuple(),

bioimageio/core/resource_io/io_.py

@@ -19,7 +19,7 @@
 )
 from bioimageio.spec.shared.raw_nodes import ResourceDescription as RawResourceDescription
 from . import nodes
-from .utils import resolve_raw_resource_description, resolve_source
+from .utils import resolve_raw_node, resolve_source
 
 serialize_raw_resource_description = spec.io_.serialize_raw_resource_description
 save_raw_resource_description = spec.io_.save_raw_resource_description
@@ -54,7 +54,7 @@ def load_resource_description(
         else:
             raise ValueError(f"Not found any of the specified weights formats {weights_priority_order}")
 
-    rd: ResourceDescription = resolve_raw_resource_description(raw_rd=raw_rd, nodes_module=nodes)
+    rd: ResourceDescription = resolve_raw_node(raw_rd=raw_rd, nodes_module=nodes)
     assert isinstance(rd, getattr(nodes, get_class_name_from_type(raw_rd.type)))
 
     return rd

bioimageio/core/resource_io/utils.py

@@ -9,7 +9,6 @@
 from bioimageio.spec.shared import raw_nodes, resolve_source, source_available
 from bioimageio.spec.shared.node_transformer import (
     GenericRawNode,
-    GenericRawRD,
     GenericResolvedNode,
     NodeTransformer,
     NodeVisitor,
@@ -105,8 +104,8 @@ def all_sources_available(
         return True
 
 
-def resolve_raw_resource_description(
-    raw_rd: GenericRawRD, nodes_module: typing.Any, uri_only_if_in_package: bool = True
+def resolve_raw_node(
+    raw_rd: GenericRawNode, nodes_module: typing.Any, uri_only_if_in_package: bool = True
 ) -> GenericResolvedNode:
     """resolve all uris and paths (that are included when packaging)"""
     rd = UriNodeTransformer(root_path=raw_rd.root_path, uri_only_if_in_package=uri_only_if_in_package).transform(raw_rd)