Add remote Gradio adapter and support multithreaded prediction

.gitignore

@@ -6,6 +6,7 @@
 .env
 .idea/
 .tox/
+Unet 2D Nuclei Broad_example/
 __pycache__/
 bioimageio_cache/
 bioimageio_unzipped_tf_weights/
@@ -15,7 +16,8 @@ coverage.xml
 dist/
 docs/api/
 dogfood/
+example/**/output/
+output.zip
 pkgs/
 site/
 typings/pooch/
-example/**/output/

changelog.md

@@ -1,3 +1,10 @@
+### 0.11.0
+
+- bump spec to 0.5.11.0
+- support ONNX provider choice (as 'devices')
+- added experimental bioimageio.core.backends.gradio_backend
+- improved prediction pipeline and model adapter interfaces
+
 ### 0.10.4
 
 - fix postprocessing order

conda-recipe/meta.yaml

@@ -55,7 +55,7 @@ test:
   requires:
     {% for dep in pyproject['project']['optional-dependencies']['dev'] %}
       {% if 'torch' not in dep %}  # can't install pytorch>=2.8 from conda-forge smh
-      - {{ dep.replace(';python_version<"3.10"', '').lower().replace('_', '-') }}
+      - {{ dep.replace('[mcp]', '').replace(';python_version>="3.10"', '').replace(';python_version<"3.10"', '').lower().replace('_', '-') }}
       {% endif %}
     {% endfor %}
   commands:

conftest.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+from typing import Any
+
+
+def pytest_ignore_collect(collection_path: Path, config: Any) -> bool:
+    if sys.version_info >= (3, 10):
+        return False
+
+    path = str(collection_path).replace("\\", "/")
+    return "/src/bioimageio/core/remote_backends/gradio/" in path

example/dataset_statistics_demo.ipynb

@@ -29,7 +29,6 @@
     "\n",
     "from pprint import pprint\n",
     "\n",
-    "import imageio\n",
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "import xarray as xr\n",
@@ -39,7 +38,6 @@
     "warnings.filterwarnings(\"ignore\")\n",
     "\n",
     "import bioimageio.core\n",
-    "from bioimageio.core.prediction import predict_with_tiling\n",
     "from bioimageio.core.prediction_pipeline import create_prediction_pipeline"
    ]
   },
@@ -196,7 +194,7 @@
    "source": [
     "def process_dataset(pp, dataset):\n",
     "    stats = pp._ipt_stats.compute_measures()[\"per_dataset\"]\n",
-    "    print(f\"initial stats:\")\n",
+    "    print(\"initial stats:\")\n",
     "    pprint(\n",
     "        None\n",
     "        if not stats\n",
@@ -205,7 +203,7 @@
     "    stats = {}\n",
     "    sample_dataset = [{\"input0\": s} for s in dataset]\n",
     "    [pp.apply_preprocessing(s, stats) for s in sample_dataset]\n",
-    "    print(f\"final stats:\")\n",
+    "    print(\"final stats:\")\n",
     "    pprint(\n",
     "        None\n",
     "        if not stats\n",

example/export_cellpose_model/cellpose_original.py

@@ -1,4 +1,6 @@
-"""Run original cellpose model and save an analog input and output for bioimageio tests"""
+"""Run original cellpose model and save an analog input and output for bioimageio tests
+Works on cellpose 4.1, but not since cellpose 4.2
+"""
 
 import os
 from pathlib import Path

example/model_usage.ipynb

@@ -495,7 +495,6 @@
    ],
    "source": [
     "from bioimageio.core.digest_spec import create_sample_for_model\n",
-    "from bioimageio.spec.utils import download\n",
     "\n",
     "input_paths = {ipt.id: ipt.test_tensor.source for ipt in model.inputs}\n",
     "print(f\"input paths: {input_paths}\")\n",

mkdocs.yaml

@@ -87,7 +87,7 @@ plugins:
         python:
           inventories:
             - https://docs.pydantic.dev/latest/objects.inv
-            - https://bioimage-io.github.io/spec-bioimage-io/v0.5.10.2/objects.inv
+            - https://bioimage-io.github.io/spec-bioimage-io/v0.5.11.0/objects.inv
           options:
             annotations_path: source
             backlinks: tree

pyproject.toml

@@ -6,7 +6,7 @@ requires-python = ">=3.9"
 readme = "README.md"
 dynamic = ["version"]
 dependencies = [
-    "bioimageio.spec ==0.5.10.2",
+    "bioimageio.spec ==0.5.11.0",
     "imagecodecs",
     "imageio>=2.10",
     "loguru",
@@ -52,9 +52,13 @@ partners = [
     # "stardist",         # for model testing and stardist postprocessing  # TODO: add updated stardist to partners env
 ]
 stardist = ["stardist"] # for stardist postprocessing
+gradio-server = ['gradio[mcp];python_version>="3.10"']
+gradio-client = ['gradio_client;python_version>="3.10"']
 dev = [
-    "cellpose",                              # for model testing
+    "cellpose<4.2",                          # for model testing
     "crick",
+    'gradio[mcp];python_version>="3.10"',
+    'gradio_client;python_version>="3.10"',
     "httpx",
     "jax",
     "jupyter",
@@ -68,7 +72,7 @@ dev = [
     'onnx_ir!=0.1.14;python_version<"3.10"', # uses typing.Concatentate which requires py>=3.10
     "packaging>=17.0",
     "pre-commit",
-    "pyright==1.1.408",
+    "pyright==1.1.410",
     "pytest-cov",
     "pytest",
     "python-dotenv",

src/bioimageio/core/__init__.py

@@ -16,12 +16,27 @@
 """
 # ruff: noqa: E402
 
-__version__ = "0.10.4"
+__version__ = "0.11.0"
 from loguru import logger
 
 logger.disable("bioimageio.core")
 
-import bioimageio.spec
+
+from bioimageio.spec import ValidationSummary as ValidationSummary
+from bioimageio.spec import build_description as build_description
+from bioimageio.spec import dump_description as dump_description
+from bioimageio.spec import load_dataset_description as load_dataset_description
+from bioimageio.spec import load_description as load_description
+from bioimageio.spec import (
+    load_description_and_validate_format_only as load_description_and_validate_format_only,
+)
+from bioimageio.spec import load_model_description as load_model_description
+from bioimageio.spec import save_bioimageio_package as save_bioimageio_package
+from bioimageio.spec import (
+    save_bioimageio_package_as_folder as save_bioimageio_package_as_folder,
+)
+from bioimageio.spec import save_bioimageio_yaml_only as save_bioimageio_yaml_only
+from bioimageio.spec import validate_format as validate_format
 
 from . import axis as axis
 from . import backends as backends
@@ -31,7 +46,6 @@
 from . import common as common
 from . import digest_spec as digest_spec
 from . import io as io
-from . import model_adapters as model_adapters
 from . import prediction as prediction
 from . import proc_ops as proc_ops
 from . import proc_setup as proc_setup
@@ -40,46 +54,40 @@
 from . import stat_measures as stat_measures
 from . import tensor as tensor
 from . import weight_converters as weight_converters
+from ._prediction_pipeline import IntermediatePrediction as IntermediatePrediction
 from ._prediction_pipeline import PredictionPipeline as PredictionPipeline
+from ._prediction_pipeline import RemotePredictionPipeline as RemotePredictionPipeline
 from ._prediction_pipeline import (
     create_prediction_pipeline as create_prediction_pipeline,
 )
+from ._prediction_pipeline import (
+    create_remote_prediction_pipeline as create_remote_prediction_pipeline,
+)
 from ._resource_tests import enable_determinism as enable_determinism
 from ._resource_tests import load_description_and_test as load_description_and_test
 from ._resource_tests import test_description as test_description
 from ._resource_tests import test_model as test_model
+from ._sample_serializer import SampleSerializer as SampleSerializer
 from ._settings import Settings as Settings
 from ._settings import settings as settings
 
-# reexports from bioimageio.spec
-build_description = bioimageio.spec.build_description
-dump_description = bioimageio.spec.dump_description
-load_dataset_description = bioimageio.spec.load_dataset_description
-load_description = bioimageio.spec.load_description
-load_description_and_validate_format_only = (
-    bioimageio.spec.load_description_and_validate_format_only
-)
-load_model_description = bioimageio.spec.load_model_description
-save_bioimageio_package = bioimageio.spec.save_bioimageio_package
-save_bioimageio_package_as_folder = bioimageio.spec.save_bioimageio_package_as_folder
-save_bioimageio_yaml_only = bioimageio.spec.save_bioimageio_yaml_only
-validate_format = bioimageio.spec.validate_format
-ValidationSummary = bioimageio.spec.ValidationSummary
-
-
 # reexports from bioimageio.core submodules
-add_weights = weight_converters.add_weights
-Axis = axis.Axis
-AxisId = axis.AxisId
-BlockMeta = block_meta.BlockMeta
-compute_dataset_measures = stat_calculators.compute_dataset_measures
-create_model_adapter = backends.create_model_adapter
-MemberId = common.MemberId
-predict = prediction.predict
-predict_many = prediction.predict_many
-Sample = sample.Sample
-Stat = stat_measures.Stat
-Tensor = tensor.Tensor
+from .axis import Axis as Axis
+from .axis import AxisId as AxisId
+from .backends import create_model_adapter as create_model_adapter
+from .block_meta import BlockMeta as BlockMeta
+from .common import MemberId as MemberId
+from .prediction import predict as predict
+from .prediction import predict_many as predict_many
+from .sample import Sample as Sample
+from .sample import SampleBlock as SampleBlock
+from .sample import SampleBlockMeta as SampleBlockMeta
+from .stat_calculators import compute_dataset_measures as compute_dataset_measures
+from .stat_calculators import compute_measures as compute_measures
+from .stat_calculators import compute_sample_measures as compute_sample_measures
+from .stat_measures import Stat as Stat
+from .tensor import Tensor as Tensor
+from .weight_converters import add_weights as add_weights
 
 # aliases
 test_resource = test_description

src/bioimageio/core/__main__.py

@@ -14,7 +14,7 @@
     + "<cyan>{module}</cyan> - <level>{message}</level>",
 )
 
-from .cli import Bioimageio
+from .cli import Bioimageio  # noqa: E402
 
 
 def main():

src/bioimageio/core/_axis_annotations.py

@@ -0,0 +1,9 @@
+from typing import Annotated, TypeVar
+
+from ._common_annotations import PydanticMappingProxyAnnotation
+from .axis import PerAxis
+
+_T = TypeVar("_T")
+
+PerAxisAnno = Annotated[PerAxis[_T], PydanticMappingProxyAnnotation]
+"""PerAxis annotated with `PydanticMappingProxyAnnotation` to be compatible with pydantic models."""

src/bioimageio/core/_common_annotations.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from types import MappingProxyType
+from typing import (
+    Annotated,
+    Any,
+    Hashable,
+    Mapping,
+    TypeVar,
+)
+
+import pydantic
+from pydantic_core.core_schema import (
+    CoreSchema,
+    chain_schema,
+    is_instance_schema,
+    json_or_python_schema,
+    no_info_plain_validator_function,
+    plain_serializer_function_ser_schema,
+)
+from typing_extensions import get_args
+
+from .common import PerMember
+
+_K = TypeVar("_K", bound=Hashable)
+_V = TypeVar("_V")
+
+
+def _validate_from_mapping(d: Mapping[_K, _V]) -> MappingProxyType[_K, _V]:
+    return MappingProxyType(dict(d))
+
+
+class PydanticMappingProxyAnnotation:
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
+    ) -> CoreSchema:
+
+        k_type, v_type = get_args(source_type)
+        mapping_proxy_schema = chain_schema(
+            [
+                handler.generate_schema(dict[k_type, v_type]),
+                no_info_plain_validator_function(_validate_from_mapping),
+                is_instance_schema(MappingProxyType),
+            ]
+        )
+        return json_or_python_schema(
+            json_schema=mapping_proxy_schema,
+            python_schema=mapping_proxy_schema,
+            serialization=plain_serializer_function_ser_schema(dict),
+        )
+
+
+_T = TypeVar("_T")
+
+PerMemberAnno = Annotated[PerMember[_T], PydanticMappingProxyAnnotation]

src/bioimageio/core/_description_serializer.py

@@ -0,0 +1,69 @@
+import base64
+import hashlib
+from io import BytesIO
+from typing import Tuple
+from zipfile import ZipFile
+
+from bioimageio.spec import (
+    InvalidDescr,
+    ResourceDescr,
+    load_description,
+    save_bioimageio_package_to_stream,
+)
+from bioimageio.spec.common import Sha256
+
+
+class DescriptionSerializer:
+    """Description serializer intended for client/server communication, NOT for sharing resource descriptions.
+
+    This serializer only includes local files to keep the serialized package small.
+    """
+
+    STRING_ENCODING = "ascii"
+
+    @staticmethod
+    def serialize(rd: ResourceDescr) -> bytes:
+        stream = save_bioimageio_package_to_stream(rd, local_files_only=True)
+        _ = stream.seek(0)
+        return stream.read()
+
+    @classmethod
+    def serialize_to_string(cls, rd: ResourceDescr) -> str:
+        package_bytes = cls.serialize(rd)
+
+        safe_bytes = cls._get_safe_bytes(package_bytes)
+        serialized_str = safe_bytes.decode(cls.STRING_ENCODING)
+        if len(serialized_str) <= 2083:
+            raise RuntimeError(
+                "Serialized model description should be longer than 2083 characters to not be treated as a URL on the server side."
+            )
+        return serialized_str
+
+    @staticmethod
+    def _get_safe_bytes(raw_bytes: bytes) -> bytes:
+        return base64.b64encode(raw_bytes)
+
+    @classmethod
+    def deserialize_from_string(cls, serialized: str) -> ResourceDescr:
+        package_bytes = base64.b64decode(serialized.encode(cls.STRING_ENCODING))
+        return cls.deserialize(package_bytes)
+
+    @staticmethod
+    def deserialize(serialized: bytes) -> ResourceDescr:
+        descr = load_description(ZipFile(BytesIO(serialized)), perform_io_checks=False)
+        if isinstance(descr, InvalidDescr):
+            raise ValueError(f"invalid serialized model package: {descr.get_reason()}")
+
+        return descr
+
+    @classmethod
+    def serialize_to_string_and_hash(cls, rd: ResourceDescr) -> Tuple[str, Sha256]:
+        package_bytes = cls.serialize(rd)
+        safe_bytes = cls._get_safe_bytes(package_bytes)
+        serialized_str = safe_bytes.decode(cls.STRING_ENCODING)
+        if len(serialized_str) <= 2083:
+            raise RuntimeError(
+                "Serialized model description should be longer than 2083 characters to not be treated as a URL on the server side."
+            )
+        sha256 = Sha256(hashlib.sha256(package_bytes).hexdigest())
+        return serialized_str, sha256