fixing cyto2 CLI (#257) and masks_flows_to_seg bug (#256)

MouseLand · May 5, 2021 · 4b1eaa3 · 4b1eaa3
1 parent 4a693d2
commit 4b1eaa3
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/cellpose/__main__.py b/cellpose/__main__.py
@@ -127,7 +127,7 @@ def main():
 
         if not args.train and not args.train_size:
             tic = time.time()
-            if not (args.pretrained_model=='cyto' or args.pretrained_model=='nuclei'):
+            if not (args.pretrained_model=='cyto' or args.pretrained_model=='nuclei' or args.pretrained_model=='cyto2'):
                 cpmodel_path = args.pretrained_model
                 if not os.path.exists(cpmodel_path):
                     logger.warning('model path does not exist, using cyto model')
@@ -154,7 +154,10 @@ def main():
             logger.info('>>>> running cellpose on %d images using chan_to_seg %s and chan (opt) %s'%
                             (nimg, cstr0[channels[0]], cstr1[channels[1]]))
 
-            if args.pretrained_model=='cyto' or args.pretrained_model=='nuclei':
+            if args.pretrained_model=='cyto' or args.pretrained_model=='nuclei' or args.pretrained_model=='cyto2':
+                if args.mxnet and args.pretrained_model=='cyto2':
+                    logger.warning('cyto2 model not available in mxnet, using cyto model')
+                    args.pretrained_model = 'cyto'
                 model = models.Cellpose(gpu=gpu, device=device, model_type=args.pretrained_model, 
                                             torch=(not args.mxnet))
             else:

diff --git a/tests/test_train.py b/tests/test_train.py
@@ -2,6 +2,9 @@
 from pathlib import Path
 from subprocess import check_output, STDOUT
 import os, shutil
+import os 
+
+os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
 
 def test_class_train(data_dir, image_names):
     train_dir = str(data_dir.joinpath('2D').joinpath('train'))