Add option to only run official models in "Try all Available Models" (#207)

Author: qin-yu

plantseg/utils.py

@@ -26,7 +26,7 @@ def load_config(config_path: str) -> dict:
     return config
 
 
-def get_model_zoo() -> dict:
+def get_model_zoo(get_custom: bool = True) -> dict:
     """
     returns a dictionary of all models in the model zoo.
     example:
@@ -43,22 +43,24 @@ def get_model_zoo() -> dict:
 
     zoo_config = load_config(zoo_config)
 
-    custom_zoo_config = load_config(custom_zoo)
+    if get_custom:
+        custom_zoo_config = load_config(custom_zoo)
 
-    if custom_zoo_config is None:
-        custom_zoo_config = {}
+        if custom_zoo_config is None:
+            custom_zoo_config = {}
 
-    zoo_config.update(custom_zoo_config)
+        zoo_config.update(custom_zoo_config)
     return zoo_config
 
 
 def list_models(dimensionality_filter: list[str] = None,
                 modality_filter: list[str] = None,
-                output_type_filter: list[str] = None) -> list[str]:
+                output_type_filter: list[str] = None,
+                use_custom_models: bool = True) -> list[str]:
     """
     return a list of models in the model zoo by name
     """
-    zoo_config = get_model_zoo()
+    zoo_config = get_model_zoo(use_custom_models)
     models = list(zoo_config.keys())
 
     if dimensionality_filter is not None:

plantseg/viewer/widget/predictions.py

@@ -152,11 +152,12 @@ def _on_model_name_changed(model_name: str):
     widget_unet_predictions.model_name.tooltip = f'Select a pretrained model. Current model description: {description}'
 
 
-def _compute_multiple_predictions(image, patch_size, device):
+def _compute_multiple_predictions(image, patch_size, device, use_custom_models=True):
     out_layers = []
-    for i, model_name in enumerate(list_models()):
+    model_list = list_models(use_custom_models=use_custom_models)
+    for i, model_name in enumerate(model_list):
 
-        napari_formatted_logging(f'Running UNet Predictions: {model_name} {i}/{len(list_models())}',
+        napari_formatted_logging(f'Running UNet Predictions: {model_name} {i}/{len(model_list)}',
                                  thread='UNet Grid Predictions')
 
         out_name = create_layer_name(image.name, model_name)
@@ -182,15 +183,19 @@ def _compute_multiple_predictions(image, patch_size, device):
           patch_size={'label': 'Patch size',
                       'tooltip': 'Patch size use to processed the data.'},
           device={'label': 'Device',
-                  'choices': ALL_DEVICES}
+                  'choices': ALL_DEVICES},
+          use_custom_models={'label': 'Use custom models',
+                             'tooltip': 'If True, custom models will also be used.'}
           )
 def widget_test_all_unet_predictions(image: Image,
                                      patch_size: Tuple[int, int, int] = (80, 170, 170),
-                                     device: str = ALL_DEVICES[0]) -> Future[List[LayerDataTuple]]:
+                                     device: str = ALL_DEVICES[0],
+                                     use_custom_models: bool = True) -> Future[List[LayerDataTuple]]:
     func = thread_worker(partial(_compute_multiple_predictions,
                                  image=image,
                                  patch_size=patch_size,
-                                 device=device))
+                                 device=device,
+                                 use_custom_models=use_custom_models,))
 
     future = Future()