Merge pull request #764 from mrariden/cli_updates

Cli updates
MouseLand · Aug 23, 2023 · 95f8d98 · 95f8d98
2 parents 63c79e6 + dbb0b62
commit 95f8d98
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/cellpose/__main__.py b/cellpose/__main__.py
@@ -153,8 +153,9 @@ def main():
             for image_name in tqdm(image_names, file=tqdm_out):
                 image = io.imread(image_name)
                 out = model.eval(image, channels=channels, diameter=diameter,
-                                do_3D=args.do_3D, net_avg=(not args.fast_mode or args.net_avg),
-                                augment=False,
+                                do_3D=args.do_3D,
+                                net_avg=(not args.fast_mode or args.net_avg),
+                                augment=args.augment,
                                 resample=(not args.no_resample and not args.fast_mode),
                                 flow_threshold=args.flow_threshold,
                                 cellprob_threshold=args.cellprob_threshold,
@@ -247,7 +248,8 @@ def main():
                                            save_each=args.save_each,
                                            n_epochs=args.n_epochs,
                                            batch_size=args.batch_size, 
-                                           min_train_masks=args.min_train_masks)
+                                           min_train_masks=args.min_train_masks,
+                                           model_name=args.model_name_out)
                 model.pretrained_model = cpmodel_path
                 logger.info('>>>> model trained and saved to %s'%cpmodel_path)
 

diff --git a/cellpose/cli.py b/cellpose/cli.py
@@ -91,6 +91,8 @@ def get_arg_parser():
                                 help='anisotropy of volume in 3D')
     algorithm_args.add_argument('--exclude_on_edges', action='store_true',
                                 help='discard masks which touch edges of image')
+    algorithm_args.add_argument('--augment', action='store_true',
+                                help='tiles image with overlapping tiles and flips overlapped regions to augment')
 
     # output settings
     output_args = parser.add_argument_group("Output Arguments")
@@ -152,5 +154,8 @@ def get_arg_parser():
                                help='number of epochs to skip between saves. Default: %(default)s')
     training_args.add_argument('--save_each', action='store_true',
                                help='save the model under a different filename per --save_every epoch for later comparsion')
+    training_args.add_argument('--model_name_out', default=None, type=str,
+                               help='Name of model to save as, defaults to name describing model architecture. '
+                                    'Model is saved in the folder specified by --dir in models subfolder.')
 
     return parser