Update segmentation.py (#72)

* Update segmentation.py Changed function CellPose to CellposeModels which allows for loading of user trained models. * added default model_type if not provided to cellpose --------- Co-authored-by: Jordao Bragantini <[email protected]>
royerlab · Apr 9, 2024 · 4f46179 · 4f46179
1 parent 4a7f71c
commit 4f46179
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/ultrack/imgproc/segmentation.py b/ultrack/imgproc/segmentation.py
@@ -214,7 +214,10 @@ def inverted_edt(
 class Cellpose:
     def __init__(self, **kwargs) -> None:
         """See cellpose.models.Cellpose documentation for details."""
-        from cellpose.models import Cellpose as _Cellpose
+        from cellpose.models import CellposeModel as _Cellpose
+
+        if "pretrained_model" not in kwargs and "model_type" not in kwargs:
+            kwargs["model_type"] = "cyto"
 
         self.model = _Cellpose(**kwargs)
 
@@ -223,5 +226,5 @@ def __call__(self, image: ArrayLike, **kwargs) -> np.ndarray:
         Predicts image labels.
         See cellpose.models.Cellpose.eval documentation for details.
         """
-        labels, _, _, _ = self.model.eval(image, **kwargs)
+        labels, _, _ = self.model.eval(image, **kwargs)
         return labels