update pre- and postprocessing tests

Author: FynnBe

tests/prediction_pipeline/test_postprocessing.py

@@ -1,7 +1,7 @@
 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
+from bioimageio.core.prediction_pipeline._combined_processing import make_postprocessing
 
 
 def test_binarize_postprocessing():

tests/prediction_pipeline/test_preprocessing.py

@@ -1,31 +1,33 @@
 import numpy as np
-import pytest
 import xarray as xr
-from bioimageio.core.resource_io.nodes import Preprocessing
 
-from bioimageio.core.prediction_pipeline._preprocessing import make_preprocessing
+from bioimageio.core.prediction_pipeline._combined_processing import CombinedProcessing
 
 
 def test_scale_linear():
-    spec = Preprocessing(name="scale_linear", kwargs={"offset": [1, 2, 42], "gain": [1, 2, 3], "axes": "yx"})
-    data = xr.DataArray(np.arange(6).reshape(1, 2, 3), dims=("x", "y", "c"))
+    from bioimageio.core.prediction_pipeline._processing import ScaleLinear
+
+    preprocessing = ScaleLinear("data_name", offset=[1, 2, 42], gain=[1, 2, 3], axes="yx")
+    data = xr.DataArray(np.arange(6).reshape((1, 2, 3)), dims=("x", "y", "c"))
     expected = xr.DataArray(np.array([[[1, 4, 48], [4, 10, 57]]]), dims=("x", "y", "c"))
-    preprocessing = make_preprocessing([spec])
-    result = preprocessing(data)
+    result = preprocessing.apply(data)
     xr.testing.assert_allclose(expected, result)
 
 
 def test_scale_linear_no_channel():
-    spec = Preprocessing(name="scale_linear", kwargs={"offset": 1, "gain": 2, "axes": "yx"})
+    from bioimageio.core.prediction_pipeline._processing import ScaleLinear
+
+    preprocessing = ScaleLinear("data_name", offset=1, gain=2, axes="yx")
     data = xr.DataArray(np.arange(6).reshape(2, 3), dims=("x", "y"))
     expected = xr.DataArray(np.array([[1, 3, 5], [7, 9, 11]]), dims=("x", "y"))
-    preprocessing = make_preprocessing([spec])
-    result = preprocessing(data)
+    result = preprocessing.apply(data)
     xr.testing.assert_allclose(expected, result)
 
 
 def test_zero_mean_unit_variance_preprocessing():
-    zero_mean_spec = Preprocessing(name="zero_mean_unit_variance", kwargs={})
+    from bioimageio.core.prediction_pipeline._processing import ZeroMeanUnitVariance
+
+    preprocessing = ZeroMeanUnitVariance("data_name")
     data = xr.DataArray(np.arange(9).reshape(3, 3), dims=("x", "y"))
     expected = xr.DataArray(
         np.array(
@@ -37,14 +39,15 @@ def test_zero_mean_unit_variance_preprocessing():
         ),
         dims=("x", "y"),
     )
-    preprocessing = make_preprocessing([zero_mean_spec])
     result = preprocessing(data)
     xr.testing.assert_allclose(expected, result)
 
 
 def test_zero_mean_unit_across_axes():
-    zero_mean_spec = Preprocessing(name="zero_mean_unit_variance", kwargs={"axes": ("x", "y")})
-    data = xr.DataArray(np.arange(18).reshape(2, 3, 3), dims=("c", "x", "y"))
+    from bioimageio.core.prediction_pipeline._processing import ZeroMeanUnitVariance
+
+    preprocessing = ZeroMeanUnitVariance("data_name", axes=("x", "y"))
+    data = xr.DataArray(np.arange(18).reshape((2, 3, 3)), dims=("c", "x", "y"))
     expected = xr.DataArray(
         np.array(
             [
@@ -55,54 +58,51 @@ def test_zero_mean_unit_across_axes():
         ),
         dims=("x", "y"),
     )
-    preprocessing = make_preprocessing([zero_mean_spec])
     result = preprocessing(data)
     xr.testing.assert_allclose(expected, result[dict(c=0)])
 
 
 def test_zero_mean_unit_variance_fixed():
+    from bioimageio.core.prediction_pipeline._processing import ZeroMeanUnitVariance
+
     np_data = np.arange(9).reshape(3, 3)
     mean = np_data.mean()
     std = np_data.mean()
     eps = 1.0e-7
-    kwargs = {"mode": "fixed", "mean": mean, "std": std, "eps": eps}
-    zero_mean_spec = Preprocessing(name="zero_mean_unit_variance", kwargs=kwargs)
-    data = xr.DataArray(np_data, dims=("x", "y"))
+    preprocessing = ZeroMeanUnitVariance("data_name", mode="fixed", mean=mean, std=std, eps=eps)
 
+    data = xr.DataArray(np_data, dims=("x", "y"))
     expected = xr.DataArray((np_data - mean) / (std + eps), dims=("x", "y"))
-    preprocessing = make_preprocessing([zero_mean_spec])
     result = preprocessing(data)
     xr.testing.assert_allclose(expected, result)
 
 
 def test_binarize():
-    binarize_spec = Preprocessing(name="binarize", kwargs={"threshold": 14})
-    data = xr.DataArray(np.arange(30).reshape(2, 3, 5), dims=("x", "y", "c"))
+    from bioimageio.core.prediction_pipeline._processing import Binarize
+
+    preprocessing = Binarize("data_name", threshold=14)
+    data = xr.DataArray(np.arange(30).reshape((2, 3, 5)), dims=("x", "y", "c"))
     expected = xr.zeros_like(data)
     expected[{"x": slice(1, None)}] = 1
-    preprocessing = make_preprocessing([binarize_spec])
     result = preprocessing(data)
     xr.testing.assert_allclose(expected, result)
 
 
 def test_clip_preprocessing():
-    clip_spec = Preprocessing(name="clip", kwargs={"min": 3, "max": 5})
+    from bioimageio.core.prediction_pipeline._processing import Clip
+
+    preprocessing = Clip("data_name", min=3, max=5)
     data = xr.DataArray(np.arange(9).reshape(3, 3), dims=("x", "y"))
     expected = xr.DataArray(np.array([[3, 3, 3], [3, 4, 5], [5, 5, 5]]), dims=("x", "y"))
-    preprocessing = make_preprocessing([clip_spec])
     result = preprocessing(data)
     xr.testing.assert_equal(expected, result)
 
 
-def test_unknown_preprocessing_should_raise():
-    mypreprocessing = Preprocessing(name="mycoolpreprocessing", kwargs={"axes": ("x", "y")})
-    with pytest.raises(NotImplementedError):
-        make_preprocessing([mypreprocessing])
-
-
 def test_combination_of_preprocessing_steps_with_dims_specified():
-    zero_mean_spec = Preprocessing(name="zero_mean_unit_variance", kwargs={"axes": ("x", "y")})
-    data = xr.DataArray(np.arange(18).reshape(2, 3, 3), dims=("c", "x", "y"))
+    from bioimageio.core.prediction_pipeline._processing import ZeroMeanUnitVariance
+
+    preprocessing = ZeroMeanUnitVariance("data_name", axes=("x", "y"))
+    data = xr.DataArray(np.arange(18).reshape((2, 3, 3)), dims=("c", "x", "y"))
 
     expected = xr.DataArray(
         np.array(
@@ -115,51 +115,65 @@ def test_combination_of_preprocessing_steps_with_dims_specified():
         dims=("x", "y"),
     )
 
-    preprocessing = make_preprocessing([zero_mean_spec])
     result = preprocessing(data)
     xr.testing.assert_allclose(expected, result[dict(c=0)])
 
 
 def test_scale_range():
-    scale_range_spec = Preprocessing(name="scale_range", kwargs={})
+    from bioimageio.core.prediction_pipeline._processing import ScaleRange
 
+    preprocessing = ScaleRange("data_name")
     np_data = np.arange(9).reshape(3, 3).astype("float32")
     data = xr.DataArray(np_data, dims=("x", "y"))
+    preprocessing.set_computed_sample_statistics(
+        CombinedProcessing.compute_sample_statistics(
+            {"data_name": data}, preprocessing.get_required_sample_statistics()
+        )
+    )
 
     exp_data = (np_data - np_data.min()) / np_data.max()
     expected = xr.DataArray(exp_data, dims=("x", "y"))
 
-    preprocessing = make_preprocessing([scale_range_spec])
     result = preprocessing(data)
     xr.testing.assert_allclose(expected, result)
 
 
 def test_scale_range_axes():
+    from bioimageio.core.prediction_pipeline._processing import ScaleRange
+
     min_percentile = 1.0
     max_percentile = 99.0
-    kwargs = {"axes": ("x", "y"), "min_percentile": min_percentile, "max_percentile": max_percentile}
-    scale_range_spec = Preprocessing(name="scale_range", kwargs=kwargs)
+    preprocessing = ScaleRange(
+        "data_name", axes=("x", "y"), min_percentile=min_percentile, max_percentile=max_percentile
+    )
 
-    np_data = np.arange(18).reshape(2, 3, 3).astype("float32")
+    np_data = np.arange(18).reshape((2, 3, 3)).astype("float32")
     data = xr.DataArray(np_data, dims=("c", "x", "y"))
 
+    preprocessing.set_computed_sample_statistics(
+        CombinedProcessing.compute_sample_statistics(
+            {"data_name": data}, preprocessing.get_required_sample_statistics()
+        )
+    )
+
     p_low = np.percentile(np_data, min_percentile, axis=(1, 2), keepdims=True)
     p_up = np.percentile(np_data, max_percentile, axis=(1, 2), keepdims=True)
     exp_data = (np_data - p_low) / p_up
     expected = xr.DataArray(exp_data, dims=("c", "x", "y"))
 
-    preprocessing = make_preprocessing([scale_range_spec])
     result = preprocessing(data)
     xr.testing.assert_allclose(expected, result)
 
 
 def test_sigmoid():
+    from bioimageio.core.prediction_pipeline._processing import 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", kwargs={})])
+    sigmoid = Sigmoid("data_name")
     res = sigmoid(data)
 
     exp = xr.DataArray(1.0 / (1 + np.exp(-np_data)), dims=axes)