add predict_sample_with_fixed_blocking

Author: FynnBe

bioimageio/core/_prediction_pipeline.py

@@ -179,40 +179,17 @@ def get_output_sample_id(self, input_sample_id: SampleId):
                 self.model_description.id or self.model_description.name
             )
 
-    def predict_sample_with_blocking(
+    def predict_sample_with_fixed_blocking(
         self,
         sample: Sample,
+        input_block_shape: Mapping[MemberId, Mapping[AxisId, int]],
+        *,
         skip_preprocessing: bool = False,
         skip_postprocessing: bool = False,
-        ns: Optional[
-            Union[
-                v0_5.ParameterizedSize_N,
-                Mapping[Tuple[MemberId, AxisId], v0_5.ParameterizedSize_N],
-            ]
-        ] = None,
-        batch_size: Optional[int] = None,
     ) -> Sample:
-        """predict a sample by splitting it into blocks according to the model and the `ns` parameter"""
         if not skip_preprocessing:
             self.apply_preprocessing(sample)
 
-        if isinstance(self.model_description, v0_4.ModelDescr):
-            raise NotImplementedError(
-                "predict with blocking not implemented for v0_4.ModelDescr {self.model_description.name}"
-            )
-
-        ns = ns or self._default_ns
-        if isinstance(ns, int):
-            ns = {
-                (ipt.id, a.id): ns
-                for ipt in self.model_description.inputs
-                for a in ipt.axes
-                if isinstance(a.size, v0_5.ParameterizedSize)
-            }
-        input_block_shape = self.model_description.get_tensor_sizes(
-            ns, batch_size or self._default_batch_size
-        ).inputs
-
         n_blocks, input_blocks = sample.split_into_blocks(
             input_block_shape,
             halo=self._default_input_halo,
@@ -239,6 +216,47 @@ def predict_sample_with_blocking(
 
         return predicted_sample
 
+    def predict_sample_with_blocking(
+        self,
+        sample: Sample,
+        skip_preprocessing: bool = False,
+        skip_postprocessing: bool = False,
+        ns: Optional[
+            Union[
+                v0_5.ParameterizedSize_N,
+                Mapping[Tuple[MemberId, AxisId], v0_5.ParameterizedSize_N],
+            ]
+        ] = None,
+        batch_size: Optional[int] = None,
+    ) -> Sample:
+        """predict a sample by splitting it into blocks according to the model and the `ns` parameter"""
+
+        if isinstance(self.model_description, v0_4.ModelDescr):
+            raise NotImplementedError(
+                "`predict_sample_with_blocking` not implemented for v0_4.ModelDescr"
+                + f" {self.model_description.name}."
+                + " Consider using `predict_sample_with_fixed_blocking`"
+            )
+
+        ns = ns or self._default_ns
+        if isinstance(ns, int):
+            ns = {
+                (ipt.id, a.id): ns
+                for ipt in self.model_description.inputs
+                for a in ipt.axes
+                if isinstance(a.size, v0_5.ParameterizedSize)
+            }
+        input_block_shape = self.model_description.get_tensor_sizes(
+            ns, batch_size or self._default_batch_size
+        ).inputs
+
+        return self.predict_sample_with_fixed_blocking(
+            sample,
+            input_block_shape=input_block_shape,
+            skip_preprocessing=skip_preprocessing,
+            skip_postprocessing=skip_postprocessing,
+        )
+
     # def predict(
     #     self,
     #     inputs: Predict_IO,

bioimageio/core/prediction.py

@@ -12,6 +12,7 @@
 )
 
 import xarray as xr
+from loguru import logger
 from numpy.typing import NDArray
 from tqdm import tqdm
 
@@ -41,6 +42,7 @@ def predict(
             Mapping[Tuple[MemberId, AxisId], v0_5.ParameterizedSize_N],
         ]
     ] = None,
+    input_block_shape: Optional[Mapping[MemberId, Mapping[AxisId, int]]] = None,
     skip_preprocessing: bool = False,
     skip_postprocessing: bool = False,
     save_output_path: Optional[Union[Path, str]] = None,
@@ -53,7 +55,11 @@ def predict(
         inputs: the input sample or the named input(s) for this model as a dictionary
         sample_id: the sample id.
         blocksize_parameter: (optional) tile the input into blocks parametrized by
-            blocksize according to any parametrized axis sizes defined in the model RDF
+            blocksize according to any parametrized axis sizes defined in the model RDF.
+            Note: For a predetermined, fixed block shape use `input_block_shape`
+        input_block_shape: (optional) tile the input sample tensors into blocks.
+            Note: For a parameterized block shape, not dealing with the exact block shape,
+            use `blocksie_parameter`.
         skip_preprocessing: flag to skip the model's preprocessing
         skip_postprocessing: flag to skip the model's postprocessing
         save_output_path: A path with `{member_id}` `{sample_id}` in it
@@ -83,19 +89,33 @@ def predict(
             pp.model_description, inputs=inputs, sample_id=sample_id
         )
 
-    if blocksize_parameter is None:
-        output = pp.predict_sample_without_blocking(
+    if input_block_shape is not None:
+        if blocksize_parameter is not None:
+            logger.warning(
+                "ignoring blocksize_parameter={} in favor of input_block_shape={}",
+                blocksize_parameter,
+                input_block_shape,
+            )
+
+        output = pp.predict_sample_with_fixed_blocking(
             sample,
+            input_block_shape=input_block_shape,
             skip_preprocessing=skip_preprocessing,
             skip_postprocessing=skip_postprocessing,
         )
-    else:
+    elif blocksize_parameter is not None:
         output = pp.predict_sample_with_blocking(
             sample,
             skip_preprocessing=skip_preprocessing,
             skip_postprocessing=skip_postprocessing,
             ns=blocksize_parameter,
         )
+    else:
+        output = pp.predict_sample_without_blocking(
+            sample,
+            skip_preprocessing=skip_preprocessing,
+            skip_postprocessing=skip_postprocessing,
+        )
     if save_output_path:
         save_sample(save_output_path, output)