Merge pull request #159 from bioimage-io/improve-test-model

Improve test-model

Author: constantinpape

bioimageio/core/resource_tests.py

@@ -1,14 +1,22 @@
 import traceback
 from pathlib import Path
-from typing import List, Optional, Union
+from typing import List, Optional, Tuple, Union
 
+import numpy
 import numpy as np
 import xarray as xr
+from marshmallow import ValidationError
 
 from bioimageio.core import load_resource_description
 from bioimageio.core.prediction import predict
 from bioimageio.core.prediction_pipeline import create_prediction_pipeline
-from bioimageio.core.resource_io.nodes import Model, ResourceDescription, URI
+from bioimageio.core.resource_io.nodes import (
+    ImplicitOutputShape,
+    Model,
+    ParametrizedInputShape,
+    ResourceDescription,
+    URI,
+)
 from bioimageio.spec.model.raw_nodes import WeightsFormat
 from bioimageio.spec.shared.raw_nodes import ResourceDescription as RawResourceDescription
 
@@ -45,32 +53,93 @@ def test_resource(
     tb: Optional = None
 
     try:
-        model = load_resource_description(model_rdf)
+        rd = load_resource_description(model_rdf)
     except Exception as e:
         error = str(e)
         tb = traceback.format_tb(e.__traceback__)
     else:
-        if isinstance(model, Model):
+        if isinstance(rd, Model):
+            model = rd
             try:
-                prediction_pipeline = create_prediction_pipeline(
-                    bioimageio_model=model, devices=devices, weight_format=weight_format
-                )
                 inputs = [np.load(str(in_path)) for in_path in model.test_inputs]
-                results = predict(prediction_pipeline, inputs)
-                if isinstance(results, (np.ndarray, xr.DataArray)):
-                    results = [results]
-
                 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):
+                        raise ValidationError(
+                            f"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):
+                        error = (error or "") + (
+                            f"Shape of test output {idx} '{out_spec.name}' does not match "
+                            f"output shape description: {out_spec.shape}.\n"
+                        )
+
+                with create_prediction_pipeline(
+                    bioimageio_model=model, devices=devices, weight_format=weight_format
+                ) as prediction_pipeline:
+                    results = predict(prediction_pipeline, inputs)
+
                 if len(results) != len(expected):
-                    error = (
+                    error = (error or "") + (
                         f"Number of outputs and number of expected outputs disagree: {len(results)} != {len(expected)}"
                     )
                 else:
                     for res, exp in zip(results, expected):
                         try:
                             np.testing.assert_array_almost_equal(res, exp, decimal=decimal)
                         except AssertionError as e:
-                            error = f"Output and expected output disagree:\n {e}"
+                            error = (error or "") + f"Output and expected output disagree:\n {e}"
             except Exception as e:
                 error = str(e)
                 tb = traceback.format_tb(e.__traceback__)

tests/conftest.py

@@ -1,11 +1,17 @@
+import logging
+
 import pytest
+
 from bioimageio.core import export_resource_package
 
+logger = logging.getLogger(__name__)
+
 # test models for various frameworks
-torch_models = ["unet2d_fixed_shape", "unet2d_multi_tensor", "unet2d_nuclei_broad_model"]
+torch_models = []
+torch_models_pre_3_10 = ["unet2d_fixed_shape", "unet2d_multi_tensor", "unet2d_nuclei_broad_model"]
 torchscript_models = ["unet2d_multi_tensor", "unet2d_nuclei_broad_model"]
 onnx_models = ["unet2d_multi_tensor", "unet2d_nuclei_broad_model", "hpa_densenet"]
-tensorflow1_models = ["FruNet_model"]
+tensorflow1_models = ["FruNet_model", "stardist"]
 tensorflow2_models = []
 keras_models = ["FruNet_model"]
 tensorflow_js_models = ["FruNet_model"]
@@ -28,12 +34,23 @@
     "hpa_densenet": (
         "https://raw.githubusercontent.com/bioimage-io/spec-bioimage-io/main/example_specs/models/hpa-densenet/rdf.yaml"
     ),
+    "stardist": (
+        "https://raw.githubusercontent.com/bioimage-io/spec-bioimage-io/main/example_specs/models/stardist_example_model/rdf.yaml"
+    ),
+    "stardist_wrong_shape": (
+        "https://raw.githubusercontent.com/bioimage-io/spec-bioimage-io/main/example_specs/models/stardist_example_model/rdf_wrong_shape.yaml"
+    ),
 }
 
 try:
     import torch
+
+    torch_version = tuple(map(int, torch.__version__.split(".")[:2]))
+    logger.warning(f"detected torch version {torch_version}.x")
 except ImportError:
     torch = None
+    torch_version = None
+
 skip_torch = torch is None
 
 try:
@@ -64,6 +81,9 @@
 # load all model packages we need for testing
 load_model_packages = set()
 if not skip_torch:
+    if torch_version < (3, 10):
+        torch_models += torch_models_pre_3_10
+
     load_model_packages |= set(torch_models + torchscript_models)
 
 if not skip_onnx:
@@ -74,6 +94,7 @@
     load_model_packages |= set(tensorflow_js_models)
     if tf_major_version == 1:
         load_model_packages |= set(tensorflow1_models)
+        load_model_packages.add("stardist_wrong_shape")
     elif tf_major_version == 2:
         load_model_packages |= set(tensorflow2_models)
 
@@ -93,7 +114,7 @@ def pytest_configure():
 #
 
 # written as model group to automatically skip on missing torch
-@pytest.fixture(params=[] if skip_torch else ["unet2d_nuclei_broad_model"])
+@pytest.fixture(params=[] if skip_torch or torch_version >= (3, 10) else ["unet2d_nuclei_broad_model"])
 def unet2d_nuclei_broad_model(request):
     return pytest.model_packages[request.param]
 
@@ -104,6 +125,12 @@ def FruNet_model(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_shape"])
+def stardist_wrong_shape(request):
+    return pytest.model_packages[request.param]
+
+
 @pytest.fixture(params=[] if skip_torch else torch_models)
 def any_torch_model(request):
     return pytest.model_packages[request.param]
@@ -146,11 +173,15 @@ def any_model(request):
 # temporary fixtures to test not with all, but only a manual selection of models
 # (models/functionality should be improved to get rid of this specific model group)
 #
-@pytest.fixture(params=[] if skip_torch else ["unet2d_nuclei_broad_model", "unet2d_fixed_shape"])
+@pytest.fixture(
+    params=[] if skip_torch or torch_version >= (3, 10) else ["unet2d_nuclei_broad_model", "unet2d_fixed_shape"]
+)
 def unet2d_fixed_shape_or_not(request):
     return pytest.model_packages[request.param]
 
 
-@pytest.fixture(params=[] if skip_torch else ["unet2d_nuclei_broad_model", "unet2d_multi_tensor"])
+@pytest.fixture(
+    params=[] if skip_torch or torch_version >= (3, 10) else ["unet2d_nuclei_broad_model", "unet2d_multi_tensor"]
+)
 def unet2d_multi_tensor_or_not(request):
     return pytest.model_packages[request.param]

tests/test_resource_tests/test_test_model.py

@@ -0,0 +1,8 @@
+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])"
+    )