diff --git a/cellpose/__main__.py b/cellpose/__main__.py index 751756fd..c89a8781 100644 --- 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 index 56495daa..9e544c0b 100644 --- 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