From 79fb89b33011ac3476e2055b061dca29901fe74d Mon Sep 17 00:00:00 2001 From: fynnbe Date: Mon, 31 Oct 2022 20:17:08 +0100 Subject: [PATCH 01/19] clean up nodes and add Workflow --- bioimageio/core/resource_io/nodes.py | 65 +++++++++++++++++++--------- 1 file changed, 45 insertions(+), 20 deletions(-) diff --git a/bioimageio/core/resource_io/nodes.py b/bioimageio/core/resource_io/nodes.py index 47e2035f..7d8465ca 100644 --- a/bioimageio/core/resource_io/nodes.py +++ b/bioimageio/core/resource_io/nodes.py @@ -6,10 +6,12 @@ from marshmallow import missing from marshmallow.utils import _Missing +from bioimageio.spec.collection import raw_nodes as collection_raw_nodes +from bioimageio.spec.dataset import raw_nodes as dataset_raw_nodes from bioimageio.spec.model import raw_nodes as model_raw_nodes from bioimageio.spec.rdf import raw_nodes as rdf_raw_nodes -from bioimageio.spec.collection import raw_nodes as collection_raw_nodes from bioimageio.spec.shared import raw_nodes +from bioimageio.spec.workflow import raw_nodes as workflow_raw_nodes @dataclass @@ -48,12 +50,12 @@ class CiteEntry(Node, rdf_raw_nodes.CiteEntry): @dataclass -class Author(Node, model_raw_nodes.Author): +class Author(Node, rdf_raw_nodes.Author): pass @dataclass -class Maintainer(Node, model_raw_nodes.Maintainer): +class Maintainer(Node, rdf_raw_nodes.Maintainer): pass @@ -62,10 +64,19 @@ class Badge(Node, rdf_raw_nodes.Badge): pass +@dataclass +class Attachments(Node, rdf_raw_nodes.Attachments): + files: Union[_Missing, List[Path]] = missing + unknown: Union[_Missing, Dict[str, Any]] = missing + + @dataclass class RDF(rdf_raw_nodes.RDF, ResourceDescription): + authors: Union[_Missing, List[Author]] = missing + attachments: Union[_Missing, Attachments] = missing badges: Union[_Missing, List[Badge]] = missing - covers: Union[_Missing, List[Path]] = missing + cite: Union[_Missing, List[CiteEntry]] = missing + maintainers: Union[_Missing, List[Maintainer]] = missing @dataclass @@ -74,17 +85,22 @@ class CollectionEntry(Node, collection_raw_nodes.CollectionEntry): @dataclass -class LinkedDataset(Node, model_raw_nodes.LinkedDataset): +class Collection(collection_raw_nodes.Collection, RDF): + collection: List[CollectionEntry] = missing + + +@dataclass +class Dataset(Node, dataset_raw_nodes.Dataset): pass @dataclass -class ModelParent(Node, model_raw_nodes.ModelParent): +class LinkedDataset(Node, model_raw_nodes.LinkedDataset): pass @dataclass -class Collection(collection_raw_nodes.Collection, RDF): +class ModelParent(Node, model_raw_nodes.ModelParent): pass @@ -106,6 +122,7 @@ class Postprocessing(Node, model_raw_nodes.Postprocessing): @dataclass class InputTensor(Node, model_raw_nodes.InputTensor): axes: Tuple[str, ...] = missing + preprocessing: Union[_Missing, List[Preprocessing]] = missing def __post_init__(self): super().__post_init__() @@ -116,6 +133,7 @@ def __post_init__(self): @dataclass class OutputTensor(Node, model_raw_nodes.OutputTensor): axes: Tuple[str, ...] = missing + postprocessing: Union[_Missing, List[Postprocessing]] = missing def __post_init__(self): super().__post_init__() @@ -132,40 +150,39 @@ def __call__(self, *args, **kwargs): @dataclass -class KerasHdf5WeightsEntry(Node, model_raw_nodes.KerasHdf5WeightsEntry): - source: Path = missing +class WeightsEntryBase(model_raw_nodes.WeightsEntryBase): + dependencies: Union[_Missing, Dependencies] = missing @dataclass -class OnnxWeightsEntry(Node, model_raw_nodes.OnnxWeightsEntry): +class KerasHdf5WeightsEntry(WeightsEntryBase, model_raw_nodes.KerasHdf5WeightsEntry): source: Path = missing @dataclass -class PytorchStateDictWeightsEntry(Node, model_raw_nodes.PytorchStateDictWeightsEntry): +class OnnxWeightsEntry(WeightsEntryBase, model_raw_nodes.OnnxWeightsEntry): source: Path = missing - architecture: Union[_Missing, ImportedSource] = missing @dataclass -class TorchscriptWeightsEntry(Node, model_raw_nodes.TorchscriptWeightsEntry): +class PytorchStateDictWeightsEntry(WeightsEntryBase, model_raw_nodes.PytorchStateDictWeightsEntry): source: Path = missing + architecture: Union[_Missing, ImportedSource] = missing @dataclass -class TensorflowJsWeightsEntry(Node, model_raw_nodes.TensorflowJsWeightsEntry): +class TorchscriptWeightsEntry(WeightsEntryBase, model_raw_nodes.TorchscriptWeightsEntry): source: Path = missing @dataclass -class TensorflowSavedModelBundleWeightsEntry(Node, model_raw_nodes.TensorflowSavedModelBundleWeightsEntry): +class TensorflowJsWeightsEntry(WeightsEntryBase, model_raw_nodes.TensorflowJsWeightsEntry): source: Path = missing @dataclass -class Attachments(Node, rdf_raw_nodes.Attachments): - files: Union[_Missing, List[Path]] = missing - unknown: Union[_Missing, Dict[str, Any]] = missing +class TensorflowSavedModelBundleWeightsEntry(WeightsEntryBase, model_raw_nodes.TensorflowSavedModelBundleWeightsEntry): + source: Path = missing WeightsEntry = Union[ @@ -180,8 +197,16 @@ class Attachments(Node, rdf_raw_nodes.Attachments): @dataclass class Model(model_raw_nodes.Model, RDF): - authors: List[Author] = missing - maintainers: Union[_Missing, List[Maintainer]] = missing + inputs: List[InputTensor] = missing + outputs: List[OutputTensor] = missing + parent: Union[_Missing, ModelParent] = missing + run_mode: Union[_Missing, RunMode] = missing test_inputs: List[Path] = missing test_outputs: List[Path] = missing + training_data: Union[_Missing, Dataset, LinkedDataset] = missing weights: Dict[model_raw_nodes.WeightsFormat, WeightsEntry] = missing + + +@dataclass +class Workflow(workflow_raw_nodes.Workflow, RDF): + pass From 6373c886e0f9f32ebaaf686b5f4a06fdf2bb1ebf Mon Sep 17 00:00:00 2001 From: fynnbe Date: Fri, 4 Nov 2022 14:10:45 +0100 Subject: [PATCH 02/19] update workflow nodes --- bioimageio/core/resource_io/nodes.py | 56 +++++++++++++++++++++++++++- 1 file changed, 55 insertions(+), 1 deletion(-) diff --git a/bioimageio/core/resource_io/nodes.py b/bioimageio/core/resource_io/nodes.py index 7d8465ca..b369a97f 100644 --- a/bioimageio/core/resource_io/nodes.py +++ b/bioimageio/core/resource_io/nodes.py @@ -208,5 +208,59 @@ class Model(model_raw_nodes.Model, RDF): @dataclass -class Workflow(workflow_raw_nodes.Workflow, RDF): +class Axis(Node, workflow_raw_nodes.Axis): + pass + + +@dataclass +class BatchAxis(Node, workflow_raw_nodes.BatchAxis): + pass + + +@dataclass +class ChannelAxis(Node, workflow_raw_nodes.ChannelAxis): + pass + + +@dataclass +class IndexAxis(Node, workflow_raw_nodes.IndexAxis): + pass + + +@dataclass +class SpaceAxis(Node, workflow_raw_nodes.SpaceAxis): + pass + + +@dataclass +class TimeAxis(Node, workflow_raw_nodes.TimeAxis): pass + + +@dataclass +class InputSpec(Node, workflow_raw_nodes.InputSpec): + pass + + +@dataclass +class OptionSpec(Node, workflow_raw_nodes.OptionSpec): + pass + + +@dataclass +class OutputSpec(Node, workflow_raw_nodes.OutputSpec): + pass + + +@dataclass +class Step(Node, workflow_raw_nodes.Step): + pass + + +@dataclass +class Workflow(workflow_raw_nodes.Workflow, RDF): + inputs: List[InputSpec] = missing + options: List[OptionSpec] = missing + outputs: List[OutputSpec] = missing + steps: List[Step] = missing + test_steps: List[Step] = missing From d473faa84927db435b1cf40f881840fa6cfa6c8f Mon Sep 17 00:00:00 2001 From: fynnbe Date: Fri, 4 Nov 2022 14:11:40 +0100 Subject: [PATCH 03/19] allow for node defaults (don't init with missing) --- bioimageio/core/resource_io/utils.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/bioimageio/core/resource_io/utils.py b/bioimageio/core/resource_io/utils.py index 430d30fc..ffeefbaa 100644 --- a/bioimageio/core/resource_io/utils.py +++ b/bioimageio/core/resource_io/utils.py @@ -6,6 +6,8 @@ import typing from types import ModuleType +from marshmallow import missing + from bioimageio.spec.shared import raw_nodes, resolve_source, source_available from bioimageio.spec.shared.node_transformer import ( GenericRawNode, @@ -85,7 +87,9 @@ def __init__(self, nodes_module: ModuleType): def generic_transformer(self, node: GenericRawNode) -> GenericResolvedNode: if isinstance(node, raw_nodes.RawNode): resolved_data = { - field.name: self.transform(getattr(node, field.name)) for field in dataclasses.fields(node) + field.name: self.transform(getattr(node, field.name)) + for field in dataclasses.fields(node) + if getattr(node, field.name) is not missing # exclude missing fields to respect for node defaults } resolved_node_type: typing.Type[GenericResolvedNode] = getattr(self.nodes, node.__class__.__name__) return resolved_node_type(**resolved_data) # type: ignore From 91f7d973f099394eb2e1f810808cf974c6cf18b9 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Fri, 4 Nov 2022 19:32:32 +0100 Subject: [PATCH 04/19] make image_helper agnostic to axis letter vs name --- bioimageio/core/image_helper.py | 22 ++++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/bioimageio/core/image_helper.py b/bioimageio/core/image_helper.py index 0468b61f..2bbb3229 100644 --- a/bioimageio/core/image_helper.py +++ b/bioimageio/core/image_helper.py @@ -13,7 +13,7 @@ # -def transform_input_image(image: np.ndarray, tensor_axes: str, image_axes: Optional[str] = None): +def transform_input_image(image: np.ndarray, tensor_axes: Sequence[str], image_axes: Optional[str] = None): """Transform input image into output tensor with desired axes. Args: @@ -35,7 +35,16 @@ def transform_input_image(image: np.ndarray, tensor_axes: str, image_axes: Optio image_axes = "bczyx" else: raise ValueError(f"Invalid number of image dimensions: {ndim}") - tensor = DataArray(image, dims=tuple(image_axes)) + + # instead of 'b' we might want 'batch', etc... + axis_letter_map = { + letter: name + for letter, name in {"b": "batch", "c": "channel", "i": "index", "t": "time"} + if name in tensor_axes # only do this mapping if the full name is in the desired tensor_axes + } + image_axes = tuple(axis_letter_map.get(a, a) for a in image_axes) + + tensor = DataArray(image, dims=image_axes) # expand the missing image axes missing_axes = tuple(set(tensor_axes) - set(image_axes)) tensor = tensor.expand_dims(dim=missing_axes) @@ -75,9 +84,10 @@ def transform_output_tensor(tensor: np.ndarray, tensor_axes: str, output_axes: s def to_channel_last(image): - chan_id = image.dims.index("c") + c = "c" if "c" in image.dims else "channel" + chan_id = image.dims.index(c) if chan_id != image.ndim - 1: - target_axes = tuple(ax for ax in image.dims if ax != "c") + ("c",) + target_axes = tuple(ax for ax in image.dims if ax != c) + (c,) image = image.transpose(*target_axes) return image @@ -113,9 +123,9 @@ def save_image(out_path, image): squeeze = {ax: 0 if (ax in "bc" and sh == 1) else slice(None) for ax, sh in zip(image.dims, image.shape)} image = image[squeeze] - if "b" in image.dims: + if "b" in image.dims or "batch" in image.dims: raise RuntimeError(f"Cannot save prediction with batchsize > 1 as {ext}-file") - if "c" in image.dims: # image formats need channel last + if "c" in image.dims or "channel" in image.dims: # image formats need channel last image = to_channel_last(image) save_function = imageio.volsave if is_volume else imageio.imsave From e3371997189fc92d998ffde332af6f047e802b37 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Sat, 5 Nov 2022 00:30:26 +0100 Subject: [PATCH 05/19] accept pathlib.Path in save_image --- bioimageio/core/image_helper.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/bioimageio/core/image_helper.py b/bioimageio/core/image_helper.py index 2bbb3229..fdbead71 100644 --- a/bioimageio/core/image_helper.py +++ b/bioimageio/core/image_helper.py @@ -112,10 +112,10 @@ def load_tensors(sources, tensor_specs: List[Union[InputTensor, OutputTensor]]) return [load_image(s, sspec.axes) for s, sspec in zip(sources, tensor_specs)] -def save_image(out_path, image): - ext = os.path.splitext(out_path)[1] +def save_image(out_path: os.PathLike, image): + ext = os.path.splitext(str(out_path))[1] if ext == ".npy": - np.save(out_path, image) + np.save(str(out_path), image) else: is_volume = "z" in image.dims From 8a52f600fb2f4eed895fa0195ad3f03312e3ee93 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Sat, 5 Nov 2022 00:31:11 +0100 Subject: [PATCH 06/19] fix nodes --- bioimageio/core/resource_io/nodes.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/bioimageio/core/resource_io/nodes.py b/bioimageio/core/resource_io/nodes.py index b369a97f..1c850795 100644 --- a/bioimageio/core/resource_io/nodes.py +++ b/bioimageio/core/resource_io/nodes.py @@ -8,7 +8,7 @@ from bioimageio.spec.collection import raw_nodes as collection_raw_nodes from bioimageio.spec.dataset import raw_nodes as dataset_raw_nodes -from bioimageio.spec.model import raw_nodes as model_raw_nodes +from bioimageio.spec.model.v0_4 import raw_nodes as model_raw_nodes from bioimageio.spec.rdf import raw_nodes as rdf_raw_nodes from bioimageio.spec.shared import raw_nodes from bioimageio.spec.workflow import raw_nodes as workflow_raw_nodes @@ -150,7 +150,7 @@ def __call__(self, *args, **kwargs): @dataclass -class WeightsEntryBase(model_raw_nodes.WeightsEntryBase): +class WeightsEntryBase(model_raw_nodes._WeightsEntryBase): dependencies: Union[_Missing, Dependencies] = missing @@ -259,8 +259,8 @@ class Step(Node, workflow_raw_nodes.Step): @dataclass class Workflow(workflow_raw_nodes.Workflow, RDF): - inputs: List[InputSpec] = missing - options: List[OptionSpec] = missing - outputs: List[OutputSpec] = missing + inputs_spec: List[InputSpec] = missing + options_spec: List[OptionSpec] = missing + outputs_spec: List[OutputSpec] = missing steps: List[Step] = missing test_steps: List[Step] = missing From f028bd28cdd0dfcd4570a95e75e89be607ff7ad1 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Sat, 5 Nov 2022 00:33:09 +0100 Subject: [PATCH 07/19] add cli command run --- bioimageio/core/__main__.py | 99 ++++++++++++++++++++++++++++++++++--- 1 file changed, 93 insertions(+), 6 deletions(-) diff --git a/bioimageio/core/__main__.py b/bioimageio/core/__main__.py index a26f43d1..8be385ff 100644 --- a/bioimageio/core/__main__.py +++ b/bioimageio/core/__main__.py @@ -3,19 +3,23 @@ import os import sys import warnings +from argparse import ArgumentParser +from functools import partial from glob import glob - from pathlib import Path -from pprint import pformat, pprint -from typing import List, Optional +from pprint import pformat +from typing import List, Optional, Union import typer -from bioimageio.core import __version__, prediction, commands, resource_tests, load_raw_resource_description +from bioimageio.core import __version__, commands, load_raw_resource_description, prediction, resource_tests from bioimageio.core.common import TestSummary -from bioimageio.core.prediction_pipeline import get_weight_formats +from bioimageio.core.image_helper import load_image, save_image +from bioimageio.core.resource_io import nodes +from bioimageio.core.workflow.operators import run_workflow from bioimageio.spec.__main__ import app, help_version as help_version_spec from bioimageio.spec.model.raw_nodes import WeightsFormat +from bioimageio.spec.workflow.raw_nodes import Workflow try: from typing import get_args @@ -244,7 +248,7 @@ def predict_images( tiling = json.loads(tiling.replace("'", '"')) assert isinstance(tiling, dict) - # this is a weird typer bug: default devices are empty tuple although they should be None + # this is a weird typer bug: default devices are empty tuple, although they should be None if len(devices) == 0: devices = None prediction.predict_images( @@ -310,5 +314,88 @@ def convert_keras_weights_to_tensorflow( ) +@app.command(context_settings=dict(allow_extra_args=True, ignore_unknown_options=True), add_help_option=False) +def run( + rdf_source: str = typer.Argument(..., help="BioImage.IO RDF id/url/path."), + *, + output_folder: Path = Path("outputs"), + output_tensor_extension: str = ".npy", + ctx: typer.Context, +): + resource = load_raw_resource_description(rdf_source, update_to_format="latest") + if not isinstance(resource, Workflow): + raise NotImplementedError(f"Non-workflow RDFs not yet supported (got type {resource.type})") + + map_type = dict( + any=str, + boolean=bool, + float=float, + int=int, + list=str, + string=str, + ) + wf = resource + parser = ArgumentParser(description=f"CLI for {wf.name}") + + # replicate typer args to show up in help + parser.add_argument( + metavar="rdf-source", + dest="rdf_source", + help="BioImage.IO RDF id/url/path. The optional arguments below are RDF specific.", + ) + parser.add_argument( + metavar="output-folder", dest="output_folder", help="Folder to save outputs to.", default=Path("outputs") + ) + parser.add_argument( + metavar="output-tensor-extension", + dest="output_tensor_extension", + help="Output tensor extension.", + default=".npy", + ) + + def add_param_args(params): + for param in params: + argument_kwargs = {} + if param.type == "tensor": + argument_kwargs["type"] = partial(load_image, axes=[a.name or a.type for a in param.axes]) + else: + argument_kwargs["type"] = map_type[param.type] + + if param.type == "list": + argument_kwargs["nargs"] = "*" + + argument_kwargs["help"] = param.description or "" + if hasattr(param, "default"): + argument_kwargs["default"] = param.default + else: + argument_kwargs["required"] = True + + argument_kwargs["metavar"] = param.name[0].capitalize() + parser.add_argument("--" + param.name.replace("_", "-"), **argument_kwargs) + + def prepare_parameter(value, param: Union[nodes.InputSpec, nodes.OptionSpec]): + if param.type == "tensor": + return load_image(value, [a.name or a.type for a in param.axes]) + else: + return value + + add_param_args(wf.inputs_spec) + add_param_args(wf.options_spec) + args = parser.parse_args([rdf_source, str(output_folder), output_tensor_extension] + list(ctx.args)) + outputs = run_workflow( + rdf_source, + inputs=[prepare_parameter(getattr(args, ipt.name), ipt) for ipt in wf.inputs_spec], + options={opt.name: prepare_parameter(getattr(args, opt.name), opt) for opt in wf.options_spec}, + ) + output_folder.mkdir(parents=True, exist_ok=True) + for out_spec, out in zip(wf.outputs_spec, outputs): + out_path = output_folder / out_spec.name + if out_spec.type == "tensor": + save_image(out_path.with_suffix(output_tensor_extension), out) + else: + with out_path.with_suffix(".json").open("w") as f: + json.dump(out, f) + + if __name__ == "__main__": app() From 7d58fa7a740584e541ccd689e54b6d8398dd54be Mon Sep 17 00:00:00 2001 From: fynnbe Date: Sat, 5 Nov 2022 00:34:38 +0100 Subject: [PATCH 08/19] add workflow operators --- bioimageio/core/workflow/__init__.py | 0 .../core/workflow/operators/__init__.py | 4 + bioimageio/core/workflow/operators/_assert.py | 8 + .../core/workflow/operators/_generate.py | 15 ++ bioimageio/core/workflow/operators/_run.py | 162 ++++++++++++++++++ .../core/workflow/operators/_various.py | 57 ++++++ 6 files changed, 246 insertions(+) create mode 100644 bioimageio/core/workflow/__init__.py create mode 100644 bioimageio/core/workflow/operators/__init__.py create mode 100644 bioimageio/core/workflow/operators/_assert.py create mode 100644 bioimageio/core/workflow/operators/_generate.py create mode 100644 bioimageio/core/workflow/operators/_run.py create mode 100644 bioimageio/core/workflow/operators/_various.py diff --git a/bioimageio/core/workflow/__init__.py b/bioimageio/core/workflow/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/bioimageio/core/workflow/operators/__init__.py b/bioimageio/core/workflow/operators/__init__.py new file mode 100644 index 00000000..87e447e4 --- /dev/null +++ b/bioimageio/core/workflow/operators/__init__.py @@ -0,0 +1,4 @@ +from ._assert import assert_shape +from ._generate import generate_random_uniform_tensor +from ._run import run_model_inference, run_workflow +from ._various import binarize, load_tensors, log, select_outputs diff --git a/bioimageio/core/workflow/operators/_assert.py b/bioimageio/core/workflow/operators/_assert.py new file mode 100644 index 00000000..fdf54302 --- /dev/null +++ b/bioimageio/core/workflow/operators/_assert.py @@ -0,0 +1,8 @@ +from typing import Sequence + +import xarray as xr + + +def assert_shape(tensor: xr.DataArray, shape: Sequence[int]) -> xr.DataArray: + assert tensor.shape == tuple(shape) + return tensor diff --git a/bioimageio/core/workflow/operators/_generate.py b/bioimageio/core/workflow/operators/_generate.py new file mode 100644 index 00000000..c5f09237 --- /dev/null +++ b/bioimageio/core/workflow/operators/_generate.py @@ -0,0 +1,15 @@ +from typing import Sequence, Union + +import numpy as np +import xarray as xr + + +def generate_random_uniform_tensor( + shape: Sequence[Union[int, str]], axes: Sequence[str], *, low: Union[int, float] = 0, high: Union[int, float] = 1 +) -> xr.DataArray: + """generate a tensor with uniformly distributed samples in the interval [low, high) + Returns: + xr.DataArray: random tensor + """ + assert len(shape) == len(axes) + return xr.DataArray(np.random.uniform(low=low, high=high, size=[int(s) for s in shape]), dims=tuple(axes)) diff --git a/bioimageio/core/workflow/operators/_run.py b/bioimageio/core/workflow/operators/_run.py new file mode 100644 index 00000000..1ef173c4 --- /dev/null +++ b/bioimageio/core/workflow/operators/_run.py @@ -0,0 +1,162 @@ +from os import PathLike +from typing import Any, Dict, IO, List, Optional, Sequence, Tuple, Union + +import numpy as np +import xarray as xr +from marshmallow import missing + +from bioimageio.core import load_resource_description +from bioimageio.core.prediction_pipeline import create_prediction_pipeline +from bioimageio.core.resource_io import nodes +from bioimageio.spec import load_raw_resource_description +from bioimageio.spec.model import raw_nodes +from bioimageio.spec.shared.raw_nodes import ResourceDescription as RawResourceDescription + + +def run_model_inference( + rdf_source: Union[dict, PathLike, IO, str, bytes, raw_nodes.URI, RawResourceDescription], + *tensors, + enable_preprocessing: bool = True, + enable_postprocessing: bool = True, + devices: Optional[Sequence[str]] = None, +) -> List[xr.DataArray]: + """run model inference + + Returns: + list: model outputs + """ + model = load_raw_resource_description(rdf_source, update_to_format="latest") + assert isinstance(model, raw_nodes.Model) + # remove pre-/postprocessing if specified + if not enable_preprocessing: + for ipt in model.inputs: + if ipt.preprocessing: + ipt.preprocessing = missing + if not enable_postprocessing: + for ipt in model.outputs: + if ipt.postprocessing: + ipt.postprocessing = missing + + with create_prediction_pipeline(model, devices=devices) as pred_pipeline: + return pred_pipeline.forward(*tensors) + + +def run_workflow( + rdf_source: Union[dict, PathLike, IO, str, bytes, raw_nodes.URI, RawResourceDescription], + inputs: Sequence = tuple(), + options: Dict[str, Any] = None, +) -> Sequence: + return _run_workflow(rdf_source, test_steps=False, inputs=inputs, options=options) + + +def run_workflow_test( + rdf_source: Union[dict, PathLike, IO, str, bytes, raw_nodes.URI, RawResourceDescription], +) -> Sequence: + return _run_workflow(rdf_source, test_steps=True) + + +def _run_workflow( + rdf_source: Union[dict, PathLike, IO, str, bytes, raw_nodes.URI, RawResourceDescription], + *, + test_steps: bool, + inputs: Sequence = tuple(), + options: Dict[str, Any] = None, +) -> Tuple: + import bioimageio.core.workflow.operators as ops + + workflow = load_resource_description(rdf_source) + assert isinstance(workflow, nodes.Workflow) + wf_options = {opt.name: opt.default for opt in workflow.options_spec} + if test_steps: + assert not inputs + assert not options + wf_inputs = {} + steps = workflow.test_steps + else: + if not len(workflow.inputs_spec) == len(inputs): + raise ValueError(f"Expected {len(workflow.inputs_spec)} inputs, but got {len(inputs)}.") + + wf_inputs = {ipt_spec.name: ipt for ipt_spec, ipt in zip(workflow.inputs_spec, inputs)} + for k, v in options.items(): + if k not in wf_options: + raise ValueError(f"Got unknown option {k}, expected one of {set(wf_options)}.") + + wf_options[k] = v + + steps = workflow.steps + + named_outputs = {} # for later referencing + + def map_ref(value): + assert isinstance(workflow, nodes.Workflow) + if isinstance(value, str) and value.startswith("${{") and value.endswith("}}"): + ref = value[4:-2].strip() + if ref.startswith("self.inputs."): + ref = ref[len("self.inputs.") :] + if ref not in wf_inputs: + raise ValueError(f"Invalid workflow input reference {value}.") + + return wf_inputs[ref] + elif ref.startswith("self.options."): + ref = ref[len("self.options.") :] + if ref not in wf_options: + raise ValueError(f"Invalid workflow option reference {value}.") + + return wf_options[ref] + elif ref == "self.rdf_source": + assert workflow.rdf_source is not missing + return str(workflow.rdf_source) + elif ref in named_outputs: + return named_outputs[ref] + else: + raise ValueError(f"Invalid reference {value}.") + else: + return value + + # implicit inputs to a step are the outputs of the previous step. + # For the first step these are the workflow inputs. + outputs = inputs + for step in steps: + if not hasattr(ops, step.op): + raise NotImplementedError(f"{step.op} not implemented in {ops}") + + op = getattr(ops, step.op) + if step.inputs is missing: + inputs = outputs + else: + inputs = [map_ref(ipt) for ipt in step.inputs] + + options = {k: map_ref(v) for k, v in (step.options or {}).items()} + outputs = op(*inputs, **options) + if not isinstance(outputs, tuple): + outputs = (outputs,) + + if step.outputs: + assert step.id is not missing + if len(step.outputs) != len(outputs): + raise ValueError( + f"Got {len(step.outputs)} step output name{'s' if len(step.outputs) > 1 else ''} ({step.id}.outputs), " + f"but op {step.op} returned {len(outputs)} outputs." + ) + + named_outputs.update({f"{step.id}.outputs.{out_name}": out for out_name, out in zip(step.outputs, outputs)}) + + if len(workflow.outputs_spec) != len(outputs): + raise ValueError(f"Expected {len(workflow.outputs_spec)} outputs from last step, but got {len(outputs)}.") + + def tensor_as_xr(tensor, axes: Sequence[nodes.Axis]): + spec_axes = [a.name or a.type for a in axes] + if isinstance(tensor, xr.DataArray): + if list(tensor.dims) != spec_axes: + raise ValueError( + f"Last workflow step returned xarray.DataArray with dims {tensor.dims}, but expected dims {spec_axes}." + ) + + return tensor + else: + return xr.DataArray(tensor, dims=spec_axes) + + return [ + tensor_as_xr(out, out_spec.axes) if out_spec.type == "tensor" else out + for out_spec, out in zip(workflow.outputs_spec, outputs) + ] diff --git a/bioimageio/core/workflow/operators/_various.py b/bioimageio/core/workflow/operators/_various.py new file mode 100644 index 00000000..37bb3520 --- /dev/null +++ b/bioimageio/core/workflow/operators/_various.py @@ -0,0 +1,57 @@ +import logging +from typing import List, Sequence, Tuple + +import numpy as np +import xarray as xr +from imageio import imread + +logger = logging.getLogger(__name__) + + +def binarize(tensor: xr.DataArray, threshold: float): + return tensor > threshold + + +def select_outputs(*args) -> Tuple: + """helper to select workflow outputs (to be used as a final step in a workflow) + + Returns: + tuple: selected outputs (inputs to this op) + + """ + + return args + + +def log(*args, log_level: int = logging.INFO, **kwargs) -> Tuple: + """log any key word arguments (kwargs/options) + + Returns: + tuple: positional inputs to this op + + """ + for k, v in kwargs.items(): + logger.log( + log_level, + f"{k}: %s", + f"{v.shape} mean: {v.mean().item():.4f} std: {v.std().item():.4f}" + if isinstance(v, (np.ndarray, xr.DataArray)) + else v, + ) + + return args + + +def load_tensors(sources: List[str], axes: Sequence[str]) -> List[xr.DataArray]: + """load tensors""" + assert len(sources) == len(axes) + tensors = [] + for source, ax in zip(sources, axes): + if source.split(".")[-1] == ".npy": + data = np.load(str(source)) + else: + data = imread(source) + + tensors.append(xr.DataArray(data, dims=ax)) + + return tensors From 82f28f5c12df3243688b8070944da4bf89d2a324 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Tue, 8 Nov 2022 11:51:21 +0100 Subject: [PATCH 09/19] add iterate_workflow_steps and iterate_test_workflow_steps --- bioimageio/core/workflow/operators/_run.py | 66 +++++++++++++++++----- 1 file changed, 53 insertions(+), 13 deletions(-) diff --git a/bioimageio/core/workflow/operators/_run.py b/bioimageio/core/workflow/operators/_run.py index 1ef173c4..d803b88f 100644 --- a/bioimageio/core/workflow/operators/_run.py +++ b/bioimageio/core/workflow/operators/_run.py @@ -1,5 +1,6 @@ +from dataclasses import dataclass from os import PathLike -from typing import Any, Dict, IO, List, Optional, Sequence, Tuple, Union +from typing import Any, Dict, Generator, IO, List, Optional, Sequence, Tuple, Union import numpy as np import xarray as xr @@ -45,32 +46,64 @@ def run_workflow( rdf_source: Union[dict, PathLike, IO, str, bytes, raw_nodes.URI, RawResourceDescription], inputs: Sequence = tuple(), options: Dict[str, Any] = None, -) -> Sequence: - return _run_workflow(rdf_source, test_steps=False, inputs=inputs, options=options) +) -> tuple: + outputs = tuple() + for state in _iterate_workflow_steps_impl(rdf_source, test_steps=False, inputs=inputs, options=options): + outputs = state.outputs + + return outputs def run_workflow_test( rdf_source: Union[dict, PathLike, IO, str, bytes, raw_nodes.URI, RawResourceDescription], -) -> Sequence: - return _run_workflow(rdf_source, test_steps=True) +) -> tuple: + outputs = tuple() + for state in _iterate_workflow_steps_impl(rdf_source, test_steps=True): + outputs = state.outputs + + return outputs + + +@dataclass +class WorkflowState: + wf_inputs: Dict[str, Any] + wf_options: Dict[str, Any] + inputs: tuple + outputs: tuple + named_outputs: Dict[str, Any] + + +def iterate_workflow_steps( + rdf_source: Union[dict, PathLike, IO, str, bytes, raw_nodes.URI, RawResourceDescription], + *, + inputs: Sequence = tuple(), + options: Dict[str, Any] = None, +) -> Generator[WorkflowState]: + yield from _iterate_workflow_steps_impl(rdf_source, inputs=inputs, options=options, test_steps=False) + + +def iterate_test_workflow_steps( + rdf_source: Union[dict, PathLike, IO, str, bytes, raw_nodes.URI, RawResourceDescription] +) -> Generator[WorkflowState]: + yield from _iterate_workflow_steps_impl(rdf_source, test_steps=True) -def _run_workflow( +def _iterate_workflow_steps_impl( rdf_source: Union[dict, PathLike, IO, str, bytes, raw_nodes.URI, RawResourceDescription], *, test_steps: bool, inputs: Sequence = tuple(), options: Dict[str, Any] = None, -) -> Tuple: +) -> Generator[WorkflowState]: import bioimageio.core.workflow.operators as ops workflow = load_resource_description(rdf_source) assert isinstance(workflow, nodes.Workflow) - wf_options = {opt.name: opt.default for opt in workflow.options_spec} + wf_options: Dict[str, Any] = {opt.name: opt.default for opt in workflow.options_spec} if test_steps: assert not inputs assert not options - wf_inputs = {} + wf_inputs: Dict[str, Any] = {} steps = workflow.test_steps else: if not len(workflow.inputs_spec) == len(inputs): @@ -115,7 +148,7 @@ def map_ref(value): # implicit inputs to a step are the outputs of the previous step. # For the first step these are the workflow inputs. - outputs = inputs + outputs = tuple(inputs) for step in steps: if not hasattr(ops, step.op): raise NotImplementedError(f"{step.op} not implemented in {ops}") @@ -124,8 +157,9 @@ def map_ref(value): if step.inputs is missing: inputs = outputs else: - inputs = [map_ref(ipt) for ipt in step.inputs] + inputs = tuple(map_ref(ipt) for ipt in step.inputs) + assert isinstance(inputs, tuple) options = {k: map_ref(v) for k, v in (step.options or {}).items()} outputs = op(*inputs, **options) if not isinstance(outputs, tuple): @@ -141,6 +175,9 @@ def map_ref(value): named_outputs.update({f"{step.id}.outputs.{out_name}": out for out_name, out in zip(step.outputs, outputs)}) + yield WorkflowState( + wf_inputs=wf_inputs, wf_options=wf_options, inputs=inputs, outputs=outputs, named_outputs=named_outputs + ) if len(workflow.outputs_spec) != len(outputs): raise ValueError(f"Expected {len(workflow.outputs_spec)} outputs from last step, but got {len(outputs)}.") @@ -156,7 +193,10 @@ def tensor_as_xr(tensor, axes: Sequence[nodes.Axis]): else: return xr.DataArray(tensor, dims=spec_axes) - return [ + outputs = tuple( tensor_as_xr(out, out_spec.axes) if out_spec.type == "tensor" else out for out_spec, out in zip(workflow.outputs_spec, outputs) - ] + ) + yield WorkflowState( + wf_inputs=wf_inputs, wf_options=wf_options, inputs=inputs, outputs=outputs, named_outputs=named_outputs + ) From 2c482bbce64d01d8aae19192830b2daedddf5418 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Wed, 1 Feb 2023 19:22:19 +0100 Subject: [PATCH 10/19] update workflow nodes --- bioimageio/core/resource_io/nodes.py | 41 +++++----------------------- 1 file changed, 7 insertions(+), 34 deletions(-) diff --git a/bioimageio/core/resource_io/nodes.py b/bioimageio/core/resource_io/nodes.py index 1c850795..9f441e22 100644 --- a/bioimageio/core/resource_io/nodes.py +++ b/bioimageio/core/resource_io/nodes.py @@ -213,54 +213,27 @@ class Axis(Node, workflow_raw_nodes.Axis): @dataclass -class BatchAxis(Node, workflow_raw_nodes.BatchAxis): +class BatchAxis(Node, workflow_raw_nodes.Axis): pass @dataclass -class ChannelAxis(Node, workflow_raw_nodes.ChannelAxis): +class Input(Node, workflow_raw_nodes.Input): pass @dataclass -class IndexAxis(Node, workflow_raw_nodes.IndexAxis): +class Option(Node, workflow_raw_nodes.Option): pass @dataclass -class SpaceAxis(Node, workflow_raw_nodes.SpaceAxis): - pass - - -@dataclass -class TimeAxis(Node, workflow_raw_nodes.TimeAxis): - pass - - -@dataclass -class InputSpec(Node, workflow_raw_nodes.InputSpec): - pass - - -@dataclass -class OptionSpec(Node, workflow_raw_nodes.OptionSpec): - pass - - -@dataclass -class OutputSpec(Node, workflow_raw_nodes.OutputSpec): - pass - - -@dataclass -class Step(Node, workflow_raw_nodes.Step): +class Output(Node, workflow_raw_nodes.Output): pass @dataclass class Workflow(workflow_raw_nodes.Workflow, RDF): - inputs_spec: List[InputSpec] = missing - options_spec: List[OptionSpec] = missing - outputs_spec: List[OutputSpec] = missing - steps: List[Step] = missing - test_steps: List[Step] = missing + inputs: List[Input] = missing + options: List[Option] = missing + outputs: List[Output] = missing From 540a1945d9e39a4410eb823dad2627f4b5c26d7e Mon Sep 17 00:00:00 2001 From: fynnbe Date: Wed, 1 Feb 2023 19:47:18 +0100 Subject: [PATCH 11/19] ImportedSource -> ImportedCallable --- .../_model_adapters/_pytorch_model_adapter.py | 2 +- bioimageio/core/resource_io/nodes.py | 8 ++++---- bioimageio/core/resource_io/utils.py | 8 ++++---- tests/resource_io/test_utils.py | 2 +- 4 files changed, 10 insertions(+), 10 deletions(-) diff --git a/bioimageio/core/prediction_pipeline/_model_adapters/_pytorch_model_adapter.py b/bioimageio/core/prediction_pipeline/_model_adapters/_pytorch_model_adapter.py index f47aa1d7..b3709f30 100644 --- a/bioimageio/core/prediction_pipeline/_model_adapters/_pytorch_model_adapter.py +++ b/bioimageio/core/prediction_pipeline/_model_adapters/_pytorch_model_adapter.py @@ -55,7 +55,7 @@ def _unload(self) -> None: def get_nn_instance(model_node: nodes.Model, **kwargs): weight_spec = model_node.weights.get("pytorch_state_dict") assert weight_spec is not None - assert isinstance(weight_spec.architecture, nodes.ImportedSource) + assert isinstance(weight_spec.architecture, nodes.ImportedCallable) model_kwargs = weight_spec.kwargs joined_kwargs = {} if model_kwargs is missing else dict(model_kwargs) joined_kwargs.update(kwargs) diff --git a/bioimageio/core/resource_io/nodes.py b/bioimageio/core/resource_io/nodes.py index 47e2035f..03144ff4 100644 --- a/bioimageio/core/resource_io/nodes.py +++ b/bioimageio/core/resource_io/nodes.py @@ -124,11 +124,11 @@ def __post_init__(self): @dataclass -class ImportedSource(Node): - factory: Callable +class ImportedCallable(Node): + call: Callable def __call__(self, *args, **kwargs): - return self.factory(*args, **kwargs) + return self.call(*args, **kwargs) @dataclass @@ -144,7 +144,7 @@ class OnnxWeightsEntry(Node, model_raw_nodes.OnnxWeightsEntry): @dataclass class PytorchStateDictWeightsEntry(Node, model_raw_nodes.PytorchStateDictWeightsEntry): source: Path = missing - architecture: Union[_Missing, ImportedSource] = missing + architecture: Union[_Missing, ImportedCallable] = missing @dataclass diff --git a/bioimageio/core/resource_io/utils.py b/bioimageio/core/resource_io/utils.py index 430d30fc..81bf46e8 100644 --- a/bioimageio/core/resource_io/utils.py +++ b/bioimageio/core/resource_io/utils.py @@ -60,21 +60,21 @@ def __enter__(self): def __exit__(self, exc_type, exc_value, traceback): sys.path.remove(self.path) - def transform_LocalImportableModule(self, node: raw_nodes.LocalImportableModule) -> nodes.ImportedSource: + def transform_LocalCallableFromModule(self, node: raw_nodes.LocalCallableFromModule) -> nodes.ImportedCallable: with self.TemporaryInsertionIntoPythonPath(str(node.root_path)): module = importlib.import_module(node.module_name) - return nodes.ImportedSource(factory=getattr(module, node.callable_name)) + return nodes.ImportedCallable(call=getattr(module, node.callable_name)) @staticmethod - def transform_ResolvedImportableSourceFile(node: raw_nodes.ResolvedImportableSourceFile) -> nodes.ImportedSource: + def transform_ResolvedImportableSourceFile(node: raw_nodes.ResolvedImportableSourceFile) -> nodes.ImportedCallable: module_path = resolve_source(node.source_file) module_name = f"module_from_source.{module_path.stem}" importlib_spec = importlib.util.spec_from_file_location(module_name, module_path) assert importlib_spec is not None dep = importlib.util.module_from_spec(importlib_spec) importlib_spec.loader.exec_module(dep) # type: ignore # todo: possible to use "loader.load_module"? - return nodes.ImportedSource(factory=getattr(dep, node.callable_name)) + return nodes.ImportedCallable(call=getattr(dep, node.callable_name)) class RawNodeTypeTransformer(NodeTransformer): diff --git a/tests/resource_io/test_utils.py b/tests/resource_io/test_utils.py index 30889a1d..a16252ef 100644 --- a/tests/resource_io/test_utils.py +++ b/tests/resource_io/test_utils.py @@ -13,7 +13,7 @@ def test_resolve_import_path(tmpdir): node = raw_nodes.ImportableSourceFile(source_file=source_file, callable_name="Foo") uri_transformed = utils.UriNodeTransformer(root_path=tmpdir).transform(node) source_transformed = utils.SourceNodeTransformer().transform(uri_transformed) - assert isinstance(source_transformed, nodes.ImportedSource) + assert isinstance(source_transformed, nodes.ImportedCallable) Foo = source_transformed.factory assert Foo.__name__ == "Foo" assert isinstance(Foo, type) From a7d7ae257522ba2569d0e58b3f6e0a87f79126cc Mon Sep 17 00:00:00 2001 From: fynnbe Date: Wed, 1 Feb 2023 20:09:35 +0100 Subject: [PATCH 12/19] SourceNodeTransformer -> CallableNodeTransformer --- bioimageio/core/resource_io/utils.py | 8 ++++---- tests/resource_io/test_utils.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/bioimageio/core/resource_io/utils.py b/bioimageio/core/resource_io/utils.py index 81bf46e8..3463fb39 100644 --- a/bioimageio/core/resource_io/utils.py +++ b/bioimageio/core/resource_io/utils.py @@ -44,9 +44,9 @@ def visit_WindowsPath(self, leaf: pathlib.WindowsPath): self._visit_source(leaf) -class SourceNodeTransformer(NodeTransformer): +class CallableNodeTransformer(NodeTransformer): """ - Imports all source callables + Import all callables note: Requires previous transformation by UriNodeTransformer """ @@ -67,7 +67,7 @@ def transform_LocalCallableFromModule(self, node: raw_nodes.LocalCallableFromMod return nodes.ImportedCallable(call=getattr(module, node.callable_name)) @staticmethod - def transform_ResolvedImportableSourceFile(node: raw_nodes.ResolvedImportableSourceFile) -> nodes.ImportedCallable: + def transform_ResolvedCallableFromSourceFile(node: raw_nodes.ResolvedCallableFromSourceFile) -> nodes.ImportedCallable: module_path = resolve_source(node.source_file) module_name = f"module_from_source.{module_path.stem}" importlib_spec = importlib.util.spec_from_file_location(module_name, module_path) @@ -109,6 +109,6 @@ def resolve_raw_node( ) -> 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) - rd = SourceNodeTransformer().transform(rd) + rd = CallableNodeTransformer().transform(rd) rd = RawNodeTypeTransformer(nodes_module).transform(rd) return rd diff --git a/tests/resource_io/test_utils.py b/tests/resource_io/test_utils.py index a16252ef..1e0be92d 100644 --- a/tests/resource_io/test_utils.py +++ b/tests/resource_io/test_utils.py @@ -12,7 +12,7 @@ def test_resolve_import_path(tmpdir): (tmpdir / str(source_file)).write_text("class Foo: pass", encoding="utf8") node = raw_nodes.ImportableSourceFile(source_file=source_file, callable_name="Foo") uri_transformed = utils.UriNodeTransformer(root_path=tmpdir).transform(node) - source_transformed = utils.SourceNodeTransformer().transform(uri_transformed) + source_transformed = utils.CallableNodeTransformer().transform(uri_transformed) assert isinstance(source_transformed, nodes.ImportedCallable) Foo = source_transformed.factory assert Foo.__name__ == "Foo" From efb5f5194abc5f7744353708b73e8f332b0c4d98 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Thu, 9 Feb 2023 12:35:09 +0100 Subject: [PATCH 13/19] remove workflow operators --- bioimageio/core/workflow/__init__.py | 0 .../core/workflow/operators/__init__.py | 4 - bioimageio/core/workflow/operators/_assert.py | 8 - .../core/workflow/operators/_generate.py | 15 -- bioimageio/core/workflow/operators/_run.py | 202 ------------------ .../core/workflow/operators/_various.py | 57 ----- 6 files changed, 286 deletions(-) delete mode 100644 bioimageio/core/workflow/__init__.py delete mode 100644 bioimageio/core/workflow/operators/__init__.py delete mode 100644 bioimageio/core/workflow/operators/_assert.py delete mode 100644 bioimageio/core/workflow/operators/_generate.py delete mode 100644 bioimageio/core/workflow/operators/_run.py delete mode 100644 bioimageio/core/workflow/operators/_various.py diff --git a/bioimageio/core/workflow/__init__.py b/bioimageio/core/workflow/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/bioimageio/core/workflow/operators/__init__.py b/bioimageio/core/workflow/operators/__init__.py deleted file mode 100644 index 87e447e4..00000000 --- a/bioimageio/core/workflow/operators/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from ._assert import assert_shape -from ._generate import generate_random_uniform_tensor -from ._run import run_model_inference, run_workflow -from ._various import binarize, load_tensors, log, select_outputs diff --git a/bioimageio/core/workflow/operators/_assert.py b/bioimageio/core/workflow/operators/_assert.py deleted file mode 100644 index fdf54302..00000000 --- a/bioimageio/core/workflow/operators/_assert.py +++ /dev/null @@ -1,8 +0,0 @@ -from typing import Sequence - -import xarray as xr - - -def assert_shape(tensor: xr.DataArray, shape: Sequence[int]) -> xr.DataArray: - assert tensor.shape == tuple(shape) - return tensor diff --git a/bioimageio/core/workflow/operators/_generate.py b/bioimageio/core/workflow/operators/_generate.py deleted file mode 100644 index c5f09237..00000000 --- a/bioimageio/core/workflow/operators/_generate.py +++ /dev/null @@ -1,15 +0,0 @@ -from typing import Sequence, Union - -import numpy as np -import xarray as xr - - -def generate_random_uniform_tensor( - shape: Sequence[Union[int, str]], axes: Sequence[str], *, low: Union[int, float] = 0, high: Union[int, float] = 1 -) -> xr.DataArray: - """generate a tensor with uniformly distributed samples in the interval [low, high) - Returns: - xr.DataArray: random tensor - """ - assert len(shape) == len(axes) - return xr.DataArray(np.random.uniform(low=low, high=high, size=[int(s) for s in shape]), dims=tuple(axes)) diff --git a/bioimageio/core/workflow/operators/_run.py b/bioimageio/core/workflow/operators/_run.py deleted file mode 100644 index d803b88f..00000000 --- a/bioimageio/core/workflow/operators/_run.py +++ /dev/null @@ -1,202 +0,0 @@ -from dataclasses import dataclass -from os import PathLike -from typing import Any, Dict, Generator, IO, List, Optional, Sequence, Tuple, Union - -import numpy as np -import xarray as xr -from marshmallow import missing - -from bioimageio.core import load_resource_description -from bioimageio.core.prediction_pipeline import create_prediction_pipeline -from bioimageio.core.resource_io import nodes -from bioimageio.spec import load_raw_resource_description -from bioimageio.spec.model import raw_nodes -from bioimageio.spec.shared.raw_nodes import ResourceDescription as RawResourceDescription - - -def run_model_inference( - rdf_source: Union[dict, PathLike, IO, str, bytes, raw_nodes.URI, RawResourceDescription], - *tensors, - enable_preprocessing: bool = True, - enable_postprocessing: bool = True, - devices: Optional[Sequence[str]] = None, -) -> List[xr.DataArray]: - """run model inference - - Returns: - list: model outputs - """ - model = load_raw_resource_description(rdf_source, update_to_format="latest") - assert isinstance(model, raw_nodes.Model) - # remove pre-/postprocessing if specified - if not enable_preprocessing: - for ipt in model.inputs: - if ipt.preprocessing: - ipt.preprocessing = missing - if not enable_postprocessing: - for ipt in model.outputs: - if ipt.postprocessing: - ipt.postprocessing = missing - - with create_prediction_pipeline(model, devices=devices) as pred_pipeline: - return pred_pipeline.forward(*tensors) - - -def run_workflow( - rdf_source: Union[dict, PathLike, IO, str, bytes, raw_nodes.URI, RawResourceDescription], - inputs: Sequence = tuple(), - options: Dict[str, Any] = None, -) -> tuple: - outputs = tuple() - for state in _iterate_workflow_steps_impl(rdf_source, test_steps=False, inputs=inputs, options=options): - outputs = state.outputs - - return outputs - - -def run_workflow_test( - rdf_source: Union[dict, PathLike, IO, str, bytes, raw_nodes.URI, RawResourceDescription], -) -> tuple: - outputs = tuple() - for state in _iterate_workflow_steps_impl(rdf_source, test_steps=True): - outputs = state.outputs - - return outputs - - -@dataclass -class WorkflowState: - wf_inputs: Dict[str, Any] - wf_options: Dict[str, Any] - inputs: tuple - outputs: tuple - named_outputs: Dict[str, Any] - - -def iterate_workflow_steps( - rdf_source: Union[dict, PathLike, IO, str, bytes, raw_nodes.URI, RawResourceDescription], - *, - inputs: Sequence = tuple(), - options: Dict[str, Any] = None, -) -> Generator[WorkflowState]: - yield from _iterate_workflow_steps_impl(rdf_source, inputs=inputs, options=options, test_steps=False) - - -def iterate_test_workflow_steps( - rdf_source: Union[dict, PathLike, IO, str, bytes, raw_nodes.URI, RawResourceDescription] -) -> Generator[WorkflowState]: - yield from _iterate_workflow_steps_impl(rdf_source, test_steps=True) - - -def _iterate_workflow_steps_impl( - rdf_source: Union[dict, PathLike, IO, str, bytes, raw_nodes.URI, RawResourceDescription], - *, - test_steps: bool, - inputs: Sequence = tuple(), - options: Dict[str, Any] = None, -) -> Generator[WorkflowState]: - import bioimageio.core.workflow.operators as ops - - workflow = load_resource_description(rdf_source) - assert isinstance(workflow, nodes.Workflow) - wf_options: Dict[str, Any] = {opt.name: opt.default for opt in workflow.options_spec} - if test_steps: - assert not inputs - assert not options - wf_inputs: Dict[str, Any] = {} - steps = workflow.test_steps - else: - if not len(workflow.inputs_spec) == len(inputs): - raise ValueError(f"Expected {len(workflow.inputs_spec)} inputs, but got {len(inputs)}.") - - wf_inputs = {ipt_spec.name: ipt for ipt_spec, ipt in zip(workflow.inputs_spec, inputs)} - for k, v in options.items(): - if k not in wf_options: - raise ValueError(f"Got unknown option {k}, expected one of {set(wf_options)}.") - - wf_options[k] = v - - steps = workflow.steps - - named_outputs = {} # for later referencing - - def map_ref(value): - assert isinstance(workflow, nodes.Workflow) - if isinstance(value, str) and value.startswith("${{") and value.endswith("}}"): - ref = value[4:-2].strip() - if ref.startswith("self.inputs."): - ref = ref[len("self.inputs.") :] - if ref not in wf_inputs: - raise ValueError(f"Invalid workflow input reference {value}.") - - return wf_inputs[ref] - elif ref.startswith("self.options."): - ref = ref[len("self.options.") :] - if ref not in wf_options: - raise ValueError(f"Invalid workflow option reference {value}.") - - return wf_options[ref] - elif ref == "self.rdf_source": - assert workflow.rdf_source is not missing - return str(workflow.rdf_source) - elif ref in named_outputs: - return named_outputs[ref] - else: - raise ValueError(f"Invalid reference {value}.") - else: - return value - - # implicit inputs to a step are the outputs of the previous step. - # For the first step these are the workflow inputs. - outputs = tuple(inputs) - for step in steps: - if not hasattr(ops, step.op): - raise NotImplementedError(f"{step.op} not implemented in {ops}") - - op = getattr(ops, step.op) - if step.inputs is missing: - inputs = outputs - else: - inputs = tuple(map_ref(ipt) for ipt in step.inputs) - - assert isinstance(inputs, tuple) - options = {k: map_ref(v) for k, v in (step.options or {}).items()} - outputs = op(*inputs, **options) - if not isinstance(outputs, tuple): - outputs = (outputs,) - - if step.outputs: - assert step.id is not missing - if len(step.outputs) != len(outputs): - raise ValueError( - f"Got {len(step.outputs)} step output name{'s' if len(step.outputs) > 1 else ''} ({step.id}.outputs), " - f"but op {step.op} returned {len(outputs)} outputs." - ) - - named_outputs.update({f"{step.id}.outputs.{out_name}": out for out_name, out in zip(step.outputs, outputs)}) - - yield WorkflowState( - wf_inputs=wf_inputs, wf_options=wf_options, inputs=inputs, outputs=outputs, named_outputs=named_outputs - ) - if len(workflow.outputs_spec) != len(outputs): - raise ValueError(f"Expected {len(workflow.outputs_spec)} outputs from last step, but got {len(outputs)}.") - - def tensor_as_xr(tensor, axes: Sequence[nodes.Axis]): - spec_axes = [a.name or a.type for a in axes] - if isinstance(tensor, xr.DataArray): - if list(tensor.dims) != spec_axes: - raise ValueError( - f"Last workflow step returned xarray.DataArray with dims {tensor.dims}, but expected dims {spec_axes}." - ) - - return tensor - else: - return xr.DataArray(tensor, dims=spec_axes) - - outputs = tuple( - tensor_as_xr(out, out_spec.axes) if out_spec.type == "tensor" else out - for out_spec, out in zip(workflow.outputs_spec, outputs) - ) - yield WorkflowState( - wf_inputs=wf_inputs, wf_options=wf_options, inputs=inputs, outputs=outputs, named_outputs=named_outputs - ) diff --git a/bioimageio/core/workflow/operators/_various.py b/bioimageio/core/workflow/operators/_various.py deleted file mode 100644 index 37bb3520..00000000 --- a/bioimageio/core/workflow/operators/_various.py +++ /dev/null @@ -1,57 +0,0 @@ -import logging -from typing import List, Sequence, Tuple - -import numpy as np -import xarray as xr -from imageio import imread - -logger = logging.getLogger(__name__) - - -def binarize(tensor: xr.DataArray, threshold: float): - return tensor > threshold - - -def select_outputs(*args) -> Tuple: - """helper to select workflow outputs (to be used as a final step in a workflow) - - Returns: - tuple: selected outputs (inputs to this op) - - """ - - return args - - -def log(*args, log_level: int = logging.INFO, **kwargs) -> Tuple: - """log any key word arguments (kwargs/options) - - Returns: - tuple: positional inputs to this op - - """ - for k, v in kwargs.items(): - logger.log( - log_level, - f"{k}: %s", - f"{v.shape} mean: {v.mean().item():.4f} std: {v.std().item():.4f}" - if isinstance(v, (np.ndarray, xr.DataArray)) - else v, - ) - - return args - - -def load_tensors(sources: List[str], axes: Sequence[str]) -> List[xr.DataArray]: - """load tensors""" - assert len(sources) == len(axes) - tensors = [] - for source, ax in zip(sources, axes): - if source.split(".")[-1] == ".npy": - data = np.load(str(source)) - else: - data = imread(source) - - tensors.append(xr.DataArray(data, dims=ax)) - - return tensors From 1f402fe4b5fdea66f2abcae35dec3361539727cf Mon Sep 17 00:00:00 2001 From: fynnbe Date: Thu, 9 Feb 2023 12:47:46 +0100 Subject: [PATCH 14/19] remove workflow run command moved to bioimageio.workflows --- bioimageio/core/__main__.py | 85 ------------------------------------- 1 file changed, 85 deletions(-) diff --git a/bioimageio/core/__main__.py b/bioimageio/core/__main__.py index 8be385ff..8aa498cb 100644 --- a/bioimageio/core/__main__.py +++ b/bioimageio/core/__main__.py @@ -16,7 +16,6 @@ from bioimageio.core.common import TestSummary from bioimageio.core.image_helper import load_image, save_image from bioimageio.core.resource_io import nodes -from bioimageio.core.workflow.operators import run_workflow from bioimageio.spec.__main__ import app, help_version as help_version_spec from bioimageio.spec.model.raw_nodes import WeightsFormat from bioimageio.spec.workflow.raw_nodes import Workflow @@ -196,7 +195,6 @@ def predict_image( weight_format: Optional[WeightFormatEnum] = typer.Option(None, help="The weight format to use."), devices: Optional[List[str]] = typer.Option(None, help="Devices for running the model."), ): - if isinstance(padding, str): padding = json.loads(padding.replace("'", '"')) assert isinstance(padding, dict) @@ -314,88 +312,5 @@ def convert_keras_weights_to_tensorflow( ) -@app.command(context_settings=dict(allow_extra_args=True, ignore_unknown_options=True), add_help_option=False) -def run( - rdf_source: str = typer.Argument(..., help="BioImage.IO RDF id/url/path."), - *, - output_folder: Path = Path("outputs"), - output_tensor_extension: str = ".npy", - ctx: typer.Context, -): - resource = load_raw_resource_description(rdf_source, update_to_format="latest") - if not isinstance(resource, Workflow): - raise NotImplementedError(f"Non-workflow RDFs not yet supported (got type {resource.type})") - - map_type = dict( - any=str, - boolean=bool, - float=float, - int=int, - list=str, - string=str, - ) - wf = resource - parser = ArgumentParser(description=f"CLI for {wf.name}") - - # replicate typer args to show up in help - parser.add_argument( - metavar="rdf-source", - dest="rdf_source", - help="BioImage.IO RDF id/url/path. The optional arguments below are RDF specific.", - ) - parser.add_argument( - metavar="output-folder", dest="output_folder", help="Folder to save outputs to.", default=Path("outputs") - ) - parser.add_argument( - metavar="output-tensor-extension", - dest="output_tensor_extension", - help="Output tensor extension.", - default=".npy", - ) - - def add_param_args(params): - for param in params: - argument_kwargs = {} - if param.type == "tensor": - argument_kwargs["type"] = partial(load_image, axes=[a.name or a.type for a in param.axes]) - else: - argument_kwargs["type"] = map_type[param.type] - - if param.type == "list": - argument_kwargs["nargs"] = "*" - - argument_kwargs["help"] = param.description or "" - if hasattr(param, "default"): - argument_kwargs["default"] = param.default - else: - argument_kwargs["required"] = True - - argument_kwargs["metavar"] = param.name[0].capitalize() - parser.add_argument("--" + param.name.replace("_", "-"), **argument_kwargs) - - def prepare_parameter(value, param: Union[nodes.InputSpec, nodes.OptionSpec]): - if param.type == "tensor": - return load_image(value, [a.name or a.type for a in param.axes]) - else: - return value - - add_param_args(wf.inputs_spec) - add_param_args(wf.options_spec) - args = parser.parse_args([rdf_source, str(output_folder), output_tensor_extension] + list(ctx.args)) - outputs = run_workflow( - rdf_source, - inputs=[prepare_parameter(getattr(args, ipt.name), ipt) for ipt in wf.inputs_spec], - options={opt.name: prepare_parameter(getattr(args, opt.name), opt) for opt in wf.options_spec}, - ) - output_folder.mkdir(parents=True, exist_ok=True) - for out_spec, out in zip(wf.outputs_spec, outputs): - out_path = output_folder / out_spec.name - if out_spec.type == "tensor": - save_image(out_path.with_suffix(output_tensor_extension), out) - else: - with out_path.with_suffix(".json").open("w") as f: - json.dump(out, f) - - if __name__ == "__main__": app() From 904e38a3dd5805f15e25d30c1d5e22e995d8fdcc Mon Sep 17 00:00:00 2001 From: fynnbe Date: Thu, 9 Feb 2023 20:47:49 +0100 Subject: [PATCH 15/19] fix transform_input_image --- bioimageio/core/image_helper.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/bioimageio/core/image_helper.py b/bioimageio/core/image_helper.py index fdbead71..d67ae990 100644 --- a/bioimageio/core/image_helper.py +++ b/bioimageio/core/image_helper.py @@ -39,7 +39,7 @@ def transform_input_image(image: np.ndarray, tensor_axes: Sequence[str], image_a # instead of 'b' we might want 'batch', etc... axis_letter_map = { letter: name - for letter, name in {"b": "batch", "c": "channel", "i": "index", "t": "time"} + for letter, name in {"b": "batch", "c": "channel", "i": "index", "t": "time"}.items() if name in tensor_axes # only do this mapping if the full name is in the desired tensor_axes } image_axes = tuple(axis_letter_map.get(a, a) for a in image_axes) @@ -167,7 +167,6 @@ def pad(image, axes: Sequence[str], padding, pad_right=True) -> Tuple[np.ndarray pad_width = [] crop = {} for ax, dlen, pr in zip(axes, image.shape, pad_right): - if ax in "zyx": pad_to = padding_[ax] From 4a3c4f65818876d29a6be8519b36cf085c85daa1 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Thu, 9 Feb 2023 23:20:47 +0100 Subject: [PATCH 16/19] fix tests --- tests/build_spec/test_add_weights.py | 4 ++++ tests/resource_io/test_utils.py | 2 +- tests/test_cli.py | 4 ++-- 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/build_spec/test_add_weights.py b/tests/build_spec/test_add_weights.py index 2f8300b0..482d167f 100644 --- a/tests/build_spec/test_add_weights.py +++ b/tests/build_spec/test_add_weights.py @@ -1,4 +1,7 @@ import os + +import pytest + from bioimageio.core import export_resource_package, load_raw_resource_description, load_resource_description from bioimageio.core.resource_tests import test_model as _test_model @@ -45,5 +48,6 @@ def test_add_torchscript(unet2d_nuclei_broad_model, tmp_path): _test_add_weights(unet2d_nuclei_broad_model, tmp_path, "pytorch_state_dict", "torchscript") +@pytest.mark.skipif(pytest.skip_onnx, reason="onnx") def test_add_onnx(unet2d_nuclei_broad_model, tmp_path): _test_add_weights(unet2d_nuclei_broad_model, tmp_path, "pytorch_state_dict", "onnx", opset_version=12) diff --git a/tests/resource_io/test_utils.py b/tests/resource_io/test_utils.py index 1e0be92d..7174a8d5 100644 --- a/tests/resource_io/test_utils.py +++ b/tests/resource_io/test_utils.py @@ -14,7 +14,7 @@ def test_resolve_import_path(tmpdir): uri_transformed = utils.UriNodeTransformer(root_path=tmpdir).transform(node) source_transformed = utils.CallableNodeTransformer().transform(uri_transformed) assert isinstance(source_transformed, nodes.ImportedCallable) - Foo = source_transformed.factory + Foo = source_transformed.call assert Foo.__name__ == "Foo" assert isinstance(Foo, type) diff --git a/tests/test_cli.py b/tests/test_cli.py index c0de99d4..55f205f5 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -98,7 +98,7 @@ def _test_cli_predict_images(model, tmp_path, extra_kwargs=None): expected_outputs.append(out_folder / f"im-{i}.npy") input_pattern = str(in_folder / "*.npy") - cmd = ["bioimageio", "predict-images", model, input_pattern, str(out_folder)] + cmd = ["bioimageio", "predict-images", str(model), input_pattern, str(out_folder)] if extra_kwargs is not None: cmd.extend(extra_kwargs) ret = run_subprocess(cmd) @@ -126,7 +126,7 @@ def test_torch_to_torchscript(unet2d_nuclei_broad_model, tmp_path): assert out_path.exists() -@pytest.mark.skipif(pytest.skip_onnx, reason="requires torch and onnx") +@pytest.mark.skipif(pytest.skip_onnx, reason="requires onnx") def test_torch_to_onnx(unet2d_nuclei_broad_model, tmp_path): out_path = tmp_path.with_suffix(".onnx") ret = run_subprocess(["bioimageio", "convert-torch-weights-to-onnx", str(unet2d_nuclei_broad_model), str(out_path)]) From 0a2a27f8de774955dcce6545f909d48db30003be Mon Sep 17 00:00:00 2001 From: fynnbe Date: Fri, 3 Mar 2023 11:49:55 +0100 Subject: [PATCH 17/19] add root_path arg to resolve_raw_node() --- bioimageio/core/resource_io/utils.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/bioimageio/core/resource_io/utils.py b/bioimageio/core/resource_io/utils.py index fd16a9c9..c8582141 100644 --- a/bioimageio/core/resource_io/utils.py +++ b/bioimageio/core/resource_io/utils.py @@ -79,7 +79,9 @@ def transform_LocalCallableFromModule(self, node: raw_nodes.LocalCallableFromMod return nodes.ImportedCallable(call=getattr(module, node.callable_name)) @staticmethod - def transform_ResolvedCallableFromSourceFile(node: raw_nodes.ResolvedCallableFromSourceFile) -> nodes.ImportedCallable: + def transform_ResolvedCallableFromSourceFile( + node: raw_nodes.ResolvedCallableFromSourceFile, + ) -> nodes.ImportedCallable: module_path = resolve_source(node.source_file) module_name = f"module_from_source.{module_path.stem}" importlib_spec = importlib.util.spec_from_file_location(module_name, module_path) @@ -119,10 +121,15 @@ def all_sources_available( def resolve_raw_node( - raw_rd: GenericRawNode, nodes_module: typing.Any, uri_only_if_in_package: bool = True + raw_rd: GenericRawNode, + nodes_module: typing.Any, + uri_only_if_in_package: bool = True, + root_path: typing.Optional[pathlib.Path] = None, ) -> 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) + rd = UriNodeTransformer( + root_path=root_path or raw_rd.root_path, uri_only_if_in_package=uri_only_if_in_package + ).transform(raw_rd) rd = CallableNodeTransformer().transform(rd) rd = RawNodeTypeTransformer(nodes_module).transform(rd) return rd From 2d51272932e8f94914e0d182ccb333081ef1eda1 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Fri, 3 Mar 2023 13:11:58 +0100 Subject: [PATCH 18/19] fix load_image takes axis tuple --- bioimageio/core/image_helper.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/bioimageio/core/image_helper.py b/bioimageio/core/image_helper.py index d67ae990..5f2389f5 100644 --- a/bioimageio/core/image_helper.py +++ b/bioimageio/core/image_helper.py @@ -13,13 +13,13 @@ # -def transform_input_image(image: np.ndarray, tensor_axes: Sequence[str], image_axes: Optional[str] = None): +def transform_input_image(image: np.ndarray, tensor_axes: Sequence[str], image_axes: Optional[Sequence[str]] = None): """Transform input image into output tensor with desired axes. Args: image: the input image tensor_axes: the desired tensor axes - input_axes: the axes of the input image (optional) + image_axes: the axes of the input image (optional) """ # if the image axes are not given deduce them from the required axes and image shape if image_axes is None: @@ -105,7 +105,7 @@ def load_image(in_path, axes: Sequence[str]) -> DataArray: is_volume = "z" in axes im = imageio.volread(in_path) if is_volume else imageio.imread(in_path) im = transform_input_image(im, axes) - return DataArray(im, dims=axes) + return DataArray(im, dims=tuple(axes)) def load_tensors(sources, tensor_specs: List[Union[InputTensor, OutputTensor]]) -> List[DataArray]: From a4eb0f5cd591d1e80fcb98e8c7a1a9cbb1f0d16b Mon Sep 17 00:00:00 2001 From: fynnbe Date: Wed, 15 Mar 2023 11:17:20 +0100 Subject: [PATCH 19/19] tidy up imports --- bioimageio/core/__main__.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/bioimageio/core/__main__.py b/bioimageio/core/__main__.py index 8aa498cb..1ae89a2b 100644 --- a/bioimageio/core/__main__.py +++ b/bioimageio/core/__main__.py @@ -3,22 +3,17 @@ import os import sys import warnings -from argparse import ArgumentParser -from functools import partial from glob import glob from pathlib import Path from pprint import pformat -from typing import List, Optional, Union +from typing import List, Optional import typer -from bioimageio.core import __version__, commands, load_raw_resource_description, prediction, resource_tests +from bioimageio.core import __version__, commands, prediction, resource_tests from bioimageio.core.common import TestSummary -from bioimageio.core.image_helper import load_image, save_image -from bioimageio.core.resource_io import nodes from bioimageio.spec.__main__ import app, help_version as help_version_spec from bioimageio.spec.model.raw_nodes import WeightsFormat -from bioimageio.spec.workflow.raw_nodes import Workflow try: from typing import get_args