Update submodules to the latest changes

- bioimageio.spec==0.5.3.post4
- bioimageio.core==0.6.7

.gitignore

@@ -9,3 +9,4 @@ tiktorch/__pycache/
 /.#wrapper.py#
 .py~
 unet_sample.zip
+build

README.md

@@ -18,7 +18,7 @@ Reasons to install the server component:
 
 We have a how-to for the set up of a tiktorch prediction server (powerful machine with nvidia GPU) to connect to from your client (e.g. ilastik on a laptop): [Installation](#installation).
 
-If you are interested in running Model Zoo networks from Python directly, have a look at [`bioimageio.core`](https://github.com/bioimage-io/core-bioimage-io-python), where we have an [example notebook](https://github.com/bioimage-io/core-bioimage-io-python/blob/main/example/bioimageio-core-usage.ipynb) to get you started.
+If you are interested in running Model Zoo networks from Python directly, have a look at [`bioimageio.core`](https://github.com/bioimage-io/core-bioimage-io-python), where we have an [example notebook](https://github.com/bioimage-io/core-bioimage-io-python/blob/main/example/model_usage.ipynb) to get you started.
 
 
 ## Installation

setup.py

@@ -27,8 +27,8 @@
     ],
     packages=find_packages(exclude=["tests"]),  # Required
     install_requires=[
-        "bioimageio.spec==0.4.9.post5",
-        "bioimageio.core==0.5.11",
+        "bioimageio.spec==0.5.3.post4",
+        "bioimageio.core==0.6.7",
         "grpcio>=1.31",
         "numpy<2",  # pytorch 2.2.2-py3.9_0 for macos is compiled with numpy 1.*
         "protobuf",

tests/conftest.py

@@ -1,41 +1,22 @@
 import faulthandler
-import io
 import logging.handlers
 import multiprocessing as mp
 import signal
 import sys
 import threading
-from collections import namedtuple
 from os import getenv
-from pathlib import Path
 from random import randint
-from zipfile import ZipFile
+from typing import Generator, List, Tuple
+from unittest.mock import create_autospec, patch
 
 import numpy as np
 import pytest
-from bioimageio.core.resource_io import export_resource_package
+import xarray as xr
+from bioimageio.core import AxisId, PredictionPipeline, Sample, Tensor
+from bioimageio.spec.model import v0_5
+from bioimageio.spec.model.v0_5 import TensorId
 
-TEST_DATA = "data"
-TEST_BIOIMAGEIO_ZIPFOLDER = "unet2d"
-TEST_BIOIMAGEIO_ONNX = "unet2d_onnx"
-TEST_BIOIMAGEIO_DUMMY_EXPLICIT = "dummy"
-TEST_BIOIMAGEIO_DUMMY_EXPLICIT_RDF = f"{TEST_BIOIMAGEIO_DUMMY_EXPLICIT}/Dummy.model.yaml"
-TEST_BIOIMAGEIO_DUMMY_PARAM_RDF = "dummy_param/Dummy.model_param.yaml"
-TEST_BIOIMAGEIO_TENSORFLOW_DUMMY = "dummy_tensorflow"
-TEST_BIOIMAGEIO_TORCHSCRIPT = "unet2d_torchscript"
-
-NNModel = namedtuple("NNModel", ["model", "state"])
-
-
-@pytest.fixture
-def data_path():
-    conf_path = Path(__file__).parent
-    return conf_path / TEST_DATA
-
-
-def read_bytes(filename):
-    with open(filename, "rb") as file:
-        return file.read()
+from tiktorch.server.session import process
 
 
 @pytest.fixture
@@ -84,121 +65,115 @@ def assert_threads_cleanup():
         pytest.fail("Threads still running:\n\t%s" % "\n\t".join(running_threads))
 
 
-@pytest.fixture
-def bioimageio_model_bytes(data_path):
-    rdf_source = data_path / TEST_BIOIMAGEIO_ZIPFOLDER / "UNet2DNucleiBroad.model.yaml"
-    data = io.BytesIO()
-    export_resource_package(rdf_source, output_path=data)
-    return data
-
-
-@pytest.fixture
-def bioimageio_model_zipfile(bioimageio_model_bytes):
-    with ZipFile(bioimageio_model_bytes, mode="r") as zf:
-        yield zf
-
+MockedPredictionPipeline = Generator[Tuple[PredictionPipeline, Sample], None, None]
 
-@pytest.fixture
-def bioimageio_dummy_model_filepath(data_path, tmpdir):
-    bioimageio_net_dir = Path(data_path) / TEST_BIOIMAGEIO_DUMMY_EXPLICIT
-    path = tmpdir / "dummy_model.zip"
-
-    with ZipFile(path, mode="w") as zip_model:
-        for f_path in bioimageio_net_dir.iterdir():
-            if str(f_path.name).startswith("__"):
-                continue
-
-            with f_path.open(mode="rb") as f:
-                zip_model.writestr(f_path.name, f.read())
-
-    return path
-
-
-@pytest.fixture
-def bioimageio_dummy_explicit_model_bytes(data_path):
-    rdf_source = data_path / TEST_BIOIMAGEIO_DUMMY_EXPLICIT_RDF
-    return _bioimageio_package(rdf_source)
-
-
-@pytest.fixture
-def bioimageio_dummy_param_model_bytes(data_path):
-    rdf_source = data_path / TEST_BIOIMAGEIO_DUMMY_PARAM_RDF
-    return _bioimageio_package(rdf_source)
 
-
-@pytest.fixture(params=[(TEST_BIOIMAGEIO_DUMMY_PARAM_RDF, "param"), (TEST_BIOIMAGEIO_DUMMY_EXPLICIT_RDF, "input")])
-def bioimageio_dummy_model(request, data_path):
-    path, tensor_id = request.param
-    yield _bioimageio_package(data_path / path), tensor_id
-
-
-def _bioimageio_package(rdf_source):
-    data = io.BytesIO()
-    export_resource_package(rdf_source, output_path=data)
-    return data
-
-
-def archive(directory):
-    result = io.BytesIO()
-
-    with ZipFile(result, mode="w") as zip_model:
-
-        def _archive(path_to_archive):
-            for path in path_to_archive.iterdir():
-                if str(path.name).startswith("__"):
-                    continue
-
-                if path.is_dir():
-                    _archive(path)
-
-                else:
-                    with path.open(mode="rb") as f:
-                        zip_model.writestr(str(path).replace(str(directory), ""), f.read())
-
-        _archive(directory)
-
-    return result
-
-
-@pytest.fixture
-def bioimageio_dummy_tensorflow_model_bytes(data_path):
-    rdf_source = data_path / TEST_BIOIMAGEIO_TENSORFLOW_DUMMY / "rdf.yaml"
-    data = io.BytesIO()
-    export_resource_package(rdf_source, output_path=data)
-    return data
-
-
-@pytest.fixture
-def bioimageio_unet2d_onnx_bytes(data_path):
-    bioimageio_net_dir = Path(data_path) / TEST_BIOIMAGEIO_ONNX
-    return archive(bioimageio_net_dir)
-
-
-@pytest.fixture
-def bioimageio_unet2d_onnx_test_data(data_path):
-    bioimageio_net_dir = Path(data_path) / TEST_BIOIMAGEIO_ONNX
-    test_input = bioimageio_net_dir / "test_input.npy"
-    test_output = bioimageio_net_dir / "test_output.npy"
-    return {"test_input": test_input, "test_output": test_output}
+def patched_prediction_pipeline(mocked_prediction_pipeline: PredictionPipeline):
+    return patch.object(process, "_get_prediction_pipeline_from_model_bytes", lambda *args: mocked_prediction_pipeline)
 
 
 @pytest.fixture
-def npy_zeros_file(tmpdir):
-    path = str(tmpdir / "zeros.npy")
-    zeros = np.zeros(shape=(64, 64))
-    np.save(path, zeros)
-    return path
+def bioimage_model_explicit_siso() -> MockedPredictionPipeline:
+    mocked_prediction_pipeline, mocked_output_sample = _bioimage_model_siso(
+        [
+            v0_5.BatchAxis(),
+            v0_5.ChannelAxis(channel_names=["channel1", "channel2"]),
+            v0_5.SpaceInputAxis(id="x", size=10),
+            v0_5.SpaceInputAxis(id="y", size=10),
+        ]
+    )
+    with patched_prediction_pipeline(mocked_prediction_pipeline):
+        yield mocked_prediction_pipeline, mocked_output_sample
 
 
 @pytest.fixture
-def bioimageio_unet2d_torchscript_bytes(data_path):
-    bioimageio_net_dir = Path(data_path) / TEST_BIOIMAGEIO_TORCHSCRIPT
-    return archive(bioimageio_net_dir)
+def bioimage_model_param_siso() -> MockedPredictionPipeline:
+    mocked_prediction_pipeline, mocked_output_sample = _bioimage_model_siso(
+        [
+            v0_5.BatchAxis(),
+            v0_5.ChannelAxis(channel_names=["channel1", "channel2"]),
+            v0_5.SpaceInputAxis(id="x", size=v0_5.ParameterizedSize(min=10, step=2)),
+            v0_5.SpaceInputAxis(id="y", size=v0_5.ParameterizedSize(min=20, step=3)),
+        ]
+    )
+    with patched_prediction_pipeline(mocked_prediction_pipeline):
+        yield mocked_prediction_pipeline, mocked_output_sample
+
+
+def _bioimage_model_siso(input_axes: List[v0_5.InputAxis]) -> Tuple[PredictionPipeline, Sample]:
+    """
+    Mocked bioimageio prediction pipeline with single input single output
+    """
+
+    mocked_input = create_autospec(v0_5.InputTensorDescr)
+    mocked_input.id = "input"
+    mocked_input.axes = input_axes
+    return _bioimage_model([mocked_input])
 
 
 @pytest.fixture
-def bioimageio_unet2d_torchscript_test_data(data_path):
-    bioimageio_net_dir = Path(data_path) / TEST_BIOIMAGEIO_TORCHSCRIPT
-    test_input = bioimageio_net_dir / "test_input.npy"
-    test_output = bioimageio_net_dir / "test_output.npy"
-    return {"test_input": test_input, "test_output": test_output}
+def bioimage_model_miso() -> MockedPredictionPipeline:
+    """
+    Mocked bioimageio prediction pipeline with three inputs single output
+    """
+
+    mocked_input1 = create_autospec(v0_5.InputTensorDescr)
+    mocked_input1.id = "input1"
+    mocked_input1.axes = [
+        v0_5.BatchAxis(),
+        v0_5.ChannelAxis(channel_names=["channel1", "channel2"]),
+        v0_5.SpaceInputAxis(id=AxisId("x"), size=10),
+        v0_5.SpaceInputAxis(id=AxisId("y"), size=v0_5.SizeReference(tensor_id="input3", axis_id="x")),
+    ]
+
+    mocked_input2 = create_autospec(v0_5.InputTensorDescr)
+    mocked_input2.id = "input2"
+    mocked_input2.axes = [
+        v0_5.BatchAxis(),
+        v0_5.ChannelAxis(channel_names=["channel1", "channel2"]),
+        v0_5.SpaceInputAxis(id=AxisId("x"), size=v0_5.ParameterizedSize(min=10, step=2)),
+        v0_5.SpaceInputAxis(id=AxisId("y"), size=v0_5.ParameterizedSize(min=10, step=5)),
+    ]
+
+    mocked_input3 = create_autospec(v0_5.InputTensorDescr)
+    mocked_input3.id = "input3"
+    mocked_input3.axes = [
+        v0_5.BatchAxis(),
+        v0_5.ChannelAxis(channel_names=["channel1", "channel2"]),
+        v0_5.SpaceInputAxis(id="x", size=v0_5.SizeReference(tensor_id="input2", axis_id="x")),
+        v0_5.SpaceInputAxis(id="y", size=10),
+    ]
+
+    mocked_prediction_pipeline, mocked_output_sample = _bioimage_model([mocked_input1, mocked_input2, mocked_input3])
+    with patched_prediction_pipeline(mocked_prediction_pipeline):
+        yield mocked_prediction_pipeline, mocked_output_sample
+
+
+def _bioimage_model(inputs: List[v0_5.InputTensorDescr]) -> Tuple[PredictionPipeline, Sample]:
+    mocked_descr = create_autospec(v0_5.ModelDescr)
+
+    mocked_output = create_autospec(v0_5.OutputTensorDescr)
+    mocked_output.id = "output"
+    mocked_output.axes = [
+        v0_5.BatchAxis(),
+        v0_5.ChannelAxis(channel_names=["channel1", "channel2"]),
+        v0_5.SpaceInputAxis(id=AxisId("x"), size=20),
+        v0_5.SpaceInputAxis(id=AxisId("y"), size=20),
+    ]
+    mocked_descr.inputs = inputs
+    mocked_descr.outputs = [mocked_output]
+
+    mocked_output_sample = Sample(
+        members={
+            TensorId("output"): Tensor.from_xarray(
+                xr.DataArray(np.arange(2 * 20 * 20).reshape((1, 2, 20, 20)), dims=["batch", "channel", "x", "y"])
+            )
+        },
+        id=None,
+        stat={},
+    )
+
+    mocked_prediction_pipeline = create_autospec(PredictionPipeline)
+    mocked_prediction_pipeline.model_description = mocked_descr
+    mocked_prediction_pipeline.predict_sample_without_blocking.return_value = mocked_output_sample
+    return mocked_prediction_pipeline, mocked_output_sample

tests/test_converters.py

@@ -1,9 +1,18 @@
 import numpy as np
 import pytest
 import xarray as xr
+from bioimageio.spec.model.v0_5 import TensorId
 from numpy.testing import assert_array_equal
 
-from tiktorch.converters import Sample, numpy_to_pb_tensor, pb_tensor_to_numpy, pb_tensor_to_xarray, xarray_to_pb_tensor
+from tiktorch.converters import (
+    numpy_to_pb_tensor,
+    pb_tensor_to_numpy,
+    pb_tensor_to_xarray,
+    pb_tensors_to_sample,
+    sample_to_pb_tensors,
+    xarray_to_pb_tensor,
+    xr_tensors_to_sample,
+)
 from tiktorch.proto import inference_pb2
 
 
@@ -167,7 +176,7 @@ def test_should_same_data(self, shape):
 
 
 class TestSample:
-    def test_create_sample_from_pb_tensors(self):
+    def test_pb_tensors_to_sample(self):
         arr_1 = np.arange(32 * 32, dtype=np.int64).reshape(32, 32)
         tensor_1 = inference_pb2.Tensor(
             dtype="int64",
@@ -184,30 +193,30 @@ def test_create_sample_from_pb_tensors(self):
             shape=[inference_pb2.NamedInt(name="x", size=64), inference_pb2.NamedInt(name="y", size=64)],
         )
 
-        sample = Sample.from_pb_tensors([tensor_1, tensor_2])
-        assert len(sample.tensors) == 2
-        assert sample.tensors["input1"].equals(xr.DataArray(arr_1, dims=["x", "y"]))
-        assert sample.tensors["input2"].equals(xr.DataArray(arr_2, dims=["x", "y"]))
+        sample = pb_tensors_to_sample([tensor_1, tensor_2])
+        assert len(sample.members) == 2
+        assert sample.members[TensorId("input1")].data.equals(xr.DataArray(arr_1, dims=["x", "y"]))
+        assert sample.members[TensorId("input2")].data.equals(xr.DataArray(arr_2, dims=["x", "y"]))
 
-    def test_create_sample_from_raw_data(self):
+    def test_xr_tensors_to_sample(self):
         arr_1 = np.arange(32 * 32, dtype=np.int64).reshape(32, 32)
         tensor_1 = xr.DataArray(arr_1, dims=["x", "y"])
         arr_2 = np.arange(64 * 64, dtype=np.int64).reshape(64, 64)
         tensor_2 = xr.DataArray(arr_2, dims=["x", "y"])
         tensors_ids = ["input1", "input2"]
-        actual_sample = Sample.from_xr_tensors(tensors_ids, [tensor_1, tensor_2])
 
-        expected_dict = {tensors_ids[0]: tensor_1, tensors_ids[1]: tensor_2}
-        expected_sample = Sample(expected_dict)
-        assert actual_sample == expected_sample
+        actual_sample = xr_tensors_to_sample(tensors_ids, [tensor_1, tensor_2])
+        assert len(actual_sample.members) == 2
+        assert actual_sample.members[TensorId("input1")].data.equals(tensor_1)
+        assert actual_sample.members[TensorId("input2")].data.equals(tensor_2)
 
     def test_sample_to_pb_tensors(self):
         arr_1 = np.arange(32 * 32, dtype=np.int64).reshape(32, 32)
         tensor_1 = xr.DataArray(arr_1, dims=["x", "y"])
         arr_2 = np.arange(64 * 64, dtype=np.int64).reshape(64, 64)
         tensor_2 = xr.DataArray(arr_2, dims=["x", "y"])
         tensors_ids = ["input1", "input2"]
-        sample = Sample.from_xr_tensors(tensors_ids, [tensor_1, tensor_2])
+        sample = xr_tensors_to_sample(tensors_ids, [tensor_1, tensor_2])
 
         pb_tensor_1 = inference_pb2.Tensor(
             dtype="int64",
@@ -223,5 +232,5 @@ def test_sample_to_pb_tensors(self):
         )
         expected_tensors = [pb_tensor_1, pb_tensor_2]
 
-        actual_tensors = sample.to_pb_tensors()
+        actual_tensors = sample_to_pb_tensors(sample)
         assert expected_tensors == actual_tensors