synapse_net.tools.segmentation_widget

  1import inspect
  2import re
  3from typing import Optional, Union
  4
  5import napari
  6import numpy as np
  7import torch
  8
  9from napari.utils.notifications import show_info
 10from qtpy.QtWidgets import QCheckBox, QComboBox, QLabel, QPushButton, QVBoxLayout, QWidget
 11
 12from .base_widget import BaseWidget
 13from ..inference.active_zone import segment_active_zone
 14from ..inference.compartments import segment_compartments
 15from ..inference.cristae import segment_cristae
 16from ..inference.inference import (
 17    _get_model_registry,
 18    _segment_ribbon_AZ,
 19    compute_scale_from_voxel_size,
 20    get_model,
 21    get_model_training_patch_shape,
 22    get_segmentation_function,
 23    run_segmentation,
 24)
 25from ..inference.mitochondria import segment_mitochondria
 26from ..inference.util import get_default_tiling, get_device
 27from ..inference.vesicles import VESICLE_SEGMENTATION_MODES, segment_vesicles
 28
 29
 30_MAX_MIN_SIZE = 100_000_000
 31_POSTPROCESSING_PARAMETER_SPECS = {
 32    segment_vesicles: {
 33        "min_size": {"type": "int", "min": 0, "max": _MAX_MIN_SIZE, "step": 1},
 34        "mode": {"type": "choice", "options": list(VESICLE_SEGMENTATION_MODES)},
 35        "threshold": {"type": "float", "min": 0.0, "max": 1.0, "step": 0.01, "decimals": 2},
 36    },
 37    segment_mitochondria: {
 38        "min_size": {"type": "int", "min": 0, "max": _MAX_MIN_SIZE, "step": 1},
 39        "seed_distance": {"type": "int", "min": 0, "max": 10_000, "step": 1},
 40    },
 41    segment_active_zone: {
 42        "min_size": {"type": "int", "min": 0, "max": _MAX_MIN_SIZE, "step": 1},
 43        "foreground_threshold": {"type": "float", "min": 0.0, "max": 1.0, "step": 0.01, "decimals": 2},
 44    },
 45    segment_compartments: {
 46        "boundary_threshold": {"type": "float", "min": 0.0, "max": 1.0, "step": 0.01, "decimals": 2},
 47        "n_slices_exclude": {"type": "int", "min": 0, "max": 10_000, "step": 1},
 48        "min_z_extent": {"type": "int", "min": 0, "max": 10_000, "step": 1},
 49    },
 50    _segment_ribbon_AZ: {
 51        "threshold": {"type": "float", "min": 0.0, "max": 1.0, "step": 0.01, "decimals": 2},
 52        "n_slices_exclude": {"type": "int", "min": 0, "max": 10_000, "step": 1},
 53        "min_membrane_size": {
 54            "type": "int", "min": 0, "max": _MAX_MIN_SIZE, "step": 1, "default": 50_000,
 55        },
 56        "n_ribbons": {"type": "int", "min": 1, "max": 1_000, "step": 1},
 57    },
 58    segment_cristae: {
 59        "min_size": {
 60            "type": "int", "min": 0, "max": _MAX_MIN_SIZE, "step": 1,
 61            "tooltip": "Minimum cristae size in voxels. Smaller connected components are removed — "
 62                       "raise this to filter out small, detached cristae on the mitochondria membranes.",
 63        },
 64        "foreground_threshold": {
 65            "type": "float", "min": 0.0, "max": 1.0, "step": 0.01, "decimals": 2,
 66            "tooltip": "Probability cutoff for the cristae foreground. Higher values keep only more "
 67                       "confident predictions (less foreground).",
 68        },
 69        "erosion_distance_nm": {
 70            "type": "float", "min": 0.0, "max": 1_000.0, "step": 0.1, "decimals": 1,
 71            "tooltip": "Distance in nm to shrink the mitochondria mask inward before restricting "
 72                       "cristae to it. 0 disables erosion (rely on Minimum Size instead).",
 73        },
 74    },
 75}
 76
 77
 78def _load_custom_model(model_path: str, device: Optional[Union[str, torch.device]] = None) -> torch.nn.Module:
 79    model_path = _clean_filepath(model_path)
 80    if device is None:
 81        device = get_device(device)
 82    try:
 83        model = torch.load(model_path, map_location=torch.device(device), weights_only=False)
 84    except Exception as e:
 85        print(e)
 86        print("model path", model_path)
 87        return None
 88    return model
 89
 90
 91def _available_devices():
 92    available_devices = []
 93    for i in ["cuda", "mps", "cpu"]:
 94        try:
 95            device = get_device(i)
 96        except RuntimeError:
 97            pass
 98        else:
 99            available_devices.append(device)
100    return available_devices
101
102
103def _get_current_tiling(tiling_params: dict, model_type: str):
104    # Read the tiling values from the spin boxes.
105    tiling = {name: {axis: param.value() for axis, param in params.items()} for name, params in tiling_params.items()}
106    if "2d" in model_type:  # A 2d model predicts single slices.
107        tiling["tile"]["z"] = 1
108        tiling["halo"]["z"] = 0
109    show_info(f"Using tiling: {tiling}")
110    return tiling
111
112
113def _clean_filepath(filepath):
114    """Cleans a given filepath by:
115    - Removing newline characters (\n)
116    - Removing escape sequences
117    - Stripping the 'file://' prefix if present
118
119    Args:
120        filepath (str): The original filepath
121
122    Returns:
123        str: The cleaned filepath
124    """
125    # Remove 'file://' prefix if present
126    if filepath.startswith("file://"):
127        filepath = filepath[7:]
128
129    # Remove escape sequences and newlines
130    filepath = re.sub(r'\\.', '', filepath)
131    filepath = filepath.replace('\n', '').replace('\r', '')
132
133    return filepath
134
135
136class SegmentationWidget(BaseWidget):
137    def __init__(self):
138        super().__init__()
139
140        self.viewer = napari.current_viewer()
141        layout = QVBoxLayout()
142
143        # Create the image selection dropdown.
144        self.image_selector_name = "Image data"
145        self.image_selector_widget = self._create_layer_selector(self.image_selector_name, layer_type="Image")
146
147        # Create buttons and widgets.
148        self.predict_button = QPushButton("Run Segmentation")
149        self.predict_button.clicked.connect(self.on_predict)
150        self.model_selector_widget = self.load_model_widget()
151        self.settings = self._create_settings_widget()
152
153        # Add the widgets to the layout.
154        layout.addWidget(self.image_selector_widget)
155        layout.addWidget(self.model_selector_widget)
156        layout.addWidget(self.settings)
157        layout.addWidget(self.predict_button)
158
159        self.setLayout(layout)
160
161    @staticmethod
162    def _clear_layout(layout):
163        while layout.count():
164            item = layout.takeAt(0)
165            widget = item.widget()
166            child_layout = item.layout()
167            if widget is not None:
168                widget.deleteLater()
169            elif child_layout is not None:
170                SegmentationWidget._clear_layout(child_layout)
171                child_layout.deleteLater()
172
173    def _update_tiling(self, model_type):
174        patch_shape = None if model_type == "- choose -" else get_model_training_patch_shape(model_type)
175        tiling = get_default_tiling(patch_shape=patch_shape)
176        for name, values in tiling.items():
177            for axis, value in values.items():
178                self.tiling[name][axis].setValue(value)
179
180    def _update_postprocessing_settings(self, model_type):
181        self._clear_layout(self.postprocessing_settings_layout)
182        self.postprocessing_parameter_widgets = {}
183        if model_type == "- choose -":
184            return
185
186        segmentation_function = get_segmentation_function(model_type)
187        parameter_specs = _POSTPROCESSING_PARAMETER_SPECS.get(segmentation_function, {})
188        function_parameters = inspect.signature(segmentation_function).parameters
189
190        for name, spec in parameter_specs.items():
191            if name not in function_parameters or function_parameters[name].default is inspect.Parameter.empty:
192                raise ValueError(
193                    f"Configured post-processing parameter '{name}' is not an optional parameter "
194                    f"of {segmentation_function.__name__}."
195                )
196            default = spec.get("default", function_parameters[name].default)
197            tooltip = spec.get("tooltip")
198            if spec["type"] == "int":
199                parameter_widget, parameter_layout = self._add_int_param(
200                    name, default, min_val=spec["min"], max_val=spec["max"], step=spec["step"], tooltip=tooltip
201                )
202                self.postprocessing_settings_layout.addLayout(parameter_layout)
203            elif spec["type"] == "float":
204                parameter_widget, parameter_layout = self._add_float_param(
205                    name,
206                    default,
207                    min_val=spec["min"],
208                    max_val=spec["max"],
209                    step=spec["step"],
210                    decimals=spec["decimals"],
211                    tooltip=tooltip,
212                )
213                self.postprocessing_settings_layout.addLayout(parameter_layout)
214            elif spec["type"] == "bool":
215                parameter_widget = self._add_boolean_param(name, default, tooltip=tooltip)
216                self.postprocessing_settings_layout.addWidget(parameter_widget)
217            elif spec["type"] == "choice":
218                parameter_widget, parameter_layout = self._add_choice_param(
219                    name, default, spec["options"], tooltip=tooltip
220                )
221                self.postprocessing_settings_layout.addLayout(parameter_layout)
222            else:
223                raise ValueError(f"Unsupported post-processing parameter type: {spec['type']}")
224            self.postprocessing_parameter_widgets[name] = parameter_widget
225
226    def _get_postprocessing_kwargs(self):
227        kwargs = {}
228        for name, widget in self.postprocessing_parameter_widgets.items():
229            if isinstance(widget, QCheckBox):
230                value = widget.isChecked()
231            elif isinstance(widget, QComboBox):
232                value = widget.currentText()
233            else:
234                value = widget.value()
235            kwargs[name] = value
236        return kwargs
237
238    def load_model_widget(self):
239        model_widget = QWidget()
240        title_label = QLabel("Select Model:")
241
242        # Exclude the models that are only offered through the CLI and not in the plugin.
243        model_list = set(_get_model_registry().urls.keys())
244        # These are the models exlcuded due to their specificity and to keep the menu simple.
245        # TODO: we should at some point update the logic here, to make it easier to support further models
246        # without cluttering the UI.
247        excluded_models = ["vesicles_2d_maus"]
248        model_list = [name for name in model_list if name not in excluded_models]
249
250        models = ["- choose -"] + model_list
251        self.model_selector = QComboBox()
252        self.model_selector.addItems(models)
253        # Create a layout and add the title label and combo box
254        layout = QVBoxLayout()
255        layout.addWidget(title_label)
256        layout.addWidget(self.model_selector)
257
258        # Set layout on the model widget
259        model_widget.setLayout(layout)
260        return model_widget
261
262    def on_predict(self):
263        # Get the model and postprocessing settings.
264        model_type = self.model_selector.currentText()
265        custom_model_path = self.checkpoint_param.text()
266        if model_type == "- choose -":
267            show_info("INFO: Please choose a model.")
268            return
269
270        device = get_device(self.device_dropdown.currentText())
271
272        # Get the image data.
273        image = self._get_layer_selector_data(self.image_selector_name)
274        if image is None:
275            show_info("INFO: Please choose an image.")
276            return
277
278        with self._computing(
279            self.predict_button, "Computing…", "Run Segmentation", "INFO: Running segmentation…"
280        ):
281            # Load the model. Override if user chose custom model.
282            rescale_input = True
283            if custom_model_path:
284                model = _load_custom_model(custom_model_path, device)
285                rescale_input = False
286                if model:
287                    show_info(f"INFO: Using custom model from path: {custom_model_path}")
288                else:
289                    show_info(f"ERROR: Failed to load custom model from path: {custom_model_path}")
290                    return
291            else:
292                model = get_model(model_type, device)
293
294            # Get the current tiling.
295            tiling = _get_current_tiling(self.tiling, model_type)
296
297            # Get the voxel size.
298            metadata = self._get_layer_selector_data(self.image_selector_name, return_metadata=True)
299            voxel_size = self._handle_resolution(metadata, self.voxel_size_param, image.ndim, return_as_list=False)
300
301            # Determine the scaling based on the voxel size.
302            scale = None
303            if voxel_size and rescale_input:
304                # Calculate scale so voxel_size is the same as in training.
305                scale = compute_scale_from_voxel_size(voxel_size, model_type)
306                scale_info = list(map(lambda x: np.round(x, 2), scale))
307                show_info(f"INFO: Rescaled the image by {scale_info} to optimize for the selected model.")
308
309            # Some models require an additional segmentation for inference or postprocessing.
310            # For these models we read out the 'Extra Segmentation' widget.
311            if model_type == "ribbon":  # Currently only the ribbon model needs the extra seg.
312                extra_seg = self._get_layer_selector_data(self.extra_seg_selector_name)
313                resolution = tuple(voxel_size[ax] for ax in "zyx")
314                kwargs = {"extra_segmentation": extra_seg, "resolution": resolution}
315            elif "cristae" in model_type:  # Cristae model expects 2 3D volumes
316                kwargs = {
317                    "extra_segmentation": self._get_layer_selector_data(self.extra_seg_selector_name),
318                    "with_channels": True,
319                    "channels_to_standardize": [0]
320                }
321            else:
322                kwargs = {}
323            kwargs.update(self._get_postprocessing_kwargs())
324            segmentation = run_segmentation(
325                image, model=model, model_type=model_type, tiling=tiling, scale=scale, **kwargs
326            )
327
328            # Add the segmentation layer(s).
329            if isinstance(segmentation, dict):
330                for name, seg in segmentation.items():
331                    self.viewer.add_labels(seg, name=name, metadata=metadata)
332            else:
333                self.viewer.add_labels(segmentation, name=f"{model_type}", metadata=metadata)
334
335        show_info(f"INFO: Segmentation of {model_type} added to layers.")
336
337    def _create_settings_widget(self):
338        setting_values = QWidget()
339        # setting_values.setToolTip(get_tooltip("embedding", "settings"))
340        setting_values.setLayout(QVBoxLayout())
341
342        # Create UI for the device.
343        device = "auto"
344        device_options = ["auto"] + _available_devices()
345
346        self.device_dropdown, layout = self._add_choice_param("device", device, device_options)
347        setting_values.layout().addLayout(layout)
348
349        # Create UI for the tile shape. The values are set to the default tiling of the selected model.
350        self.tiling = {"tile": {}, "halo": {}}
351        self.tiling["tile"]["x"], self.tiling["tile"]["y"], self.tiling["tile"]["z"], layout = self._add_shape_param(
352            ("tile_x", "tile_y", "tile_z"), (0, 0, 0), min_val=0, max_val=2048, step=16,
353            # tooltip=get_tooltip("embedding", "tiling")
354        )
355        setting_values.layout().addLayout(layout)
356
357        # Create UI for the halo.
358        self.tiling["halo"]["x"], self.tiling["halo"]["y"], self.tiling["halo"]["z"], layout = self._add_shape_param(
359            ("halo_x", "halo_y", "halo_z"), (0, 0, 0), min_val=0, max_val=512,
360            # tooltip=get_tooltip("embedding", "halo")
361        )
362        setting_values.layout().addLayout(layout)
363        self.model_selector.currentTextChanged.connect(self._update_tiling)
364        self._update_tiling(self.model_selector.currentText())
365
366        # Read voxel size from layer metadata.
367        self.voxel_size_param, layout = self._add_float_param(
368            "voxel_size", 0.0, min_val=0.0, max_val=100.0,
369        )
370        setting_values.layout().addLayout(layout)
371
372        self.checkpoint_param, layout = self._add_string_param(
373            name="checkpoint", value="", title="Load Custom Model",
374            placeholder="path/to/checkpoint.pt",
375        )
376        setting_values.layout().addLayout(layout)
377
378        # Add selection UI for additional segmentation, which some models require for inference or postproc.
379        self.extra_seg_selector_name = "Extra Segmentation"
380        self.extra_selector_widget = self._create_layer_selector(self.extra_seg_selector_name, layer_type="Labels")
381        setting_values.layout().addWidget(self.extra_selector_widget)
382
383        # Add model-specific post-processing settings that are updated when the selected model changes.
384        setting_values.layout().addWidget(QLabel("Post-processing:"))
385        self.postprocessing_settings_widget = QWidget()
386        self.postprocessing_settings_layout = QVBoxLayout()
387        self.postprocessing_settings_layout.setContentsMargins(0, 0, 0, 0)
388        self.postprocessing_settings_widget.setLayout(self.postprocessing_settings_layout)
389        setting_values.layout().addWidget(self.postprocessing_settings_widget)
390        self.postprocessing_parameter_widgets = {}
391        self.model_selector.currentTextChanged.connect(self._update_postprocessing_settings)
392        self._update_postprocessing_settings(self.model_selector.currentText())
393
394        settings = self._make_collapsible(widget=setting_values, title="Advanced Settings")
395        return settings