Merge pull request #60 from computational-cell-analytics/56-add-support-for-custom-model-paths-in-segmentation-plugin

added custom model path with target selector to choose between desire…

Author: constantinpape

synaptic_reconstruction/tools/cli.py

@@ -6,6 +6,7 @@
 
 
 # TODO: handle kwargs
+# TODO: add custom model path
 def segmentation_cli():
     parser = argparse.ArgumentParser(description="Run segmentation.")
     parser.add_argument(
@@ -16,6 +17,7 @@ def segmentation_cli():
         "--output_path", "-o", required=True,
         help="The filepath to directory where the segmentations will be saved."
     )
+    # TODO: list the availabel models here by parsing the keys of the model registry
     parser.add_argument(
         "--model", "-m", required=True, help="The model type."
     )

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,20 @@ 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("INFO: The image is not rescaled for a 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.")
 
         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 +151,11 @@ 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)
+
         settings = self._make_collapsible(widget=setting_values, title="Advanced Settings")
         return settings

synaptic_reconstruction/tools/util.py

@@ -1,15 +1,28 @@
 import os
+import re
 from typing import Dict, List, Optional, Union
 
 import torch
 import numpy as np
 import pooch
-import warnings
 
 from ..inference.vesicles import segment_vesicles
 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:
+        model = torch.load(model_path, map_location=torch.device(device), weights_only=False)
+    except Exception as e:
+        print(e)
+        print("model path", model_path)
+        return None
+    return model
+
+
 def get_model_path(model_type: str) -> str:
     """Get the local path to a given model.
 
@@ -35,19 +48,14 @@ 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_path = get_model_path(model_type)
-    warnings.filterwarnings(
-        "ignore",
-        message="You are using `torch.load` with `weights_only=False`",
-        category=FutureWarning
-    )
-    model = torch.load(model_path)
+    model = torch.load(model_path, weights_only=False)
     model.to(device)
     return model
 
 
-# TODO: distinguish between 2d and 3d vesicle model segmentation
 def run_segmentation(
     image: np.ndarray,
     model: torch.nn.Module,
@@ -60,12 +68,15 @@ def run_segmentation(
     """Run synaptic structure segmentation.
 
     Args:
-        image: ...
-        model: ...
-        model_type: ...
-        tiling: ...
-        scale: ...
-        verbose: ...
+        image: The input image or image volume.
+        model: The segmentation model.
+        model_type: The model type. This will determine which segmentation
+            post-processing is used.
+        tiling: The tiling settings for inference.
+        scale: A scale factor for resizing the input before applying the model.
+            The output will be scaled back to the initial size.
+        verbose: Whether to print detailed information about the prediction and segmentation.
+        kwargs: Optional parameter for the segmentation function.
 
     Returns:
         The segmentation.
@@ -234,3 +245,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