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