added custom model path with target selector to choose between desired segmentation functions

Author: lufre1

synaptic_reconstruction/tools/segmentation_widget.py

@@ -4,7 +4,7 @@
 
 from .base_widget import BaseWidget
 from .util import (run_segmentation, get_model, get_model_registry, _available_devices, get_device,
-                   get_current_tiling, compute_scale_from_voxel_size)
+                   get_current_tiling, compute_scale_from_voxel_size, load_custom_model)
 from synaptic_reconstruction.inference.util import get_default_tiling
 import copy
 
@@ -54,27 +54,34 @@ def load_model_widget(self):
     def on_predict(self):
         # Get the model and postprocessing settings.
         model_type = self.model_selector.currentText()
-        if model_type == "- choose -":
-            show_info("Please choose a model.")
+        custom_model_path = self.checkpoint_param.text()
+        if model_type == "- choose -" and custom_model_path is None:
+            show_info("INFO: Please choose a model.")
             return
 
-        # Load the model.
         device = get_device(self.device_dropdown.currentText())
-        model = get_model(model_type, device)
+
+        # Load the model. Override if user chose custom model
+        if custom_model_path:
+            model = load_custom_model(custom_model_path, device)
+            if model:
+                show_info(f"INFO: Using custom model from path: {custom_model_path}")
+                model_type = "custom"
+            else:
+                show_info(f"ERROR: Failed to load custom model from path: {custom_model_path}")
+                return
+        else:
+            model = get_model(model_type, 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.")
+            show_info("INFO: Please choose an image.")
             return
 
         # load current tiling
         self.tiling = get_current_tiling(self.tiling, self.default_tiling, model_type)
 
-        # TODO: Use scale derived from the image resolution.
-        # get avg image shape from training of the selected model
-        # wichmann data avg voxel size = 17.53
-
         metadata = self._get_layer_selector_data(self.image_selector_name, return_metadata=True)
         voxel_size = metadata.get("voxel_size", None)
         scale = None
@@ -90,17 +97,24 @@ def on_predict(self):
                 voxel_size["x"] = self.voxel_size_param.value()
                 voxel_size["y"] = self.voxel_size_param.value()
         if voxel_size:
-            # calculate scale so voxel_size is the same as in training
-            scale = compute_scale_from_voxel_size(voxel_size, model_type)
-            show_info(f"Rescaled the image by {scale} to optimize for the selected model.")
+            if model_type == "custom":
+                show_info(f"INFO: Scaling image not avaialable for custom model.")
+            else:
+                # calculate scale so voxel_size is the same as in training
+                scale = compute_scale_from_voxel_size(voxel_size, model_type)
+                show_info(f"INFO: Rescaled the image by {scale} to optimize for the selected model.")
+
+        if model_type == "custom":
+            # choose appropriate segmentation
+            model_type = self.model_target_selector.currentText()
 
         segmentation = run_segmentation(
             image, model=model, model_type=model_type, tiling=self.tiling, scale=scale
         )
 
         # Add the segmentation layer
         self.viewer.add_labels(segmentation, name=f"{model_type}-segmentation", metadata=metadata)
-        show_info(f"Segmentation of {model_type} added to layers.")
+        show_info(f"INFO: Segmentation of {model_type} added to layers.")
 
     def _create_settings_widget(self):
         setting_values = QWidget()
@@ -141,5 +155,18 @@ def _create_settings_widget(self):
         )
         setting_values.layout().addLayout(layout)
 
+        self.checkpoint_param, layout = self._add_string_param(
+            name="checkpoint", value="", title="Load Custom Model",
+            placeholder="path/to/checkpoint.pt",
+        )
+        setting_values.layout().addLayout(layout)
+        
+        self.model_target_selector, layout = self._add_choice_param(
+            name="Model Target", value="vesicles", options=[
+                "vesicles", "mitochondria", "active_zone", "compartments", "inner_ear_structures"
+                ]
+        )
+        setting_values.layout().addLayout(layout)
+
         settings = self._make_collapsible(widget=setting_values, title="Advanced Settings")
         return settings

synaptic_reconstruction/tools/util.py

@@ -1,4 +1,5 @@
 import os
+import re
 from typing import Dict, List, Optional, Union
 
 import torch
@@ -10,6 +11,24 @@
 from ..inference.mitochondria import segment_mitochondria
 
 
+def load_custom_model(model_path: str, device: Optional[Union[str, torch.device]] = None) -> torch.nn.Module:
+    model_path = _clean_filepath(model_path)
+    if device is None:
+        device = get_device(device)
+    try:
+        warnings.filterwarnings(
+            "ignore",
+            message="You are using `torch.load` with `weights_only=False`",
+            category=FutureWarning
+        )
+        model = torch.load(model_path, map_location=torch.device(device))
+    except Exception as e:
+        print(e)
+        print("model path", model_path)
+        return None
+    return model
+
+
 def get_model(model_type: str, device: Optional[Union[str, torch.device]] = None) -> torch.nn.Module:
     """Get the model for the given segmentation type.
 
@@ -21,7 +40,8 @@ def get_model(model_type: str, device: Optional[Union[str, torch.device]] = None
     Returns:
         The model.
     """
-    device = get_device(device)
+    if device is None:
+        device = get_device(device)
     model_registry = get_model_registry()
     model_path = model_registry.fetch(model_type)
     warnings.filterwarnings(
@@ -221,3 +241,27 @@ def compute_scale_from_voxel_size(
             voxel_size["z"] / training_voxel_size["z"]
         )
     return scale
+
+
+def _clean_filepath(filepath):
+    """
+    Cleans a given filepath by:
+    - Removing newline characters (\n)
+    - Removing escape sequences
+    - Stripping the 'file://' prefix if present
+
+    Args:
+        filepath (str): The original filepath
+
+    Returns:
+        str: The cleaned filepath
+    """
+    # Remove 'file://' prefix if present
+    if filepath.startswith("file://"):
+        filepath = filepath[7:]
+    
+    # Remove escape sequences and newlines
+    filepath = re.sub(r'\\.', '', filepath)
+    filepath = filepath.replace('\n', '').replace('\r', '')
+
+    return filepath