diff --git a/cellpose/__main__.py b/cellpose/__main__.py index 95e68266..c54df6c0 100644 --- 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 index aa6368ca..9d055456 100644 --- 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'))