Merge pull request #201 from bioimage-io/more-build-spec-tests2

Extend tests for build_spec functionality

Author: constantinpape

bioimageio/core/build_spec/add_weights.py

@@ -63,6 +63,13 @@ def add_weights(
     except Exception as e:
         raise e
     finally:
+        # clean up tmp files
+        os.remove(weight_out)
         if tmp_arch is not None:
             os.remove(tmp_arch)
+        # for some reason the weights are also copied to the cwd.
+        # not sure why this happens, but it needs to be cleaned up, unless these are the input weigths
+        weights_cwd = Path(os.path.split(weight_uri)[1])
+        if weights_cwd.exists() and weights_cwd.absolute() != Path(weight_uri).absolute():
+            os.remove(weights_cwd)
     return model

tests/build_spec/test_add_weights.py

@@ -1,4 +1,6 @@
+import os
 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
 
 
 def _test_add_weights(model, tmp_path, base_weights, added_weights, **kwargs):
@@ -28,6 +30,13 @@ def _test_add_weights(model, tmp_path, base_weights, added_weights, **kwargs):
     for weight in new_rdf.weights.values():
         assert weight.source.exists()
 
+    test_res = _test_model(out_path, added_weights)
+    test_res = _test_model(out_path)
+    assert test_res["error"] is None
+
+    # make sure the weights were cleaned from the cwd
+    assert not os.path.exists(os.path.split(weight_path)[1])
+
 
 def test_add_torchscript(unet2d_nuclei_broad_model, tmp_path):
     _test_add_weights(unet2d_nuclei_broad_model, tmp_path, "pytorch_state_dict", "torchscript")

tests/build_spec/test_build_spec.py

@@ -3,6 +3,7 @@
 from bioimageio.core import load_raw_resource_description, load_resource_description
 from bioimageio.core.resource_io import nodes
 from bioimageio.core.resource_io.utils import resolve_source
+from bioimageio.core.resource_tests import test_model as _test_model
 from marshmallow import missing
 
 
@@ -81,17 +82,19 @@ def _test_build_spec(
         output_path=out_path,
         add_deepimagej_config=add_deepimagej_config,
         maintainers=[{"github_user": "jane_doe"}],
+        input_names=[inp.name for inp in model_spec.inputs],
+        output_names=[out.name for out in model_spec.outputs],
     )
     if architecture is not None:
         kwargs["architecture"] = architecture
     if model_kwargs is not None:
-        kwargs["kwargs"] = model_kwargs
+        kwargs["model_kwargs"] = model_kwargs
     if tensorflow_version is not None:
         kwargs["tensorflow_version"] = tensorflow_version
     if opset_version is not None:
         kwargs["opset_version"] = opset_version
     if use_implicit_output_shape:
-        kwargs["input_name"] = ["input"]
+        kwargs["input_names"] = ["input"]
         kwargs["output_reference"] = ["input"]
         kwargs["output_scale"] = [[1.0, 1.0, 1.0, 1.0]]
         kwargs["output_offset"] = [[0.0, 0.0, 0.0, 0.0]]
@@ -134,6 +137,10 @@ def _test_build_spec(
             assert attachments.files is not missing
             assert n_processing == len(attachments.files)
 
+    # test inference for the model to ensure that the weights were written correctly
+    test_res = _test_model(out_path)
+    assert test_res["error"] is None
+
 
 def test_build_spec_pytorch(any_torch_model, tmp_path):
     _test_build_spec(any_torch_model, tmp_path / "model.zip", "pytorch_state_dict")

tests/weight_converter/keras/test_tensorflow.py

@@ -1,3 +1,6 @@
+import zipfile
+
+
 def test_tensorflow_converter(any_keras_model, tmp_path):
     from bioimageio.core.weight_converter.keras import convert_weights_to_tensorflow_saved_model_bundle
 
@@ -16,3 +19,9 @@ def test_tensorflow_converter_zipped(any_keras_model, tmp_path):
     ret_val = convert_weights_to_tensorflow_saved_model_bundle(any_keras_model, out_path)
     assert out_path.exists()
     assert ret_val == 0  # check for correctness is done in converter and returns 0 if it passes
+
+    # make sure that the zip package was created correctly
+    expected_names = {"saved_model.pb", "variables/variables.index"}
+    with zipfile.ZipFile(out_path, "r") as f:
+        names = set([name for name in f.namelist()])
+    assert len(expected_names - names) == 0