more get_reader updates

Author: FynnBe

bioimageio/core/cli.py

@@ -57,13 +57,12 @@
     update_hashes,
 )
 from bioimageio.spec._internal.io import is_yaml_value
-from bioimageio.spec._internal.io_basics import ZipPath
 from bioimageio.spec._internal.io_utils import open_bioimageio_yaml
 from bioimageio.spec._internal.types import NotEmpty
 from bioimageio.spec.dataset import DatasetDescr
 from bioimageio.spec.model import ModelDescr, v0_4, v0_5
 from bioimageio.spec.notebook import NotebookDescr
-from bioimageio.spec.utils import download, ensure_description_is_model, write_yaml
+from bioimageio.spec.utils import ensure_description_is_model, get_reader, write_yaml
 
 from .commands import WeightFormatArgAll, WeightFormatArgAny, package, test
 from .common import MemberId, SampleId, SupportedWeightsFormat
@@ -487,16 +486,12 @@ def _example(self):
         example_path.mkdir(exist_ok=True)
 
         for t, src in zip(input_ids, example_inputs):
-            local = download(src).path
-            dst = Path(f"{example_path}/{t}/001{''.join(local.suffixes)}")
+            reader = get_reader(src)
+            dst = Path(f"{example_path}/{t}/001{reader.suffix}")
             dst.parent.mkdir(parents=True, exist_ok=True)
             inputs001.append(dst.as_posix())
-            if isinstance(local, Path):
-                shutil.copy(local, dst)
-            elif isinstance(local, ZipPath):
-                _ = local.root.extract(local.at, path=dst)
-            else:
-                assert_never(local)
+            with dst.open("wb") as f:
+                shutil.copyfileobj(reader, f)
 
         inputs = [tuple(inputs001)]
         output_pattern = f"{example_path}/outputs/{{output_id}}/{{sample_id}}.tif"

bioimageio/core/digest_spec.py

@@ -104,13 +104,10 @@ def _import_from_file_impl(
         assert module_name.isidentifier(), module_name
 
     source_bytes = reader.read()
-    # with NamedTemporaryFile(
 
     module = sys.modules.get(module_name)
     if module is None:
         try:
-            #     local_source.write(source_bytes)
-            #     local_source.flush()
             if reader.original_file_name.endswith(".zip") or is_zipfile(reader):
                 module_path = TemporaryDirectory(
                     prefix=module_name,

bioimageio/core/io.py

@@ -23,6 +23,7 @@
 from bioimageio.spec._internal.io import get_reader, interprete_file_source
 from bioimageio.spec._internal.type_guards import is_ndarray
 from bioimageio.spec.common import (
+    BytesReader,
     FileSource,
     HttpUrl,
     PermissiveFileSource,
@@ -279,10 +280,16 @@ def load_dataset_stat(path: Path):
     return {e.measure: e.value for e in seq}
 
 
-def ensure_unzipped(source: Union[PermissiveFileSource, ZipPath], folder: Path):
-    """unzip a (downloaded) **source** to a file in **folder** if source is a zip archive.
-    Always returns the path to the unzipped source (maybe source itself)"""
-    weights_reader = get_reader(source)
+def ensure_unzipped(
+    source: Union[PermissiveFileSource, ZipPath, BytesReader], folder: Path
+):
+    """unzip a (downloaded) **source** to a file in **folder** if source is a zip archive
+    otherwise copy **source** to a file in **folder**."""
+    if isinstance(source, BytesReader):
+        weights_reader = source
+    else:
+        weights_reader = get_reader(source)
+
     out_path = folder / (
         weights_reader.original_file_name or f"file{weights_reader.suffix}"
     )

bioimageio/core/weight_converters/keras_to_tensorflow.py

@@ -1,12 +1,12 @@
 import os
 import shutil
 from pathlib import Path
+from tempfile import TemporaryDirectory
 from typing import Union, no_type_check
 from zipfile import ZipFile
 
 import tensorflow  # pyright: ignore[reportMissingTypeStubs]
 
-from bioimageio.spec._internal.io import download
 from bioimageio.spec._internal.version_type import Version
 from bioimageio.spec.common import ZipPath
 from bioimageio.spec.model.v0_5 import (
@@ -70,7 +70,7 @@ def convert(
         raise ValueError("Missing Keras Hdf5 weights to convert from.")
 
     weight_spec = model_descr.weights.keras_hdf5
-    weight_path = download(weight_spec.source).path
+    weight_reader = weight_spec.get_reader()
 
     if weight_spec.tensorflow_version:
         model_tf_major_ver = int(weight_spec.tensorflow_version.major)
@@ -79,30 +79,34 @@ def convert(
                 f"Tensorflow major versions of model {model_tf_major_ver} is not {tf_major_ver}"
             )
 
-    if tf_major_ver == 1:
-        if len(model_descr.inputs) != 1 or len(model_descr.outputs) != 1:
-            raise NotImplementedError(
-                "Weight conversion for models with multiple inputs or outputs is not yet implemented."
-            )
-
-        input_name = str(
-            d.id
-            if isinstance((d := model_descr.inputs[0]), InputTensorDescr)
-            else d.name
-        )
-        output_name = str(
-            d.id
-            if isinstance((d := model_descr.outputs[0]), OutputTensorDescr)
-            else d.name
-        )
-        return _convert_tf1(
-            ensure_unzipped(weight_path, Path("bioimageio_unzipped_tf_weights")),
-            output_path,
-            input_name,
-            output_name,
+    with TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir:
+        local_weights = ensure_unzipped(
+            weight_reader, Path(temp_dir) / "bioimageio_unzipped_tf_weights"
         )
-    else:
-        return _convert_tf2(weight_path, output_path)
+        if tf_major_ver == 1:
+            if len(model_descr.inputs) != 1 or len(model_descr.outputs) != 1:
+                raise NotImplementedError(
+                    "Weight conversion for models with multiple inputs or outputs is not yet implemented."
+                )
+
+            input_name = str(
+                d.id
+                if isinstance((d := model_descr.inputs[0]), InputTensorDescr)
+                else d.name
+            )
+            output_name = str(
+                d.id
+                if isinstance((d := model_descr.outputs[0]), OutputTensorDescr)
+                else d.name
+            )
+            return _convert_tf1(
+                ensure_unzipped(local_weights, Path("bioimageio_unzipped_tf_weights")),
+                output_path,
+                input_name,
+                output_name,
+            )
+        else:
+            return _convert_tf2(local_weights, output_path)
 
 
 def _convert_tf2(

bioimageio/core/weight_converters/torchscript_to_onnx.py

@@ -3,7 +3,6 @@
 import torch.jit
 
 from bioimageio.spec.model.v0_5 import ModelDescr, OnnxWeightsDescr
-from bioimageio.spec.utils import download
 
 from .. import __version__
 from ..digest_spec import get_member_id, get_test_inputs
@@ -55,8 +54,8 @@ def convert(
     ]
     inputs_torch = [torch.from_numpy(ipt) for ipt in inputs_numpy]
 
-    weight_path = download(torchscript_descr).path
-    model = torch.jit.load(weight_path)  # type: ignore
+    weight_reader = torchscript_descr.get_reader()
+    model = torch.jit.load(weight_reader)  # type: ignore
     model.to("cpu")
     model = model.eval()  # type: ignore
 

tests/test_bioimageio_collection.py

@@ -18,7 +18,7 @@ def _get_latest_rdf_sources():
     for entry in entries:
         version = entry["versions"][0]
         ret[f"{entry['concept']}/{version['v']}"] = (
-            HttpUrl(version["source"]),  # pyright: ignore[reportCallIssue]
+            HttpUrl(version["source"]),
             Sha256(version["sha256"]),
         )