Improved integration of cellpose

TheJacksonLaboratory · Nov 21, 2024 · e176a23 · e176a23
1 parent 8a99343
commit e176a23
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 1 deletion.
diff --git a/src/napari_activelearning/_models.py b/src/napari_activelearning/_models.py
@@ -170,6 +170,9 @@ def _fine_tune(self, train_data, train_labels, test_data, test_labels):
                 model_name=self._model_name
             )
 
+            if isinstance(self._pretrained_model, tuple):
+                self._pretrained_model = self._pretrained_model[0]
+
             self.refresh_model = True
 
     USING_CELLPOSE = True

diff --git a/src/napari_activelearning/_models_interface.py b/src/napari_activelearning/_models_interface.py
@@ -25,7 +25,23 @@ def cellpose_segmentation_parameters(
                               "cyto2",
                               "cyto3",
                               "nuclei",
-                              "tissuenet_cp3"] = "cyto3",
+                              "tissuenet_cp3",
+                              "livecell_cp3",
+                              "yeast_PhC_cp3",
+                              "yeast_BF_cp3",
+                              "bact_phase_cp3",
+                              "bact_fluor_cp3",
+                              "deepbacs_cp3",
+                              "cyto2_cp3",
+                              "CP",
+                              "CPx",
+                              "TN1",
+                              "TN2",
+                              "TN3",
+                              "LC1",
+                              "LC2",
+                              "LC3",
+                              "LC"] = "cyto3",
           gpu: bool = True
           ):
             return dict(