Merge pull request #163 from bioimage-io/fix-shape-val

Fix shape val

Author: FynnBe

bioimageio/core/__main__.py

@@ -1,6 +1,7 @@
 import enum
 import json
 import os
+import sys
 from glob import glob
 
 from pathlib import Path
@@ -53,9 +54,10 @@ def package(
     # typer bug: typer returns empty tuple instead of None if weights_order_priority is not given
     weights_priority_order = weights_priority_order or None
 
-    return commands.package(
+    ret_code = commands.package(
         rdf_source=rdf_source, path=path, weights_priority_order=weights_priority_order, verbose=verbose
     )
+    sys.exit(ret_code)
 
 
 package.__doc__ = commands.package.__doc__
@@ -89,11 +91,12 @@ def test_model(
     )
     if summary["error"] is None:
         print(f"Model test for {model_rdf} has passed.")
-        return 0
+        ret_code = 0
     else:
         print(f"Model test for {model_rdf} has FAILED!")
         print(summary)
-        return 1
+        ret_code = 1
+    sys.exit(ret_code)
 
 
 test_model.__doc__ = resource_tests.test_model.__doc__
@@ -116,11 +119,12 @@ def test_resource(
     )
     if summary["error"] is None:
         print(f"Resource test for {rdf} has passed.")
-        return 0
+        ret_code = 0
     else:
         print(f"Resource test for {rdf} has FAILED!")
         print(summary)
-        return 1
+        ret_code = 1
+    sys.exit(ret_code)
 
 
 test_resource.__doc__ = resource_tests.test_resource.__doc__
@@ -159,7 +163,6 @@ def predict_image(
     prediction.predict_image(
         model_rdf, inputs, outputs, padding, tiling, None if weight_format is None else weight_format.value, devices
     )
-    return 0
 
 
 predict_image.__doc__ = prediction.predict_image.__doc__
@@ -211,7 +214,6 @@ def predict_images(
         devices=devices,
         verbose=True,
     )
-    return 0
 
 
 predict_images.__doc__ = prediction.predict_images.__doc__
@@ -241,7 +243,8 @@ def convert_torch_weights_to_torchscript(
         output_path: Path = typer.Argument(..., help="Where to save the torchscript weights."),
         use_tracing: bool = typer.Option(True, help="Whether to use torch.jit tracing or scripting."),
     ) -> int:
-        return torch_converter.convert_weights_to_pytorch_script(model_rdf, output_path, use_tracing)
+        ret_code = torch_converter.convert_weights_to_pytorch_script(model_rdf, output_path, use_tracing)
+        sys.exit(ret_code)
 
     convert_torch_weights_to_torchscript.__doc__ = torch_converter.convert_weights_to_pytorch_script.__doc__
 

bioimageio/core/commands.py

@@ -32,6 +32,13 @@ def package(
 
     try:
         rdf_local_source = resolve_uri(rdf_source)
+    except Exception as e:
+        print(f"Failed to resolve RDF source {rdf_source}: {e}")
+        if verbose:
+            traceback.print_exc()
+        return 1
+
+    try:
         path = path.with_name(path.name.format(src_name=rdf_local_source.stem))
         shutil.move(tmp_package_path, path)
     except Exception as e:

bioimageio/core/resource_tests.py

@@ -38,6 +38,41 @@ def test_model(
         return {"error": f"Expected RDF type Model, got {type(model)} instead.", "traceback": None}
 
 
+def _validate_input_shape(shape: Tuple[int, ...], shape_spec) -> bool:
+    if isinstance(shape_spec, list):
+        if shape != tuple(shape_spec):
+            return False
+    elif isinstance(shape_spec, ParametrizedInputShape):
+        assert len(shape_spec.min) == len(shape_spec.step)
+        if len(shape) != len(shape_spec.min):
+            return False
+        min_shape = shape_spec.min
+        step = shape_spec.step
+        # check if the shape is valid for all dimension by seeing if it can be reached with an integer number of steps
+        # NOTE we allow that the valid shape is reached using a different number of steps for each axis here
+        # this is usually valid because dimensions are independent in neural networks
+        is_valid = [(sh - minsh) % st == 0 if st > 0 else sh == minsh for sh, st, minsh in zip(shape, step, min_shape)]
+        return all(is_valid)
+    else:
+        raise TypeError(f"Encountered unexpected shape description of type {type(shape_spec)}")
+
+    return True
+
+
+def _validate_output_shape(shape: Tuple[int, ...], shape_spec, input_shapes) -> bool:
+    if isinstance(shape_spec, list):
+        return shape == tuple(shape_spec)
+    elif isinstance(shape_spec, ImplicitOutputShape):
+        ipt_shape = numpy.array(input_shapes[shape_spec.reference_tensor])
+        scale = numpy.array(shape_spec.scale)
+        offset = numpy.array(shape_spec.offset)
+        exp_shape = numpy.round_(ipt_shape * scale) + 2 * offset
+
+        return shape == tuple(exp_shape)
+    else:
+        raise TypeError(f"Encountered unexpected shape description of type {type(shape_spec)}")
+
+
 def test_resource(
     model_rdf: Union[RawResourceDescription, ResourceDescription, URI, Path, str],
     *,
@@ -64,65 +99,22 @@ def test_resource(
                 inputs = [np.load(str(in_path)) for in_path in model.test_inputs]
                 expected = [np.load(str(out_path)) for out_path in model.test_outputs]
 
-                # check if test data shapes match their description
-                input_shapes = [ipt.shape for ipt in inputs]
-                output_shapes = [out.shape for out in expected]
-
-                def input_shape_is_valid(shape: Tuple[int, ...], shape_spec) -> bool:
-                    if isinstance(shape_spec, list):
-                        if shape != tuple(shape_spec):
-                            return False
-                    elif isinstance(shape_spec, ParametrizedInputShape):
-                        assert len(shape_spec.min) == len(shape_spec.step)
-                        if len(shape) != len(shape_spec.min):
-                            return False
-
-                        valid_shape = numpy.array(shape_spec.min)
-                        step = numpy.array(shape_spec.step)
-                        if (step == 0).all():
-                            return shape == tuple(valid_shape)
-
-                        shape = numpy.array(shape)
-                        while (shape <= valid_shape).all():
-                            if (shape == valid_shape).all():
-                                break
-
-                            shape += step
-                        else:
-                            return False
-
-                    else:
-                        raise TypeError(f"Encountered unexpected shape description of type {type(shape_spec)}")
-
-                    return True
-
                 assert len(inputs) == len(model.inputs)  # should be checked by validation
                 input_shapes = {}
                 for idx, (ipt, ipt_spec) in enumerate(zip(inputs, model.inputs)):
-                    if not input_shape_is_valid(ipt, ipt_spec.shape):
+                    if not _validate_input_shape(tuple(ipt.shape), ipt_spec.shape):
                         raise ValidationError(
-                            f"Shape of test input {idx} '{ipt_spec.name}' does not match "
-                            f"input shape description: {ipt_spec.shape}"
+                            f"Shape {tuple(ipt.shape)} of test input {idx} '{ipt_spec.name}' does not match "
+                            f"input shape description: {ipt_spec.shape}."
                         )
                     input_shapes[ipt_spec.name] = ipt.shape
 
-                def output_shape_is_valid(shape: Tuple[int, ...], shape_spec) -> bool:
-                    if isinstance(shape_spec, list):
-                        return shape == tuple(shape_spec)
-                    elif isinstance(shape_spec, ImplicitOutputShape):
-                        ipt_shape = numpy.array(input_shapes[shape_spec.reference_tensor])
-                        scale = numpy.array(shape_spec.scale)
-                        offset = numpy.array(shape_spec.offset)
-                        exp_shape = numpy.round_(ipt_shape * scale) + 2 * offset
-
-                        return shape == tuple(exp_shape)
-
                 assert len(expected) == len(model.outputs)  # should be checked by validation
                 for idx, (out, out_spec) in enumerate(zip(expected, model.outputs)):
-                    if not output_shape_is_valid(out, out_spec.shape):
+                    if not _validate_output_shape(tuple(out.shape), out_spec.shape, input_shapes):
                         error = (error or "") + (
-                            f"Shape of test output {idx} '{out_spec.name}' does not match "
-                            f"output shape description: {out_spec.shape}.\n"
+                            f"Shape {tuple(out.shape)} of test output {idx} '{out_spec.name}' does not match "
+                            f"output shape description: {out_spec.shape}."
                         )
 
                 with create_prediction_pipeline(

tests/conftest.py

@@ -42,6 +42,10 @@
         "https://raw.githubusercontent.com/bioimage-io/spec-bioimage-io/main/example_specs/models/"
         "stardist_example_model/rdf_wrong_shape.yaml"
     ),
+    "stardist_wrong_shape2": (
+        "https://raw.githubusercontent.com/bioimage-io/spec-bioimage-io/main/example_specs/models/"
+        "stardist_example_model/rdf_wrong_shape2.yaml"
+    ),
 }
 
 try:
@@ -94,6 +98,7 @@
     if tf_major_version == 1:
         load_model_packages |= set(tensorflow1_models)
         load_model_packages.add("stardist_wrong_shape")
+        load_model_packages.add("stardist_wrong_shape2")
     elif tf_major_version == 2:
         load_model_packages |= set(tensorflow2_models)
 
@@ -124,6 +129,12 @@ def stardist_wrong_shape(request):
     return pytest.model_packages[request.param]
 
 
+# written as model group to automatically skip on missing tensorflow 1
+@pytest.fixture(params=[] if skip_tensorflow or tf_major_version != 1 else ["stardist_wrong_shape2"])
+def stardist_wrong_shape2(request):
+    return pytest.model_packages[request.param]
+
+
 # written as model group to automatically skip on missing tensorflow 1
 @pytest.fixture(params=[] if skip_tensorflow or tf_major_version != 1 else ["stardist"])
 def stardist(request):
@@ -165,7 +176,7 @@ def any_tensorflow_js_model(request):
 # fixture to test with all models that should run in the current environment
 # we exclude stardist_wrong_shape here because it is not a valid model
 # and included only to test that validation for this model fails
-@pytest.fixture(params=load_model_packages - {"stardist_wrong_shape"})
+@pytest.fixture(params=load_model_packages - {"stardist_wrong_shape", "stardist_wrong_shape2"})
 def any_model(request):
     return pytest.model_packages[request.param]
 

tests/test_cli.py

@@ -21,6 +21,11 @@ def test_cli_test_model(unet2d_nuclei_broad_model):
     assert ret.returncode == 0
 
 
+def test_cli_test_model_fail(stardist_wrong_shape):
+    ret = subprocess.run(["bioimageio", "test-model", stardist_wrong_shape])
+    assert ret.returncode == 1
+
+
 def test_cli_test_model_with_weight_format(unet2d_nuclei_broad_model):
     ret = subprocess.run(
         ["bioimageio", "test-model", unet2d_nuclei_broad_model, "--weight-format", "pytorch_state_dict"]

tests/test_prediction.py

@@ -8,18 +8,6 @@
 from bioimageio.core.resource_io.nodes import Model
 
 
-def test_test_model(any_model):
-    from bioimageio.core.resource_tests import test_model
-
-    assert test_model(any_model)
-
-
-def test_test_resource(any_model):
-    from bioimageio.core.resource_tests import test_resource
-
-    assert test_resource(any_model)
-
-
 def test_predict_image(any_model, tmpdir):
     from bioimageio.core.prediction import predict_image
 

tests/test_resource_tests/test_test_model.py

@@ -2,7 +2,34 @@ def test_error_for_wrong_shape(stardist_wrong_shape):
     from bioimageio.core.resource_tests import test_model
 
     summary = test_model(stardist_wrong_shape)
-    assert (
-        summary["error"]
-        == "Shape of test input 0 'input' does not match input shape description: ParametrizedInputShape(min=[1, 16, 16, 1], step=[0, 16, 16, 0])"
+    expected_error_message = (
+        "Shape (1, 512, 512, 33) of test output 0 'output' does not match output shape description: "
+        "ImplicitOutputShape(reference_tensor='input', "
+        "scale=[1.0, 1.0, 1.0, 0.0], offset=[1.0, 1.0, 1.0, 33.0])."
     )
+    assert summary["error"] == expected_error_message
+
+
+def test_error_for_wrong_shape2(stardist_wrong_shape2):
+    from bioimageio.core.resource_tests import test_model
+
+    summary = test_model(stardist_wrong_shape2)
+    expected_error_message = (
+        "Shape (1, 512, 512, 1) of test input 0 'input' does not match input shape description: "
+        "ParametrizedInputShape(min=[1, 16, 16, 1], step=[0, 17, 17, 0])."
+    )
+    assert summary["error"] == expected_error_message
+
+
+def test_test_model(any_model):
+    from bioimageio.core.resource_tests import test_model
+
+    summary = test_model(any_model)
+    assert summary["error"] is None
+
+
+def test_test_resource(any_model):
+    from bioimageio.core.resource_tests import test_resource
+
+    summary = test_resource(any_model)
+    assert summary["error"] is None