fixing chan2 loading of denoise model

MouseLand · Feb 21, 2024 · ee96e17 · ee96e17
1 parent aeb8bc8
commit ee96e17
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/cellpose/denoise.py b/cellpose/denoise.py
@@ -643,8 +643,7 @@ def __init__(self, gpu=False, pretrained_model=False, nchan=1, model_type=None,
                 f">>>> model diam_mean = {self.diam_mean: .3f} (ROIs rescaled to this size during training)"
             )
             if chan2 and builtin:
-                chan2_path = "".join(str(
-                    self.pretrained_model).split("_")[:-1]) + "_nuclei"
+                chan2_path = model_path(self.pretrained_model.stem.split("_")[:-1] + "_nuclei")
                 print(f"loading model for chan2: {os.path.split(str(chan2_path)[-1])}")
                 self.net_chan2 = CPnet(self.nbase, self.nclasses, sz=3,
                                        mkldnn=self.mkldnn, max_pool=True,

diff --git a/tests/test_import.py b/tests/test_import.py
@@ -5,7 +5,7 @@ def test_cellpose_imports_without_error():
 
 
 def test_model_zoo_imports_without_error():
-    from cellpose import models
+    from cellpose import models, denoise
     for model_name in models.MODEL_NAMES:
         model = models.CellposeModel(model_type=model_name)