diff --git a/cellpose_napari/_dock_widget.py b/cellpose_napari/_dock_widget.py index c09342e..37d20f5 100644 --- a/cellpose_napari/_dock_widget.py +++ b/cellpose_napari/_dock_widget.py @@ -13,6 +13,10 @@ from magicgui import magicgui import sys +CP_models = ["cyto3", "cyto2", "cyto", "nuclei", "tissuenet_cp3", + "livecell_cp3", "yeast_PhC_cp3", "yeast_BF_cp3", "bact_phase_cp3", + "bact_fluor_cp3", "deepbacs_cp3", "cyto2_cp3"] + # initialize logger # use -v or --verbose when starting napari to increase verbosity logger = logging.getLogger(__name__) @@ -91,8 +95,9 @@ def run_cellpose(image, model_type, custom_model, channels, channel_axis, diamet @thread_worker def compute_diameter(image, channels, model_type): from cellpose import models + model_type0 = model_type if model_type in CP_models[:4] else "cyto3" - CP = models.Cellpose(model_type = model_type, gpu=True) + CP = models.Cellpose(model_type = model_type0, gpu=True) diam = CP.sz.eval(image, channels=channels, channel_axis=-1)[0] diam = np.around(diam, 2) del CP @@ -121,7 +126,7 @@ def compute_masks(masks_orig, flows_orig, cellprob_threshold, flow_threshold): @magicgui( call_button='run segmentation', layout='vertical', - model_type = dict(widget_type='ComboBox', label='model type', choices=['cyto', 'nuclei', 'cyto2', 'custom'], value='cyto', tooltip='there is a cyto model, a new cyto2 model from user submissions, and a nuclei model'), + model_type = dict(widget_type='ComboBox', label='model type', choices=[*CP_models, 'custom'], value='cyto3', tooltip='there is a cyto model, a new cyto2 model from user submissions, and a nuclei model'), custom_model = dict(widget_type='FileEdit', label='custom model path: ', tooltip='if model type is custom, specify file path to it here'), main_channel = dict(widget_type='ComboBox', label='channel to segment', choices=main_channel_choices, value=0, tooltip='choose channel with cells'), optional_nuclear_channel = dict(widget_type='ComboBox', label='optional nuclear channel', choices=optional_nuclear_channel_choices, value=0, tooltip='optional, if available, choose channel with nuclei of cells'),