Update build_model for changed output shape spec; add test

Author: constantinpape

bioimageio/core/build_spec/build_model.py

@@ -179,16 +179,16 @@ def _get_input_tensor(test_in, name, step, min_shape, data_range, axes, preproce
     return inputs
 
 
-def _get_output_tensor(test_out, name, reference_input, scale, offset, axes, data_range, postprocessing, halo):
+def _get_output_tensor(test_out, name, reference_tensor, scale, offset, axes, data_range, postprocessing, halo):
     shape = test_out.shape
-    if reference_input is None:
+    if reference_tensor is None:
         assert scale is None
         assert offset is None
         shape_description = shape
     else:
         assert scale is not None
         assert offset is not None
-        shape_description = {"reference_input": reference_input, "scale": scale, "offset": offset}
+        shape_description = {"reference_tensor": reference_tensor, "scale": scale, "offset": offset}
 
     axes = _get_axes(axes, test_out.ndim)
     data_range = _get_data_range(data_range, test_out.dtype)

tests/build_spec/test_build_spec.py

@@ -6,7 +6,7 @@
 from bioimageio.core.resource_io.io_ import load_raw_resource_description
 
 
-def _test_build_spec(path, weight_type, tensorflow_version=None):
+def _test_build_spec(path, weight_type, tensorflow_version=None, use_implicit_output_shape=False):
     from bioimageio.core.build_spec import build_model
 
     model_spec = load_raw_resource_description(path)
@@ -50,6 +50,11 @@ def _test_build_spec(path, weight_type, tensorflow_version=None):
     )
     if tensorflow_version is not None:
         kwargs["tensorflow_version"] = tensorflow_version
+    if use_implicit_output_shape:
+        kwargs["input_name"] = "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]
     raw_model = build_model(**kwargs)
     spec.model.schema.Model().dump(raw_model)
 
@@ -58,6 +63,10 @@ def test_build_spec_pytorch(any_torch_model):
     _test_build_spec(any_torch_model, "pytorch_state_dict")
 
 
+def test_build_spec_implicit_output_shape(unet2d_nuclei_broad_model):
+    _test_build_spec(unet2d_nuclei_broad_model, "pytorch_state_dict", use_implicit_output_shape=True)
+
+
 def test_build_spec_torchscript(any_torchscript_model):
     _test_build_spec(any_torchscript_model, "pytorch_script")