Merge branch 'mrariden-speedy_improvements'

cellpose/__main__.py

@@ -182,7 +182,7 @@ def main():
                     io.save_masks(image, masks, flows, image_name, png=args.save_png, tif=args.save_tif,
                                   save_flows=args.save_flows,save_outlines=args.save_outlines,
                                   save_ncolor=args.save_ncolor,dir_above=args.dir_above,savedir=args.savedir,
-                                  save_txt=args.save_txt,in_folders=args.in_folders)
+                                  save_txt=args.save_txt,in_folders=args.in_folders, save_mpl=args.save_mpl)
                 if args.save_rois:
                     io.save_rois(masks, image_name)
             logger.info('>>>> completed in %0.3f sec'%(time.time()-tic))

cellpose/cli.py

@@ -118,6 +118,9 @@ def get_arg_parser():
                              help='whether or not to save minimal "n-color" masks (disabled by default')
     output_args.add_argument('--save_txt', action='store_true',
                              help='flag to enable txt outlines for ImageJ (disabled by default)')
+    output_args.add_argument('--save_mpl', action='store_true',
+                             help='save a figure of image/mask/flows using matplotlib (disabled by default). '
+                                  'This is slow, especially with large images.')
 
     # training settings
     training_args = parser.add_argument_group("Training Arguments")

cellpose/io.py

@@ -415,7 +415,7 @@ def save_rois(masks, file_name):
 # Now saves flows, masks, etc. to separate folders.
 def save_masks(images, masks, flows, file_names, png=True, tif=False, channels=[0,0],
                suffix='',save_flows=False, save_outlines=False, save_ncolor=False, 
-               dir_above=False, in_folders=False, savedir=None, save_txt=True):
+               dir_above=False, in_folders=False, savedir=None, save_txt=False, save_mpl=False):
     """ save masks + nicely plotted segmentation image to png and/or tiff
 
     if png, masks[k] for images[k] are saved to file_names[k]+'_cp_masks.png'
@@ -449,14 +449,18 @@ def save_masks(images, masks, flows, file_names, png=True, tif=False, channels=[
         ncolor is a 4 (or 5, if 4 takes too long) index version of the labels that
         is way easier to visualize than having hundreds of unique colors that may
         be similar and touch. Any color map can be applied to it (0,1,2,3,4,...).
+
+    save_mpl: bool
+        If True, saves a matplotlib figure of the original image/segmentation/flows. Does not work for 3D.
+        This takes a long time for large images. Default is False.
     
     """
 
     if isinstance(masks, list):
         for image, mask, flow, file_name in zip(images, masks, flows, file_names):
             save_masks(image, mask, flow, file_name, png=png, tif=tif, suffix=suffix,dir_above=dir_above,
                        save_flows=save_flows,save_outlines=save_outlines,save_ncolor=save_ncolor,
-                       savedir=savedir,save_txt=save_txt,in_folders=in_folders)
+                       savedir=savedir,save_txt=save_txt,in_folders=in_folders, save_mpl=save_mpl)
         return
 
     if masks.ndim > 2 and not tif:
@@ -517,8 +521,10 @@ def save_masks(images, masks, flows, file_names, png=True, tif=False, channels=[
         for ext in exts:
 
             imsave(os.path.join(maskdir,basename + '_cp_masks' + suffix + ext), masks)
-
-    if png and MATPLOTLIB and not min(images.shape) > 3:
+
+    if save_mpl and png and MATPLOTLIB and not min(images.shape) > 3:
+        # Make and save original/segmentation/flows image
+
         img = images.copy()
         if img.ndim<3:
             img = img[:,:,np.newaxis]