Update CLI and model logic

Author: constantinpape

setup.py

@@ -13,7 +13,7 @@
     license="MIT",
     entry_points={
         "console_scripts": [
-            "synapse_net.run_segmentation = synaptic_reconstruction:tools.cli:segmentation_cli"
+            "synapse_net.run_segmentation = synaptic_reconstruction.tools.cli:segmentation_cli"
         ],
         "napari.manifest": [
             "synaptic_reconstruction = synaptic_reconstruction:napari.yaml",

synaptic_reconstruction/inference/util.py

@@ -255,21 +255,23 @@ def _get_file_paths(input_path, ext=".mrc"):
 
 
 def _load_input(img_path, extra_files, i):
-    # Load the input data data
-    with open_file(img_path, "r") as f:
-
-        # Try to automatically derive the key with the raw data.
-        keys = list(f.keys())
-        if len(keys) == 1:
-            key = keys[0]
-        elif "data" in keys:
-            key = "data"
-        elif "raw" in keys:
-            key = "raw"
-
-        input_volume = f[key][:]
-    assert input_volume.ndim == 3
+    # Load the input data.
+    if os.path.splitext(img_path)[-1] == ".tif":
+        input_volume = imageio.imread(img_path)
 
+    else:
+        with open_file(img_path, "r") as f:
+            # Try to automatically derive the key with the raw data.
+            keys = list(f.keys())
+            if len(keys) == 1:
+                key = keys[0]
+            elif "data" in keys:
+                key = "data"
+            elif "raw" in keys:
+                key = "raw"
+            input_volume = f[key][:]
+
+    assert input_volume.ndim in (2, 3)
     # For now we assume this is always tif.
     if extra_files is not None:
         extra_input = imageio.imread(extra_files[i])

synaptic_reconstruction/tools/cli.py

@@ -1,10 +1,11 @@
 import argparse
 from functools import partial
 
-from .util import run_segmentation
+from .util import run_segmentation, get_model
 from ..inference.util import inference_helper, parse_tiling
 
 
+# TODO: handle kwargs
 def segmentation_cli():
     parser = argparse.ArgumentParser(description="Run segmentation.")
     parser.add_argument(
@@ -16,7 +17,7 @@ def segmentation_cli():
         help="The filepath to directory where the segmentations will be saved."
     )
     parser.add_argument(
-        "--model_path", "-m", required=True, help="The filepath to the vesicle model."
+        "--model", "-m", required=True, help="The model type."
     )
     parser.add_argument(
         "--mask_path", help="The filepath to a tif file with a mask that will be used to restrict the segmentation."
@@ -40,10 +41,11 @@ def segmentation_cli():
     )
     args = parser.parse_args()
 
-    # TODO: preload the model!
+    model = get_model(args.model)
     tiling = parse_tiling(args.tile_shape, args.halo)
+
     segmentation_function = partial(
-        run_segmentation, model_path=args.model_path, verbose=False, tiling=tiling,
+        run_segmentation, model=model, model_type=args.model, verbose=False, tiling=tiling,
     )
     inference_helper(
         args.input_path, args.output_path, segmentation_function,

synaptic_reconstruction/tools/segmentation_widget.py

@@ -4,35 +4,27 @@
 from qtpy.QtWidgets import QWidget, QVBoxLayout, QPushButton, QLabel, QComboBox
 
 from .base_widget import BaseWidget
-from .util import run_segmentation, get_model_registry, _available_devices
+from .util import run_segmentation, get_model, get_model_registry, _available_devices
 
 
 class SegmentationWidget(BaseWidget):
     def __init__(self):
         super().__init__()
 
-        self.model = None
         self.viewer = napari.current_viewer()
         layout = QVBoxLayout()
 
-        # Create the image selection dropdown
+        # Create the image selection dropdown.
         self.image_selector_name = "Image data"
         self.image_selector_widget = self._create_layer_selector(self.image_selector_name, layer_type="Image")
 
-        # create buttons
+        # Create buttons and widgets.
         self.predict_button = QPushButton("Run Segmentation")
-
-        # Connect buttons to functions
         self.predict_button.clicked.connect(self.on_predict)
-        # self.load_model_button.clicked.connect(self.on_load_model)
-
-        # create model selector
         self.model_selector_widget = self.load_model_widget()
-
-        # create advanced settings
         self.settings = self._create_settings_widget()
 
-        # Add the widgets to the layout
+        # Add the widgets to the layout.
         layout.addWidget(self.image_selector_widget)
         layout.addWidget(self.model_selector_widget)
         layout.addWidget(self.settings)
@@ -44,8 +36,7 @@ def load_model_widget(self):
         model_widget = QWidget()
         title_label = QLabel("Select Model:")
 
-        models = list(get_model_registry().urls.keys())
-        self.model = None  # set default model
+        models = ["- choose -"] + list(get_model_registry().urls.keys())
         self.model_selector = QComboBox()
         self.model_selector.addItems(models)
         # Create a layout and add the title label and combo box
@@ -59,23 +50,20 @@ def load_model_widget(self):
 
     def on_predict(self):
         # Get the model and postprocessing settings.
-        model_key = self.model_selector.currentText()
-        if model_key == "- choose -":
+        model_type = self.model_selector.currentText()
+        if model_type == "- choose -":
             show_info("Please choose a model.")
             return
 
-        # loading model
-        model_registry = get_model_registry()
-        model_key = self.model_selector.currentText()
-        # model_path = "/home/freckmann15/.cache/synapse-net/models/vesicles"  #
-        model_path = model_registry.fetch(model_key)
-        # get image data
+        # Load the model.
+        model = get_model(model_type, self.device)
+
+        # Get the image data.
         image = self._get_layer_selector_data(self.image_selector_name)
         if image is None:
             show_info("Please choose an image.")
             return
 
-        # FIXME: don't hard-code tiling here, but figure it out centrally in the prediction function.
         # get tile shape and halo from the viewer
         tiling = {
             "tile": {
@@ -89,26 +77,16 @@ def on_predict(self):
                 "z": 1
             }
         }
-        tile_shape = (self.tile_x_param.value(), self.tile_y_param.value())
-        halo = (self.halo_x_param.value(), self.halo_y_param.value())
-        use_custom_tiling = False
-        for ts, h in zip(tile_shape, halo):
-            if ts != 0 or h != 0:  # if anything changed from default
-                use_custom_tiling = True
-        if use_custom_tiling:
-            segmentation = run_segmentation(
-                image, model_path=model_path, model_key=model_key,
-                tiling=tiling, scale=self.scale_param.value()
-                )
-        else:
-            segmentation = run_segmentation(
-                image, model_path=model_path, model_key=model_key,
-                scale=self.scale_param.value()
-                )
+
+        # TODO: Use scale derived from the image resolution.
+        scale = [self.scale_param.value()]
+        segmentation = run_segmentation(
+            image, model=model, model_type=model_type, tiling=tiling, scale=scale
+        )
 
         # Add the segmentation layer
-        self.viewer.add_labels(segmentation, name=f"{model_key}-segmentation")
-        show_info(f"Segmentation of {model_key} added to layers.")
+        self.viewer.add_labels(segmentation, name=f"{model_type}-segmentation")
+        show_info(f"Segmentation of {model_type} added to layers.")
 
     def _create_settings_widget(self):
         setting_values = QWidget()
@@ -123,6 +101,7 @@ def _create_settings_widget(self):
         setting_values.layout().addLayout(layout)
 
         # Create UI for the tile shape.
+        # TODO: make the tiling 3d and get the default values from 'inference'
         self.tile_x, self.tile_y = 512, 512  # defaults
         self.tile_x_param, self.tile_y_param, layout = self._add_shape_param(
             ("tile_x", "tile_y"), (self.tile_x, self.tile_y), min_val=0, max_val=2048, step=16,