move compute measure to Measure classes

Author: FynnBe

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

@@ -250,7 +250,7 @@ def apply(self, tensor: xr.DataArray) -> xr.DataArray:
             mean, std = self.mean, self.std
         elif self.mode == "per_sample":
             if axes:
-                mean, std = tensor.mean(axes), tensor.std(axes)
+                mean, std = Mean(axes).compute(tensor), Std(axes).compute(tensor)
             else:
                 mean, std = tensor.mean(), tensor.std()
         elif self.mode == "per_dataset":

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)