Update add_weights functionality and add tests

Author: constantinpape

bioimageio/core/build_spec/__init__.py

@@ -1 +1,2 @@
-from .build_model import add_weights, build_model
+from .add_weights import add_weights
+from .build_model import build_model

bioimageio/core/build_spec/add_weights.py

@@ -0,0 +1,68 @@
+import os
+from pathlib import Path
+from shutil import copyfile
+from typing import Dict, Optional, Union
+
+from bioimageio.core import export_resource_package, load_raw_resource_description
+from bioimageio.spec.shared.raw_nodes import ResourceDescription as RawResourceDescription
+from .build_model import _get_weights
+
+
+def add_weights(
+    model: Union[RawResourceDescription, os.PathLike, str],
+    weight_uri: Union[str, Path],
+    output_path: Union[str, Path],
+    *,
+    weight_type: Optional[str] = None,
+    architecture: Optional[str] = None,
+    model_kwargs: Optional[Dict[str, Union[int, float, str]]] = None,
+    tensorflow_version: Optional[str] = None,
+    opset_version: Optional[str] = None,
+    **weight_kwargs,
+):
+    """Add weight entry to bioimage.io model.
+
+    Args:
+        model: the resource description of the model to which the weight format is added
+        weight_uri: the weight file to be added
+        output_path: where to serialize the new model with additional weight format
+        weight_type: the format of the weights to be added
+        architecture: the file with the source code for the model architecture and the corresponding class.
+            Only required for models with pytorch_state_dict weight format.
+        model_kwargs: the keyword arguments for the model class.
+            Only required for models with pytorch_state_dict weight format.
+        tensorflow_version: the tensorflow version used for training the model.
+            Only requred for models with tensorflow or keras weight format.
+        opset_version: the opset version used in this model.
+            Only requred for models with onnx weight format.
+        weight_kwargs: additional keyword arguments for the weight.
+    """
+    model = load_raw_resource_description(model)
+
+    # copy the weight path to the input model's root, otherwise it will
+    # not be found when packaging the new model
+    weight_out = os.path.join(model.root_path, Path(weight_uri).name)
+    if Path(weight_out) != Path(weight_uri):
+        copyfile(weight_uri, weight_out)
+
+    new_weights, tmp_arch = _get_weights(
+        weight_out,
+        weight_type,
+        root=Path("."),
+        architecture=architecture,
+        model_kwargs=model_kwargs,
+        tensorflow_version=tensorflow_version,
+        opset_version=opset_version,
+        **weight_kwargs,
+    )
+    model.weights.update(new_weights)
+
+    try:
+        model_package = export_resource_package(model, output_path=output_path)
+        model = load_raw_resource_description(model_package)
+    except Exception as e:
+        raise e
+    finally:
+        if tmp_arch is not None:
+            os.remove(tmp_arch)
+    return model

bioimageio/core/build_spec/build_model.py

@@ -616,8 +616,10 @@ def build_model(
         documentation: relative file path to markdown documentation for this model.
         cite: citations for this model.
         output_path: where to save the zipped model package.
-        source: the file with the source code for the model architecture and the corresponding class.
+        architecture: the file with the source code for the model architecture and the corresponding class.
+            Only required for models with pytorch_state_dict weight format.
         model_kwargs: the keyword arguments for the model class.
+            Only required for models with pytorch_state_dict weight format.
         weight_type: the type of the weights.
         sample_inputs: list of sample inputs to demonstrate the model performance.
         sample_outputs: list of sample outputs corresponding to sample_inputs.
@@ -648,8 +650,9 @@ def build_model(
         root: optional root path for relative paths. This can be helpful when building a spec from another model spec.
         add_deepimagej_config: add the deepimagej config to the model.
         tensorflow_version: the tensorflow version used for training the model.
-            Needs to be passed for tensorflow or keras models.
-        opset_version: the opset version used in this model. Needs to be passed for onnx models.
+            Only requred for models with tensorflow or keras weight format.
+        opset_version: the opset version used in this model.
+            Only requred for models with onnx weight format.
         weight_kwargs: additional keyword arguments for this weight type.
     """
     if root is None:
@@ -848,35 +851,3 @@ def build_model(
 
     model = load_raw_resource_description(model_package)
     return model
-
-
-def add_weights(
-    model,
-    weight_uri: Union[str, Path],
-    weight_type: Optional[str] = None,
-    output_path: Optional[Union[str, Path]] = None,
-    architecture: Optional[str] = None,
-    model_kwargs: Optional[Dict[str, Union[int, float, str]]] = None,
-    tensorflow_version: Optional[str] = None,
-    opset_version: Optional[str] = None,
-    **weight_kwargs,
-):
-    """Add weight entry to bioimage.io model."""
-    # we need to pass the weight path as abs path to avoid confusion with different root directories
-    new_weights, tmp_arch = _get_weights(
-        Path(weight_uri).absolute(),
-        weight_type,
-        root=Path("."),
-        architecture=architecture,
-        model_kwargs=model_kwargs,
-        tensorflow_version=tensorflow_version,
-        opset_version=opset_version,
-        **weight_kwargs,
-    )
-    model.weights.update(new_weights)
-    if output_path is not None:
-        model_package = export_resource_package(model, output_path=output_path)
-        model = load_raw_resource_description(model_package)
-    if tmp_arch is not None:
-        os.remove(tmp_arch)
-    return model

tests/build_spec/test_add_weights.py

@@ -0,0 +1,37 @@
+from bioimageio.core import export_resource_package, load_raw_resource_description, load_resource_description
+
+
+def _test_add_weights(model, tmp_path, base_weights, added_weights, **kwargs):
+    from bioimageio.core.build_spec import add_weights
+
+    rdf = load_raw_resource_description(model)
+    assert base_weights in rdf.weights
+    assert added_weights in rdf.weights
+
+    weight_path = load_resource_description(model).weights[added_weights].source
+    assert weight_path.exists()
+
+    drop_weights = set(rdf.weights.keys()) - {base_weights}
+    for drop in drop_weights:
+        rdf.weights.pop(drop)
+    assert tuple(rdf.weights.keys()) == (base_weights,)
+
+    in_path = tmp_path / "model1.zip"
+    export_resource_package(rdf, output_path=in_path)
+
+    out_path = tmp_path / "model2.zip"
+    add_weights(in_path, weight_path, weight_type=added_weights, output_path=out_path, **kwargs)
+
+    assert out_path.exists()
+    new_rdf = load_resource_description(out_path)
+    assert set(new_rdf.weights.keys()) == {base_weights, added_weights}
+    for weight in new_rdf.weights.values():
+        assert weight.source.exists()
+
+
+def test_add_torchscript(unet2d_nuclei_broad_model, tmp_path):
+    _test_add_weights(unet2d_nuclei_broad_model, tmp_path, "pytorch_state_dict", "torchscript")
+
+
+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)