Update pre-and-postprocessing WIP

Author: constantinpape

bioimageio/core/prediction_pipeline/_postprocessing.py

@@ -3,17 +3,60 @@
 import xarray as xr
 from bioimageio.core.resource_io.nodes import Postprocessing
 
-from ._preprocessing import binarize, chain
+from . import _preprocessing as ops
 from ._types import Transform
 
 
-def sigmoid(tensor: xr.DataArray, **kwargs):
-    if kwargs:
-        raise NotImplementedError(f"Passed kwargs for sigmoid {kwargs}")
-    return 1 / (1 + xr.ufuncs.exp(-tensor))
+def scale_range(
+    tensor: xr.DataArray,
+    *,
+    reference_tensor=None,
+    mode="per_sample",
+    axes=None,
+    min_percentile=0.0,
+    max_percentile=100.0,
+) -> xr.DataArray:
 
+    # TODO if reference tensor is passed, we need to use it to compute quantiles instead of 'tensor'
+    if reference_tensor is None:
+        tensor_ = tensor
+    else:
+        raise NotImplementedError
 
-KNOWN_POSTPROCESSING = {"binarize": binarize, "sigmoid": sigmoid}
+    # valid modes according to spec: "per_sample", "per_dataset"
+    # TODO implement per_dataset
+    if mode != "per_sample":
+        raise NotImplementedError(f"Unsupported mode for scale_range: {mode}")
+
+    if axes:
+        axes = tuple(axes)
+        v_lower = tensor_.quantile(min_percentile / 100.0, dim=axes)
+        v_upper = tensor_.quantile(max_percentile / 100.0, dim=axes)
+    else:
+        v_lower = tensor_.quantile(min_percentile / 100.0)
+        v_upper = tensor_.quantile(max_percentile / 100.0)
+
+    return ops.ensure_dtype((tensor - v_lower) / v_upper, dtype="float32")
+
+
+# TODO scale the tensor s.t. it matches the mean and variance of the reference tensor
+def scale_mean_variance(tensor: xr.DataArray, *, reference_tensor, mode="per_sample"):
+    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,
+    "scale_linear": ops.scale_linear,
+    "scale_range": ops.scale_range,
+    "sigmoid": ops.sigmoid,
+    "zero_mean_unit_variance": ops.zero_mean_unit_variance
+}
 
 
 def make_postprocessing(spec: List[Postprocessing]) -> Transform:
@@ -32,4 +75,4 @@ def make_postprocessing(spec: List[Postprocessing]) -> Transform:
 
         functions.append((fn, kwargs))
 
-    return chain(*functions)
+    return ops.chain(*functions)

bioimageio/core/prediction_pipeline/_preprocessing.py

@@ -72,12 +72,19 @@ def ensure_dtype(tensor: xr.DataArray, *, dtype):
     return tensor.astype(dtype)
 
 
+def sigmoid(tensor: xr.DataArray, **kwargs):
+    if kwargs:
+        raise NotImplementedError(f"Passed kwargs for sigmoid {kwargs}")
+    return 1.0 / (1.0 + xr.ufuncs.exp(-tensor))
+
+
 KNOWN_PREPROCESSING: Dict[PreprocessingName, Transform] = {
     "scale_linear": scale_linear,
     "zero_mean_unit_variance": zero_mean_unit_variance,
     "binarize": binarize,
     "clip": clip,
-    "scale_range": scale_range
+    "scale_range": scale_range,
+    "sigmoid": sigmoid
     # "__tiktorch_ensure_dtype": ensure_dtype,
 }
 

tests/prediction_pipeline/test_preprocessing.py

@@ -151,3 +151,16 @@ def test_scale_range_axes():
     preprocessing = make_preprocessing([scale_range_spec])
     result = preprocessing(data)
     xr.testing.assert_allclose(expected, result)
+
+
+def test_sigmoid():
+    shape = (3, 32, 32)
+    axes = ("c", "y", "x")
+    np_data = np.random.rand(*shape)
+    data = xr.DataArray(np_data, dims=axes)
+
+    sigmoid = make_preprocessing([Preprocessing("sigmoid")])
+    res = sigmoid(data)
+
+    exp = 1. / (1 + np.exp(-np_data))
+    xr.testing.assert_allclose(res, exp)