Ensure output dtype in post-processing

Author: constantinpape

bioimageio/core/prediction_pipeline/_postprocessing.py

@@ -47,11 +47,6 @@ def scale_mean_variance(tensor: xr.DataArray, *, reference_tensor, mode="per_sam
     raise NotImplementedError
 
 
-# NOTE there is a subtle difference between pre-and-postprocessing:
-# pre-processing always returns float32, because the post-processing output is consumed
-# by the model. Post-processing, however, should return the dtype that is specified in the model spec
-# TODO I think the easiest way to implement this is to add dtype is an option to 'make_postprocessing'
-# and then apply 'ensure_dtype' to the result of the postprocessing chain
 KNOWN_POSTPROCESSING = {
     "binarize": ops.binarize,
     "clip": ops.clip,
@@ -62,7 +57,7 @@ def scale_mean_variance(tensor: xr.DataArray, *, reference_tensor, mode="per_sam
 }
 
 
-def make_postprocessing(spec: List[Postprocessing]) -> Transform:
+def make_postprocessing(spec: List[Postprocessing], dtype: str) -> Transform:
     """
     :param preprocessing: bioimage-io spec node
     """
@@ -78,4 +73,9 @@ def make_postprocessing(spec: List[Postprocessing]) -> Transform:
 
         functions.append((fn, kwargs))
 
+    # There is a difference between pre-and-postprocessing:
+    # Tre-processing always returns float32, because its output is consumed y the model.
+    # Post-processing, however, should return the dtype that is specified in the model spec.
+    functions.append((ops.ensure_dtype, {"dtype": dtype}))
+
     return ops.chain(*functions)

bioimageio/core/prediction_pipeline/_prediction_pipeline.py

@@ -172,7 +172,7 @@ def create_prediction_pipeline(
     postprocessing: List[Transform] = []
     for out in bioimageio_model.outputs:
         postprocessing_spec = [] if out.postprocessing is missing else out.postprocessing.copy()
-        postprocessing.append(make_postprocessing(postprocessing_spec))
+        postprocessing.append(make_postprocessing(postprocessing_spec, out.data_type))
 
     return _PredictionPipelineImpl(
         name=bioimageio_model.name,

tests/prediction_pipeline/test_postprocessing.py

@@ -0,0 +1,20 @@
+import numpy as np
+import xarray as xr
+from bioimageio.core.resource_io.nodes import Postprocessing
+from bioimageio.core.prediction_pipeline._postprocessing import make_postprocessing
+
+
+def test_binarize_postprocessing():
+    shape = (3, 32, 32)
+    axes = ("c", "y", "x")
+    np_data = np.random.rand(*shape)
+    data = xr.DataArray(np_data, dims=axes)
+
+    threshold = 0.5
+    exp = xr.DataArray(np_data > threshold, dims=axes)
+
+    for dtype in ("float32", "float64", "uint8", "uint16"):
+        binarize = make_postprocessing(spec=[Postprocessing("binarize", kwargs={"threshold": threshold})], dtype=dtype)
+        res = binarize(data)
+        assert np.dtype(res.dtype) == np.dtype(dtype)
+        xr.testing.assert_allclose(res, exp.astype(dtype))