synapse_net.tools.segmentation_widget

  1import copy
  2import inspect
  3import re
  4from typing import Optional, Union
  5
  6import napari
  7import numpy as np
  8import torch
  9
 10from napari.utils.notifications import show_info
 11from qtpy.QtWidgets import QCheckBox, QComboBox, QLabel, QPushButton, QVBoxLayout, QWidget
 12
 13from .base_widget import BaseWidget
 14from ..inference.active_zone import segment_active_zone
 15from ..inference.compartments import segment_compartments
 16from ..inference.cristae import segment_cristae
 17from ..inference.inference import (
 18    _get_model_registry,
 19    _segment_ribbon_AZ,
 20    compute_scale_from_voxel_size,
 21    get_model,
 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: dict, default_tiling: dict, model_type: str):
104    # get tiling values from qt objects
105    for k, v in tiling.items():
106        for k2, v2 in v.items():
107            if isinstance(v2, int):
108                continue
109            elif hasattr(v2, "value"):  # If it's a QSpinBox, extract the value
110                tiling[k][k2] = v2.value()
111            else:
112                raise TypeError(f"Unexpected type for tiling value: {type(v2)} at {k}/{k2}")
113    # check if user inputs tiling/halo or not
114    if default_tiling == tiling:
115        if "2d" in model_type:
116            # if its a 2d model expand x,y and set z to 1
117            tiling = {
118                "tile": {"x": 512, "y": 512, "z": 1},
119                "halo": {"x": 64, "y": 64, "z": 1},
120            }
121    else:
122        show_info(f"Using custom tiling: {tiling}")
123    if "2d" in model_type:
124        # if its a 2d model set z to 1
125        tiling["tile"]["z"] = 1
126        tiling["halo"]["z"] = 0
127        show_info(f"Using tiling: {tiling}")
128    return tiling
129
130
131def _clean_filepath(filepath):
132    """Cleans a given filepath by:
133    - Removing newline characters (\n)
134    - Removing escape sequences
135    - Stripping the 'file://' prefix if present
136
137    Args:
138        filepath (str): The original filepath
139
140    Returns:
141        str: The cleaned filepath
142    """
143    # Remove 'file://' prefix if present
144    if filepath.startswith("file://"):
145        filepath = filepath[7:]
146
147    # Remove escape sequences and newlines
148    filepath = re.sub(r'\\.', '', filepath)
149    filepath = filepath.replace('\n', '').replace('\r', '')
150
151    return filepath
152
153
154class SegmentationWidget(BaseWidget):
155    def __init__(self):
156        super().__init__()
157
158        self.viewer = napari.current_viewer()
159        layout = QVBoxLayout()
160        self.tiling = {}
161
162        # Create the image selection dropdown.
163        self.image_selector_name = "Image data"
164        self.image_selector_widget = self._create_layer_selector(self.image_selector_name, layer_type="Image")
165
166        # Create buttons and widgets.
167        self.predict_button = QPushButton("Run Segmentation")
168        self.predict_button.clicked.connect(self.on_predict)
169        self.model_selector_widget = self.load_model_widget()
170        self.settings = self._create_settings_widget()
171
172        # Add the widgets to the layout.
173        layout.addWidget(self.image_selector_widget)
174        layout.addWidget(self.model_selector_widget)
175        layout.addWidget(self.settings)
176        layout.addWidget(self.predict_button)
177
178        self.setLayout(layout)
179
180    @staticmethod
181    def _clear_layout(layout):
182        while layout.count():
183            item = layout.takeAt(0)
184            widget = item.widget()
185            child_layout = item.layout()
186            if widget is not None:
187                widget.deleteLater()
188            elif child_layout is not None:
189                SegmentationWidget._clear_layout(child_layout)
190                child_layout.deleteLater()
191
192    def _update_postprocessing_settings(self, model_type):
193        self._clear_layout(self.postprocessing_settings_layout)
194        self.postprocessing_parameter_widgets = {}
195        if model_type == "- choose -":
196            return
197
198        segmentation_function = get_segmentation_function(model_type)
199        parameter_specs = _POSTPROCESSING_PARAMETER_SPECS.get(segmentation_function, {})
200        function_parameters = inspect.signature(segmentation_function).parameters
201
202        for name, spec in parameter_specs.items():
203            if name not in function_parameters or function_parameters[name].default is inspect.Parameter.empty:
204                raise ValueError(
205                    f"Configured post-processing parameter '{name}' is not an optional parameter "
206                    f"of {segmentation_function.__name__}."
207                )
208            default = spec.get("default", function_parameters[name].default)
209            tooltip = spec.get("tooltip")
210            if spec["type"] == "int":
211                parameter_widget, parameter_layout = self._add_int_param(
212                    name, default, min_val=spec["min"], max_val=spec["max"], step=spec["step"], tooltip=tooltip
213                )
214                self.postprocessing_settings_layout.addLayout(parameter_layout)
215            elif spec["type"] == "float":
216                parameter_widget, parameter_layout = self._add_float_param(
217                    name,
218                    default,
219                    min_val=spec["min"],
220                    max_val=spec["max"],
221                    step=spec["step"],
222                    decimals=spec["decimals"],
223                    tooltip=tooltip,
224                )
225                self.postprocessing_settings_layout.addLayout(parameter_layout)
226            elif spec["type"] == "bool":
227                parameter_widget = self._add_boolean_param(name, default, tooltip=tooltip)
228                self.postprocessing_settings_layout.addWidget(parameter_widget)
229            elif spec["type"] == "choice":
230                parameter_widget, parameter_layout = self._add_choice_param(
231                    name, default, spec["options"], tooltip=tooltip
232                )
233                self.postprocessing_settings_layout.addLayout(parameter_layout)
234            else:
235                raise ValueError(f"Unsupported post-processing parameter type: {spec['type']}")
236            self.postprocessing_parameter_widgets[name] = parameter_widget
237
238    def _get_postprocessing_kwargs(self):
239        kwargs = {}
240        for name, widget in self.postprocessing_parameter_widgets.items():
241            if isinstance(widget, QCheckBox):
242                value = widget.isChecked()
243            elif isinstance(widget, QComboBox):
244                value = widget.currentText()
245            else:
246                value = widget.value()
247            kwargs[name] = value
248        return kwargs
249
250    def load_model_widget(self):
251        model_widget = QWidget()
252        title_label = QLabel("Select Model:")
253
254        # Exclude the models that are only offered through the CLI and not in the plugin.
255        model_list = set(_get_model_registry().urls.keys())
256        # These are the models exlcuded due to their specificity and to keep the menu simple.
257        # TODO: we should at some point update the logic here, to make it easier to support further models
258        # without cluttering the UI.
259        excluded_models = ["vesicles_2d_maus"]
260        model_list = [name for name in model_list if name not in excluded_models]
261
262        models = ["- choose -"] + model_list
263        self.model_selector = QComboBox()
264        self.model_selector.addItems(models)
265        # Create a layout and add the title label and combo box
266        layout = QVBoxLayout()
267        layout.addWidget(title_label)
268        layout.addWidget(self.model_selector)
269
270        # Set layout on the model widget
271        model_widget.setLayout(layout)
272        return model_widget
273
274    def on_predict(self):
275        # Get the model and postprocessing settings.
276        model_type = self.model_selector.currentText()
277        custom_model_path = self.checkpoint_param.text()
278        if model_type == "- choose -":
279            show_info("INFO: Please choose a model.")
280            return
281
282        device = get_device(self.device_dropdown.currentText())
283
284        # Get the image data.
285        image = self._get_layer_selector_data(self.image_selector_name)
286        if image is None:
287            show_info("INFO: Please choose an image.")
288            return
289
290        with self._computing(
291            self.predict_button, "Computing…", "Run Segmentation", "INFO: Running segmentation…"
292        ):
293            # Load the model. Override if user chose custom model.
294            rescale_input = True
295            if custom_model_path:
296                model = _load_custom_model(custom_model_path, device)
297                rescale_input = False
298                if model:
299                    show_info(f"INFO: Using custom model from path: {custom_model_path}")
300                else:
301                    show_info(f"ERROR: Failed to load custom model from path: {custom_model_path}")
302                    return
303            else:
304                model = get_model(model_type, device)
305
306            # Get the current tiling.
307            self.tiling = _get_current_tiling(self.tiling, self.default_tiling, model_type)
308
309            # Get the voxel size.
310            metadata = self._get_layer_selector_data(self.image_selector_name, return_metadata=True)
311            voxel_size = self._handle_resolution(metadata, self.voxel_size_param, image.ndim, return_as_list=False)
312
313            # Determine the scaling based on the voxel size.
314            scale = None
315            if voxel_size and rescale_input:
316                # Calculate scale so voxel_size is the same as in training.
317                scale = compute_scale_from_voxel_size(voxel_size, model_type)
318                scale_info = list(map(lambda x: np.round(x, 2), scale))
319                show_info(f"INFO: Rescaled the image by {scale_info} to optimize for the selected model.")
320
321            # Some models require an additional segmentation for inference or postprocessing.
322            # For these models we read out the 'Extra Segmentation' widget.
323            if model_type == "ribbon":  # Currently only the ribbon model needs the extra seg.
324                extra_seg = self._get_layer_selector_data(self.extra_seg_selector_name)
325                resolution = tuple(voxel_size[ax] for ax in "zyx")
326                kwargs = {"extra_segmentation": extra_seg, "resolution": resolution}
327            elif "cristae" in model_type:  # Cristae model expects 2 3D volumes
328                kwargs = {
329                    "extra_segmentation": self._get_layer_selector_data(self.extra_seg_selector_name),
330                    "with_channels": True,
331                    "channels_to_standardize": [0]
332                }
333            else:
334                kwargs = {}
335            kwargs.update(self._get_postprocessing_kwargs())
336            segmentation = run_segmentation(
337                image, model=model, model_type=model_type, tiling=self.tiling, scale=scale, **kwargs
338            )
339
340            # Add the segmentation layer(s).
341            if isinstance(segmentation, dict):
342                for name, seg in segmentation.items():
343                    self.viewer.add_labels(seg, name=name, metadata=metadata)
344            else:
345                self.viewer.add_labels(segmentation, name=f"{model_type}", metadata=metadata)
346
347        show_info(f"INFO: Segmentation of {model_type} added to layers.")
348
349    def _create_settings_widget(self):
350        setting_values = QWidget()
351        # setting_values.setToolTip(get_tooltip("embedding", "settings"))
352        setting_values.setLayout(QVBoxLayout())
353
354        # Create UI for the device.
355        device = "auto"
356        device_options = ["auto"] + _available_devices()
357
358        self.device_dropdown, layout = self._add_choice_param("device", device, device_options)
359        setting_values.layout().addLayout(layout)
360
361        # Create UI for the tile shape.
362        self.default_tiling = get_default_tiling()
363        self.tiling = copy.deepcopy(self.default_tiling)
364        self.tiling["tile"]["x"], self.tiling["tile"]["y"], self.tiling["tile"]["z"], layout = self._add_shape_param(
365            ("tile_x", "tile_y", "tile_z"),
366            (self.default_tiling["tile"]["x"], self.default_tiling["tile"]["y"], self.default_tiling["tile"]["z"]),
367            min_val=0, max_val=2048, step=16,
368            # tooltip=get_tooltip("embedding", "tiling")
369        )
370        setting_values.layout().addLayout(layout)
371
372        # Create UI for the halo.
373        self.tiling["halo"]["x"], self.tiling["halo"]["y"], self.tiling["halo"]["z"], layout = self._add_shape_param(
374            ("halo_x", "halo_y", "halo_z"),
375            (self.default_tiling["halo"]["x"], self.default_tiling["halo"]["y"], self.default_tiling["halo"]["z"]),
376            min_val=0, max_val=512,
377            # tooltip=get_tooltip("embedding", "halo")
378        )
379        setting_values.layout().addLayout(layout)
380
381        # Read voxel size from layer metadata.
382        self.voxel_size_param, layout = self._add_float_param(
383            "voxel_size", 0.0, min_val=0.0, max_val=100.0,
384        )
385        setting_values.layout().addLayout(layout)
386
387        self.checkpoint_param, layout = self._add_string_param(
388            name="checkpoint", value="", title="Load Custom Model",
389            placeholder="path/to/checkpoint.pt",
390        )
391        setting_values.layout().addLayout(layout)
392
393        # Add selection UI for additional segmentation, which some models require for inference or postproc.
394        self.extra_seg_selector_name = "Extra Segmentation"
395        self.extra_selector_widget = self._create_layer_selector(self.extra_seg_selector_name, layer_type="Labels")
396        setting_values.layout().addWidget(self.extra_selector_widget)
397
398        # Add model-specific post-processing settings that are updated when the selected model changes.
399        setting_values.layout().addWidget(QLabel("Post-processing:"))
400        self.postprocessing_settings_widget = QWidget()
401        self.postprocessing_settings_layout = QVBoxLayout()
402        self.postprocessing_settings_layout.setContentsMargins(0, 0, 0, 0)
403        self.postprocessing_settings_widget.setLayout(self.postprocessing_settings_layout)
404        setting_values.layout().addWidget(self.postprocessing_settings_widget)
405        self.postprocessing_parameter_widgets = {}
406        self.model_selector.currentTextChanged.connect(self._update_postprocessing_settings)
407        self._update_postprocessing_settings(self.model_selector.currentText())
408
409        settings = self._make_collapsible(widget=setting_values, title="Advanced Settings")
410        return settings