adding CLI and fixing bugs with restore

cellpose/__main__.py

@@ -96,6 +96,8 @@ def main():
                 denoise.model_path(restore_type)
             except Exception as e:
                 raise ValueError("restore_type invalid")
+            if args.train or args.train_size:
+                raise ValueError("restore_type cannot be used with training on CLI yet")
 
         model_type = None
         if pretrained_model and not os.path.exists(pretrained_model):
@@ -134,26 +136,39 @@ def main():
                 ">>>> running cellpose on %d images using chan_to_seg %s and chan (opt) %s"
                 % (nimg, cstr0[channels[0]], cstr1[channels[1]]))
 
-            # handle built-in model exceptions; bacterial ones get no size model
-            if builtin_size and restore_type is not None:
+            # handle built-in model exceptions
+            if builtin_size and restore_type is None:
                 model = models.Cellpose(gpu=gpu, device=device, model_type=model_type)
             else:
+                builtin_size = False
                 if args.all_channels:
                     channels = None
                 pretrained_model = None if model_type is not None else pretrained_model
-                model = models.CellposeModel(gpu=gpu, device=device,
-                                             pretrained_model=pretrained_model,
-                                             model_type=model_type)
+                if restore_type is None:
+                    model = models.CellposeModel(gpu=gpu, device=device,
+                                                 pretrained_model=pretrained_model,
+                                                 model_type=model_type)
+                else:
+                    model = denoise.CellposeDenoiseModel(gpu=gpu, device=device,
+                                                         pretrained_model=pretrained_model,
+                                                         model_type=model_type, 
+                                                         restore_type=restore_type,
+                                                         chan2_restore=args.chan2_restore)
 
             # handle diameters
             if args.diameter == 0:
                 if builtin_size:
                     diameter = None
                     logger.info(">>>> estimating diameter for each image")
                 else:
-                    logger.info(
-                        ">>>> not using cyto, cyto2, or nuclei model, cannot auto-estimate diameter"
-                    )
+                    if restore_type is None:
+                        logger.info(
+                            ">>>> not using cyto3, cyto, cyto2, or nuclei model, cannot auto-estimate diameter"
+                        )
+                    else:
+                        logger.info(
+                            ">>>> cannot auto-estimate diameter for image restoration"
+                        )
                     diameter = model.diam_labels
                     logger.info(">>>> using diameter %0.3f for all images" % diameter)
             else:
@@ -175,17 +190,26 @@ def main():
                     channel_axis=args.channel_axis, z_axis=args.z_axis,
                     anisotropy=args.anisotropy, niter=args.niter)
                 masks, flows = out[:2]
-                if len(out) > 3:
+                if len(out) > 3 and restore_type is None:
                     diams = out[-1]
                 else:
                     diams = diameter
+                ratio = 1.
+                if restore_type is not None:
+                    imgs_dn = out[-1]
+                    ratio = diams / model.dn.diam_mean if "upsample" in restore_type else 1.
+                    diams = model.dn.diam_mean if "upsample" in restore_type and model.dn.diam_mean > diams else diams
+                else:
+                    imgs_dn = None
                 if args.exclude_on_edges:
                     masks = utils.remove_edge_masks(masks)
                 if not args.no_npy:
-                    io.masks_flows_to_seg(image, masks, flows, image_name,
-                                          channels=channels, diams=diams)
+                    io.masks_flows_to_seg(image, masks, flows, image_name, imgs_restore=imgs_dn, 
+                                          channels=channels, diams=diams, 
+                                          restore_type=restore_type, ratio=1.)
                 if saving_something:
-                    io.save_masks(image, masks, flows, image_name, png=args.save_png,
+                    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,
                                   dir_above=args.dir_above, savedir=args.savedir,

cellpose/denoise.py

@@ -474,9 +474,9 @@ def __init__(self, gpu=False, pretrained_model=False, model_type=None,
 
     def eval(self, x, batch_size=8, channels=None, channel_axis=None, z_axis=None,
              normalize=True, rescale=None, diameter=None, tile=True, tile_overlap=0.1,
-             resample=True, invert=False, flow_threshold=0.4, cellprob_threshold=0.0, 
-             do_3D=False, anisotropy=None, stitch_threshold=0.0, min_size=15, 
-             niter=None, interp=True):
+             augment=False, resample=True, invert=False, flow_threshold=0.4, 
+             cellprob_threshold=0.0, do_3D=False, anisotropy=None, stitch_threshold=0.0, 
+             min_size=15, niter=None, interp=True):
         """
         Restore array or list of images using the image restoration model, and then segment.
 
@@ -510,6 +510,7 @@ def eval(self, x, batch_size=8, channels=None, channel_axis=None, z_axis=None,
                 if diameter is None, set to diam_mean or diam_train if available. Defaults to None.
             tile (bool, optional): tiles image to ensure GPU/CPU memory usage limited (recommended). Defaults to True.
             tile_overlap (float, optional): fraction of overlap of tiles when computing flows. Defaults to 0.1.
+            augment (bool, optional): augment tiles by flipping and averaging for segmentation. Defaults to False.
             resample (bool, optional): run dynamics at original image size (will be slower but create more accurate boundaries). Defaults to True.
             invert (bool, optional): invert image pixel intensity before running network. Defaults to False.
             flow_threshold (float, optional): flow error threshold (all cells with errors below threshold are kept) (not used for 3D). Defaults to 0.4.
@@ -549,7 +550,8 @@ def eval(self, x, batch_size=8, channels=None, channel_axis=None, z_axis=None,
         diameter = self.dn.diam_mean if self.dn.ratio > 1 else diameter
         masks, flows, styles = self.cp.eval(img_restore, batch_size=batch_size, channels=channels_new, channel_axis=-1,
                                    normalize=normalize_params, rescale=rescale, diameter=diameter,
-                                   tile=tile, tile_overlap=tile_overlap, resample=resample, invert=invert,
+                                   tile=tile, tile_overlap=tile_overlap, augment=augment, 
+                                   resample=resample, invert=invert,
                                    flow_threshold=flow_threshold, cellprob_threshold=cellprob_threshold,
                                    do_3D=do_3D, anisotropy=anisotropy, stitch_threshold=stitch_threshold,
                                    min_size=min_size, niter=niter, interp=interp)
@@ -644,7 +646,7 @@ def __init__(self, gpu=False, pretrained_model=False, nchan=1, model_type=None,
             )
             if chan2 and builtin:
                 chan2_path = model_path(os.path.split(self.pretrained_model)[-1].split("_")[0] + "_nuclei")
-                print(f"loading model for chan2: {os.path.split(str(chan2_path)[-1])}")
+                print(f"loading model for chan2: {os.path.split(str(chan2_path))[-1]}")
                 self.net_chan2 = CPnet(self.nbase, self.nclasses, sz=3,
                                        mkldnn=self.mkldnn, max_pool=True,
                                        diam_mean=17.).to(self.device)