Merge pull request #245 from bioimage-io/tile-update

Fix computation for number of tiles

Author: constantinpape

bioimageio/core/prediction.py

@@ -228,7 +228,7 @@ def load_tile(tile):
 
     if verbose:
         shape = {ax: sh for ax, sh in zip(prediction_pipeline.input_specs[0].axes, input_.shape)}
-        n_tiles = int(np.prod([np.ceil(float(shape[ax]) / tsh) for ax, tsh in tile_shape.items()]))
+        n_tiles = int(np.prod([np.ceil(float(shape[ax]) / (tsh - 2 * halo[ax])) for ax, tsh in tile_shape.items()]))
         tiles = tqdm(tiles, total=n_tiles, desc="prediction with tiling")
 
     # we need to use padded prediction for the individual tiles in case the
@@ -388,7 +388,7 @@ def check_tiling(tiling):
         spatial_axes = [ax for ax in axes if ax in "xyz"]
         halo = tiling["halo"]
         tile = tiling["tile"]
-        assert all(halo.get(ax, 0) > 0 for ax in spatial_axes)
+        assert all(halo.get(ax, 0) >= 0 for ax in spatial_axes)
         assert all(tile.get(ax, 0) > 0 for ax in spatial_axes)
 
     if isinstance(tiling, dict):
@@ -408,7 +408,8 @@ def check_tiling(tiling):
 
             halo = output_spec.halo
             if halo is None:
-                raise ValueError("Model does not provide a valid halo to use for tiling with default parameters")
+                halo = [0] * len(axes)
+            assert len(halo) == len(axes)
 
             tiling = {
                 "halo": {ax: ha for ax, ha in zip(axes, halo) if ax in "xyz"},
@@ -465,7 +466,21 @@ def predict_with_tiling(
             ref_input_shape = dict(zip(ref_input.dims, ref_input.shape))
             output_shape = tuple(int(scale[ax] * ref_input_shape[ax] + 2 * offset[ax]) for ax in output_spec.axes)
         else:
-            output_shape = tuple(output_spec.shape)
+            if len(inputs) > 1:
+                raise NotImplementedError
+            input_spec = prediction_pipeline.input_specs[0]
+            if input_spec.axes != output_spec.axes:
+                raise NotImplementedError("Tiling with a different output shape is not yet supported")
+            out_axes = output_spec.axes
+            fixed_shape = tuple(output_spec.shape)
+            if not all(fixed_shape[out_axes.index(ax)] == tile_shape for ax, tile_shape in tiling["tile"].items()):
+                raise NotImplementedError("Tiling with a different output shape is not yet supported")
+
+            output_shape = list(inputs[0].shape)
+            chan_id = out_axes.index("c")
+            if fixed_shape[chan_id] != output_shape[chan_id]:
+                output_shape[chan_id] = fixed_shape[chan_id]
+            output_shape = tuple(output_shape)
 
         outputs.append(xr.DataArray(np.zeros(output_shape, dtype=output_spec.data_type), dims=tuple(output_spec.axes)))
 

tests/conftest.py

@@ -10,8 +10,7 @@
 warnings.warn(f"testing with bioimageio.spec {bioimageio_spec_version}")
 
 # test models for various frameworks
-torch_models = []
-torch_models_pre_3_10 = [
+torch_models = [
     "unet2d_fixed_shape",
     "unet2d_multi_tensor",
     "unet2d_nuclei_broad_model",
@@ -97,9 +96,6 @@
 # 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:
@@ -130,35 +126,6 @@ def pytest_configure():
 # model groups of the form any_<weight format>_model that include all models providing a specific weight format
 #
 
-# written as model group to automatically skip on missing torch
-@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]
-
-
-# written as model group to automatically skip on missing torch
-@pytest.fixture(params=[] if skip_torch or torch_version >= (3, 10) else ["unet2d_diff_output_shape"])
-def unet2d_diff_output_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_shape"])
-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):
-    return pytest.model_packages[request.param]
-
 
 @pytest.fixture(params=[] if skip_torch else torch_models)
 def any_torch_model(request):
@@ -200,19 +167,22 @@ def any_model(request):
     return pytest.model_packages[request.param]
 
 
+# TODO it would be nice to just generate fixtures for all the individual models dynamically
 #
 # 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 or torch_version >= (3, 10) else ["unet2d_nuclei_broad_model", "unet2d_fixed_shape"]
+    params=[] if skip_torch 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 or torch_version >= (3, 10) else ["unet2d_nuclei_broad_model", "unet2d_multi_tensor"]
+    params=[] if skip_torch else ["unet2d_nuclei_broad_model", "unet2d_multi_tensor"]
 )
 def unet2d_multi_tensor_or_not(request):
     return pytest.model_packages[request.param]
@@ -221,3 +191,39 @@ def unet2d_multi_tensor_or_not(request):
 @pytest.fixture(params=[] if skip_keras else ["unet2d_keras"])
 def unet2d_keras(request):
     return pytest.model_packages[request.param]
+
+
+# written as model group to automatically skip on missing torch
+@pytest.fixture(params=[] if skip_torch else ["unet2d_nuclei_broad_model"])
+def unet2d_nuclei_broad_model(request):
+    return pytest.model_packages[request.param]
+
+
+# written as model group to automatically skip on missing torch
+@pytest.fixture(params=[] if skip_torch else ["unet2d_diff_output_shape"])
+def unet2d_diff_output_shape(request):
+    return pytest.model_packages[request.param]
+
+
+# written as model group to automatically skip on missing torch
+@pytest.fixture(params=[] if skip_torch else ["unet2d_fixed_shape"])
+def unet2d_fixed_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_shape"])
+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):
+    return pytest.model_packages[request.param]

tests/test_prediction.py

@@ -176,6 +176,10 @@ def test_predict_image_with_tiling_channel_last(stardist, tmp_path):
     _test_predict_image_with_tiling(stardist, tmp_path, 0.13)
 
 
+def test_predict_image_with_tiling_fixed_output_shape(unet2d_fixed_shape, tmp_path):
+    _test_predict_image_with_tiling(unet2d_fixed_shape, tmp_path, 0.025)
+
+
 def test_predict_images(unet2d_nuclei_broad_model, tmp_path):
     from bioimageio.core.prediction import predict_images