cellpose/__main__.py

"""
Copyright © 2023 Howard Hughes Medical Institute, Authored by Carsen Stringer and Marius Pachitariu.
"""

import sys, os, glob, pathlib, time
import numpy as np
from natsort import natsorted
from tqdm import tqdm
from cellpose import utils, models, io, version_str, train, denoise
from cellpose.cli import get_arg_parser

try:
    from cellpose.gui import gui3d, gui
    GUI_ENABLED = True
except ImportError as err:
    GUI_ERROR = err
    GUI_ENABLED = False
    GUI_IMPORT = True
except Exception as err:
    GUI_ENABLED = False
    GUI_ERROR = err
    GUI_IMPORT = False
    raise

import logging


# settings re-grouped a bit
def main():
    """ Run cellpose from command line
    """

    args = get_arg_parser().parse_args(
    )  # this has to be in a separate file for autodoc to work

    if args.version:
        print(version_str)
        return

    if args.check_mkl:
        mkl_enabled = models.check_mkl()
    else:
        mkl_enabled = True

    if len(args.dir) == 0 and len(args.image_path) == 0:
        if args.add_model:
            io.add_model(args.add_model)
        else:
            if not GUI_ENABLED:
                print("GUI ERROR: %s" % GUI_ERROR)
                if GUI_IMPORT:
                    print(
                        "GUI FAILED: GUI dependencies may not be installed, to install, run"
                    )
                    print("     pip install 'cellpose[gui]'")
            else:
                if args.Zstack:
                    gui3d.run()
                else:
                    gui.run()

    else:
        if args.verbose:
            from .io import logger_setup
            logger, log_file = logger_setup()
        else:
            print(
                ">>>> !LOGGING OFF BY DEFAULT! To see cellpose progress, set --verbose")
            print("No --verbose => no progress or info printed")
            logger = logging.getLogger(__name__)

        use_gpu = False
        channels = [args.chan, args.chan2]

        # find images
        if len(args.img_filter) > 0:
            imf = args.img_filter
        else:
            imf = None

        # Check with user if they REALLY mean to run without saving anything
        if not (args.train or args.train_size):
            saving_something = args.save_png or args.save_tif or args.save_flows or args.save_txt

        device, gpu = models.assign_device(use_torch=True, gpu=args.use_gpu,
                                           device=args.gpu_device)

        if args.pretrained_model is None or args.pretrained_model == "None" or args.pretrained_model == "False" or args.pretrained_model == "0":
            pretrained_model = False
        else:
            pretrained_model = args.pretrained_model

        restore_type = args.restore_type
        if restore_type is not None:
            try:
                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")

        if args.transformer and (restore_type is None):
            default_model = "transformer_cp3"
            backbone = "transformer"
        elif args.transformer and restore_type is not None:
            raise ValueError("no transformer based restoration")
        else:
            default_model = "cyto3"
            backbone = "default"

        if args.norm_percentile is not None:
            value1, value2 = args.norm_percentile
            normalize = {'percentile': (float(value1), float(value2))} 
        else:
            normalize = (not args.no_norm)
        

        model_type = None
        if pretrained_model and not os.path.exists(pretrained_model):
            model_type = pretrained_model if pretrained_model is not None else "cyto3"
            model_strings = models.get_user_models()
            all_models = models.MODEL_NAMES.copy()
            all_models.extend(model_strings)
            if ~np.any([model_type == s for s in all_models]):
                model_type = default_model
                logger.warning(
                    f"pretrained model has incorrect path, using {default_model}")
            if model_type == "nuclei":
                szmean = 17.
            else:
                szmean = 30.
        builtin_size = (model_type == "cyto" or model_type == "cyto2" or
                        model_type == "nuclei" or model_type == "cyto3")

        if len(args.image_path) > 0 and (args.train or args.train_size):
            raise ValueError("ERROR: cannot train model with single image input")

        if not args.train and not args.train_size:
            tic = time.time()
            if len(args.dir) > 0:
                image_names = io.get_image_files(
                    args.dir, args.mask_filter, imf=imf,
                    look_one_level_down=args.look_one_level_down)
            else:
                if os.path.exists(args.image_path):
                    image_names = [args.image_path]
                else:
                    raise ValueError(f"ERROR: no file found at {args.image_path}")
            nimg = len(image_names)

            if args.savedir:
                if not os.path.exists(args.savedir):
                    raise FileExistsError("--savedir {args.savedir} does not exist")

            cstr0 = ["GRAY", "RED", "GREEN", "BLUE"]
            cstr1 = ["NONE", "RED", "GREEN", "BLUE"]
            logger.info(
                ">>>> 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
            if builtin_size and restore_type is None and not args.pretrained_model_ortho:
                model = models.Cellpose(gpu=gpu, device=device, model_type=model_type,
                                        backbone=backbone)
            else:
                builtin_size = False
                if args.all_channels:
                    channels = None
                    img = io.imread(image_names[0])
                    if img.ndim == 3:
                        nchan = min(img.shape)
                    elif img.ndim == 2:
                        nchan = 1
                    channels = None
                else:
                    nchan = 2

                pretrained_model = None if model_type is not None else pretrained_model
                if restore_type is None:
                    pretrained_model_ortho = None if args.pretrained_model_ortho is None else args.pretrained_model_ortho
                    model = models.CellposeModel(gpu=gpu, device=device,
                                                 pretrained_model=pretrained_model,
                                                 model_type=model_type,
                                                 nchan=nchan,
                                                 backbone=backbone,
                                                 pretrained_model_ortho=pretrained_model_ortho)
                else:
                    model = denoise.CellposeDenoiseModel(
                        gpu=gpu, device=device, pretrained_model=pretrained_model,
                        model_type=model_type, restore_type=restore_type, nchan=nchan,
                        chan2_restore=args.chan2_restore)

            # handle diameters
            if args.diameter == 0:
                if builtin_size:
                    diameter = None
                    logger.info(">>>> estimating diameter for each image")
                else:
                    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:
                diameter = args.diameter
                logger.info(">>>> using diameter %0.3f for all images" % diameter)

            tqdm_out = utils.TqdmToLogger(logger, level=logging.INFO)

            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,
                    augment=args.augment, resample=(not args.no_resample),
                    flow_threshold=args.flow_threshold,
                    cellprob_threshold=args.cellprob_threshold,
                    stitch_threshold=args.stitch_threshold, min_size=args.min_size,
                    invert=args.invert, batch_size=args.batch_size,
                    interp=(not args.no_interp), normalize=normalize,
                    channel_axis=args.channel_axis, z_axis=args.z_axis,
                    anisotropy=args.anisotropy, niter=args.niter,
                    flow3D_smooth=args.flow3D_smooth)
                masks, flows = out[:2]
                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,
                                          imgs_restore=imgs_dn, channels=channels,
                                          diams=diams, restore_type=restore_type,
                                          ratio=1.)
                if saving_something:
                    suffix = "_cp_masks"
                    if args.output_name is not None: 
                        # (1) If `savedir` is not defined, then must have a non-zero `suffix`
                        if args.savedir is None and len(args.output_name) > 0:
                            suffix = args.output_name
                        elif args.savedir is not None and not os.path.samefile(args.savedir, args.dir):
                            # (2) If `savedir` is defined, and different from `dir` then                              
                            # takes the value passed as a param. (which can be empty string)
                            suffix = args.output_name

                    io.save_masks(image, masks, flows, image_name,
                                  suffix=suffix, 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,
                                  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))
        else:
            test_dir = None if len(args.test_dir) == 0 else args.test_dir
            images, labels, image_names, train_probs = None, None, None, None
            test_images, test_labels, image_names_test, test_probs = None, None, None, None
            compute_flows = False
            if len(args.file_list) > 0:
                if os.path.exists(args.file_list):
                    dat = np.load(args.file_list, allow_pickle=True).item()
                    image_names = dat["train_files"]
                    image_names_test = dat.get("test_files", None)
                    train_probs = dat.get("train_probs", None)
                    test_probs = dat.get("test_probs", None)
                    compute_flows = dat.get("compute_flows", False)
                    load_files = False
                else:
                    logger.critical(f"ERROR: {args.file_list} does not exist")
            else:
                output = io.load_train_test_data(args.dir, test_dir, imf,
                                                 args.mask_filter,
                                                 args.look_one_level_down)
                images, labels, image_names, test_images, test_labels, image_names_test = output
                load_files = True

            # training with all channels
            if args.all_channels:
                img = images[0] if images is not None else io.imread(image_names[0])
                if img.ndim == 3:
                    nchan = min(img.shape)
                elif img.ndim == 2:
                    nchan = 1
                channels = None
            else:
                nchan = 2

            # model path
            szmean = args.diam_mean
            if not os.path.exists(pretrained_model) and model_type is None:
                if not args.train:
                    error_message = "ERROR: model path missing or incorrect - cannot train size model"
                    logger.critical(error_message)
                    raise ValueError(error_message)
                pretrained_model = False
                logger.info(">>>> training from scratch")
            if args.train:
                logger.info(
                    ">>>> during training rescaling images to fixed diameter of %0.1f pixels"
                    % args.diam_mean)

            # initialize model
            model = models.CellposeModel(
                device=device, model_type=model_type, diam_mean=szmean, nchan=nchan,
                pretrained_model=pretrained_model if model_type is None else None,
                backbone=backbone)

            # train segmentation model
            if args.train:
                cpmodel_path = train.train_seg(
                    model.net, images, labels, train_files=image_names,
                    test_data=test_images, test_labels=test_labels,
                    test_files=image_names_test, train_probs=train_probs,
                    test_probs=test_probs, compute_flows=compute_flows,
                    load_files=load_files, normalize=normalize,
                    channels=channels, channel_axis=args.channel_axis, rgb=(nchan == 3),
                    learning_rate=args.learning_rate, weight_decay=args.weight_decay,
                    SGD=args.SGD, n_epochs=args.n_epochs, batch_size=args.batch_size,
                    min_train_masks=args.min_train_masks,
                    nimg_per_epoch=args.nimg_per_epoch,
                    nimg_test_per_epoch=args.nimg_test_per_epoch,
                    save_path=os.path.realpath(args.dir), save_every=args.save_every,
                    model_name=args.model_name_out)[0]
                model.pretrained_model = cpmodel_path
                logger.info(">>>> model trained and saved to %s" % cpmodel_path)

            # train size model
            if args.train_size:
                sz_model = models.SizeModel(cp_model=model, device=device)
                # data has already been normalized and reshaped
                sz_model.params = train.train_size(
                    model.net, model.pretrained_model, images, labels,
                    train_files=image_names, test_data=test_images,
                    test_labels=test_labels, test_files=image_names_test,
                    train_probs=train_probs, test_probs=test_probs,
                    load_files=load_files, channels=channels,
                    min_train_masks=args.min_train_masks,
                    channel_axis=args.channel_axis, rgb=(nchan == 3),
                    nimg_per_epoch=args.nimg_per_epoch, normalize=normalize,
                    nimg_test_per_epoch=args.nimg_test_per_epoch,
                    batch_size=args.batch_size)
                if test_images is not None:
                    test_masks = [lbl[0] for lbl in test_labels
                                 ] if test_labels is not None else test_labels
                    predicted_diams, diams_style = sz_model.eval(
                        test_images, channels=channels)
                    ccs = np.corrcoef(
                        diams_style,
                        np.array([utils.diameters(lbl)[0] for lbl in test_masks]))[0, 1]
                    cc = np.corrcoef(
                        predicted_diams,
                        np.array([utils.diameters(lbl)[0] for lbl in test_masks]))[0, 1]
                    logger.info(
                        "style test correlation: %0.4f; final test correlation: %0.4f" %
                        (ccs, cc))
                    np.save(
                        os.path.join(
                            args.test_dir,
                            "%s_predicted_diams.npy" % os.path.split(cpmodel_path)[1]),
                        {
                            "predicted_diams": predicted_diams,
                            "diams_style": diams_style
                        })


if __name__ == "__main__":
    main()