Merge pull request #226 from bioimage-io/axes_in_trfs

Fix mode fixed with axes

Author: constantinpape

bioimageio/core/prediction_pipeline/_combined_processing.py

@@ -158,16 +158,7 @@ def compute_sample_statistics(
     def _compute_tensor_statistics(tensor: xr.DataArray, measures: Set[Measure]) -> Dict[Measure, Any]:
         ret = {}
         for measure in measures:
-            if isinstance(measure, Mean):
-                v = tensor.mean(dim=measure.axes)
-            elif isinstance(measure, Std):
-                v = tensor.std(dim=measure.axes)
-            elif isinstance(measure, Percentile):
-                v = tensor.quantile(measure.n / 100.0, dim=measure.axes)
-            else:
-                raise NotImplementedError(measure)
-
-            ret[measure] = v
+            ret[measure] = measure.compute(tensor)
 
         return ret
 

bioimageio/core/prediction_pipeline/_processing.py

@@ -12,6 +12,18 @@
     from typing_extensions import Literal, get_args  # type: ignore
 
 
+def _get_fixed(
+    fixed: Union[float, Sequence[float]], tensor: xr.DataArray, axes: Optional[Sequence[str]]
+) -> Union[float, xr.DataArray]:
+    if axes is None:
+        return fixed
+
+    fixed_shape = tuple(s for d, s in tensor.sizes.items() if d not in axes)
+    fixed_dims = tuple(d for d in tensor.dims if d not in axes)
+    fixed = np.array(fixed).reshape(fixed_shape)
+    return xr.DataArray(fixed, dims=fixed_dims)
+
+
 @dataclass
 class Processing:
     """base class for all Pre- and Postprocessing transformations"""
@@ -226,8 +238,8 @@ def apply(self, tensor: xr.DataArray) -> xr.DataArray:
 @dataclass
 class ZeroMeanUnitVariance(Processing):
     mode: Literal["fixed", "per_sample", "per_dataset"] = "per_sample"
-    mean: Optional[float] = None
-    std: Optional[float] = None
+    mean: Optional[Union[float, Sequence[float]]] = None
+    std: Optional[Union[float, Sequence[float]]] = None
     axes: Optional[Sequence[str]] = None
     eps: float = 1.0e-6
 
@@ -247,12 +259,11 @@ def apply(self, tensor: xr.DataArray) -> xr.DataArray:
         axes = None if self.axes is None else tuple(self.axes)
         if self.mode == "fixed":
             assert self.mean is not None and self.std is not None
-            mean, std = self.mean, self.std
+            mean = _get_fixed(self.mean, tensor, axes)
+            std = _get_fixed(self.std, tensor, axes)
         elif self.mode == "per_sample":
-            if axes:
-                mean, std = tensor.mean(axes), tensor.std(axes)
-            else:
-                mean, std = tensor.mean(), tensor.std()
+            mean = Mean(axes).compute(tensor)
+            std = Std(axes).compute(tensor)
         elif self.mode == "per_dataset":
             mean = self.get_computed_dataset_statistics(self.tensor_name, Mean(axes))
             std = self.get_computed_dataset_statistics(self.tensor_name, Std(axes))

bioimageio/core/statistical_measures.py

@@ -1,21 +1,30 @@
 from dataclasses import dataclass
 from typing import Optional, Tuple
 
+import xarray as xr
+
 
 @dataclass(frozen=True)
 class Measure:
-    pass
+    def compute(self, tensor: xr.DataArray):
+        raise NotImplementedError(self.__class__.__name__)
 
 
 @dataclass(frozen=True)
 class Mean(Measure):
     axes: Optional[Tuple[str]] = None
 
+    def compute(self, tensor: xr.DataArray) -> xr.DataArray:
+        return tensor.mean(dim=self.axes)
+
 
 @dataclass(frozen=True)
 class Std(Measure):
     axes: Optional[Tuple[str]] = None
 
+    def compute(self, tensor: xr.DataArray) -> xr.DataArray:
+        return tensor.std(dim=self.axes)
+
 
 @dataclass(frozen=True)
 class Percentile(Measure):
@@ -25,3 +34,6 @@ class Percentile(Measure):
     def __post_init__(self):
         assert self.n >= 0
         assert self.n <= 100
+
+    def compute(self, tensor: xr.DataArray) -> xr.DataArray:
+        return tensor.quantile(self.n / 100.0, dim=self.axes)

tests/prediction_pipeline/test_preprocessing.py

@@ -43,6 +43,21 @@ def test_zero_mean_unit_variance_preprocessing():
     xr.testing.assert_allclose(expected, result)
 
 
+def test_zero_mean_unit_variance_preprocessing_fixed():
+    from bioimageio.core.prediction_pipeline._processing import ZeroMeanUnitVariance
+
+    preprocessing = ZeroMeanUnitVariance(
+        "data_name", mode="fixed", axes=["y"], mean=[1, 4, 7], std=[0.81650, 0.81650, 0.81650]
+    )
+    data = xr.DataArray(np.arange(9).reshape((1, 1, 3, 3)), dims=("b", "c", "x", "y"))
+    expected = xr.DataArray(
+        np.array([[-1.224743, 0.0, 1.224743], [-1.224743, 0.0, 1.224743], [-1.224743, 0.0, 1.224743]])[None, None],
+        dims=("b", "c", "x", "y"),
+    )
+    result = preprocessing(data)
+    xr.testing.assert_allclose(expected, result)
+
+
 def test_zero_mean_unit_across_axes():
     from bioimageio.core.prediction_pipeline._processing import ZeroMeanUnitVariance