Update grid-search WIP

Author: constantinpape

torch_em/util/grid_search.py

@@ -1,4 +1,8 @@
 import numpy as np
+import torch.nn as nn
+import xarray
+
+import bioimageio.core
 
 from micro_sam.instance_segmentation import InstanceSegmentationWithDecoder
 from micro_sam.evaluation.instance_segmentation import (
@@ -95,20 +99,46 @@ def __init__(self, model, preprocess=None, block_shape=None, halo=None):
 
         self._is_initialized = False
 
-    def initialize(self, data):
+    def _initialize_torch(self, data):
         device = next(iter(self._model.parameters())).device
 
         if self._block_shape is None:
-            scale_factors = self._model.init_kwargs["scale_factors"]
-            min_divisible = [int(np.prod([sf[i] for sf in scale_factors])) for i in range(3)]
+            if hasattr(self._model, "scale_factors"):
+                scale_factors = self._model.init_kwargs["scale_factors"]
+                min_divisible = [int(np.prod([sf[i] for sf in scale_factors])) for i in range(3)]
+            elif hasattr(self._model, "depth"):
+                depth = self._model.depth
+                min_divisible = [2**depth, 2**depth]
+            else:
+                raise RuntimeError
             input_ = self._preprocess(data)
             output = predict_with_padding(self._model, input_, min_divisible, device)
         else:
             output = predict_with_halo(
                 data, self._model, [device], self._block_shape, self._halo,
                 preprocess=self._preprocess,
             )
+        return output
+
+    def _initialize_modelzoo(self, data):
+        if self._block_shape is None:
+            with bioimageio.core.create_prediction_pipeline(self._model) as pp:
+                dims = tuple("bcyx") if data.ndim == 2 else tuple("bczyx")
+                input_ = xarray.DataArray(data[None, None], dims=dims)
+                output = bioimageio.core.prediction.predict_with_padding(pp, input_, padding=True)[0]
+                output = output.squeeze().values
+        else:
+            raise NotImplementedError
+        return output
+
+    # TODO refactor all this so that we can have a common base class that takes care of it
+    def initialize(self, data):
+        if isinstance(self._model, nn.Module):
+            output = self._initialize_torch(data)
+        else:
+            output = self._initialize_modelzoo(data)
 
+        assert output.shape[0] == 3
         self._foreground = output[0]
         self._center_distances = output[1]
         self._boundary_distances = output[2]