From 4f461793397c0a0de10919b3a3a94acfe79b6bd7 Mon Sep 17 00:00:00 2001 From: a-jacobo Date: Tue, 9 Apr 2024 10:13:39 -0700 Subject: [PATCH] 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 --- ultrack/imgproc/segmentation.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/ultrack/imgproc/segmentation.py b/ultrack/imgproc/segmentation.py index 34503b7..ea97749 100644 --- 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