initialize Percentile with axes in ScaleRange

Author: FynnBe

bioimageio/core/prediction_pipeline/_processing.py

@@ -167,14 +167,20 @@ def get_required_dataset_statistics(self) -> Dict[str, Set[Measure]]:
         if self.mode == "per_sample":
             return {}
         elif self.mode == "per_dataset":
-            measures = {Percentile(self.min_percentile), Percentile(self.max_percentile)}
+            measures = {
+                Percentile(self.min_percentile, axes=self.axes),
+                Percentile(self.max_percentile, axes=self.axes),
+            }
             return {self.reference_tensor or self.tensor_name: measures}
         else:
             raise ValueError(self.mode)
 
     def get_required_sample_statistics(self) -> Dict[str, Set[Measure]]:
         if self.mode == "per_sample":
-            measures = {Percentile(self.min_percentile), Percentile(self.max_percentile)}
+            measures = {
+                Percentile(self.min_percentile, axes=self.axes),
+                Percentile(self.max_percentile, axes=self.axes),
+            }
             return {self.reference_tensor or self.tensor_name: measures}
         elif self.mode == "per_dataset":
             return {}
@@ -200,6 +206,10 @@ def apply(self, tensor: xr.DataArray) -> xr.DataArray:
 
         return ensure_dtype((tensor - v_lower) / v_upper, dtype="float32")
 
+    def __post_init__(self):
+        super().__post_init__()
+        self.axes = None if self.axes is None else tuple(self.axes)  # make sure axes is Tuple[str] or None
+
 
 @dataclass
 class Sigmoid(Processing):