added dynamic switch between segmentation functions absed on the model used

Author: lufre1

environment.yaml

@@ -9,5 +9,8 @@ dependencies:
     - pyqt
     - magicgui
     - pytorch
+    - bioimageio.core
+    - kornia
+    - tensorboard
     - pip:
         - napari-skimage-regionprops

synaptic_reconstruction/tools/synaptic_plugin/segmentation.py

@@ -1,20 +1,13 @@
-from typing import TYPE_CHECKING
-import h5py
-from magicgui import magic_factory, widgets
 import napari
 import napari.layers
 from napari.utils.notifications import show_info
-from napari import Viewer
-import numpy as np
-from qtpy.QtWidgets import QWidget, QVBoxLayout, QHBoxLayout, QPushButton, QFileDialog, QLabel, QSpinBox, QLineEdit, QGroupBox, QFormLayout, QFrame, QComboBox
-from superqt import QCollapsible
-from elf.io import open_file
+from qtpy.QtWidgets import QWidget, QVBoxLayout, QPushButton, QLabel, QComboBox
+
 from .base_widget import BaseWidget
-import os
-from synaptic_reconstruction.inference.vesicles import segment_vesicles
+from synaptic_reconstruction.training.supervised_training import get_2d_model
 
 # Custom imports for model and prediction utilities
-from ..util import get_device, get_model_registry, run_prediction, _available_devices
+from ..util import run_segmentation, get_model_registry, _available_devices
 
 # if TYPE_CHECKING:
 #     import napari
@@ -28,6 +21,7 @@ def __init__(self):
         super().__init__()
 
         self.model = None
+        self.image = None
         self.viewer = napari.current_viewer()
         layout = QVBoxLayout()
 
@@ -137,18 +131,19 @@ def on_predict(self):
         if model_key == "- choose -":
             show_info("Please choose a model.")
             return
+        # loading model
 
-
         model_registry = get_model_registry()
-        model_path = model_registry.fetch(model_key)
+        model_key = self.model_selector.currentText()
+        model_path = "/home/freckmann15/.cache/synapse-net/models/vesicles"  # model_registry.fetch(model_key)
+        # model = get_2d_model(out_channels=2)
+        # model = load_model_weights(model=model, model_path=model_path)
 
         if self.image is None:
             show_info("Please choose an image.")
             return
 
         # get tile shape and halo from the viewer
-        tile_shape = (self.tile_x_param.value(), self.tile_y_param.value())
-        halo = (self.halo_x_param.value(), self.halo_y_param.value())
         tiling = {
             "tile": {
                 "x": self.tile_x_param.value(),
@@ -161,13 +156,24 @@ def on_predict(self):
                 "z": 1
                 }
             }
-        segmentation = segment_vesicles(self.image, model_path=model_path)  #tiling=tiling
+        tile_shape = (self.tile_x_param.value(), self.tile_y_param.value())
+        halo = (self.halo_x_param.value(), self.halo_y_param.value())
+        use_custom_tiling = False
+        for ts, h in zip(tile_shape, halo):
+            if ts != 0 or h != 0:  # if anything is changed from default
+                use_custom_tiling = True
+        if use_custom_tiling:
+            segmentation = run_segmentation(self.image, model_path=model_path, model_key=model_key, tiling=tiling)
+        else:
+            segmentation = run_segmentation(self.image, model_path=model_path, model_key=model_key)
+        # segmentation = np.random.randint(0, 256, size=self.image.shape, dtype=np.uint8)
+        self.viewer.add_image(segmentation, name="Segmentation", colormap="inferno", blending="additive")
         # Add predictions to Napari as separate layers
         # for i, pred in enumerate(segmentation):
         #     layer_name = f"Prediction {i+1}"
         #     self.viewer.add_image(pred, name=layer_name, colormap="inferno", blending="additive")
-        layer_kwargs = {"colormap": "inferno", "blending": "additive"}
-        return segmentation, layer_kwargs
+        # layer_kwargs = {"colormap": "inferno", "blending": "additive"}
+        # return segmentation, layer_kwargs
 
     def _create_settings_widget(self):
         setting_values = QWidget()
@@ -183,15 +189,15 @@ def _create_settings_widget(self):
         setting_values.layout().addLayout(layout)
 
         # Create UI for the tile shape.
-        self.tile_x, self.tile_y = 256, 256  # defaults
+        self.tile_x, self.tile_y = 0, 0  # defaults
         self.tile_x_param, self.tile_y_param, layout = self._add_shape_param(
             ("tile_x", "tile_y"), (self.tile_x, self.tile_y), min_val=0, max_val=2048, step=16,
             # tooltip=get_tooltip("embedding", "tiling")
         )
         setting_values.layout().addLayout(layout)
 
         # Create UI for the halo.
-        self.halo_x, self.halo_y = 32, 32  # defaults
+        self.halo_x, self.halo_y = 0, 0  # defaults
         self.halo_x_param, self.halo_y_param, layout = self._add_shape_param(
             ("halo_x", "halo_y"), (self.halo_x, self.halo_y), min_val=0, max_val=512,
             # tooltip=get_tooltip("embedding", "halo")

synaptic_reconstruction/tools/util.py

@@ -7,6 +7,46 @@
 import pooch
 import requests
 from torch_em.util.prediction import predict_with_halo
+from synaptic_reconstruction.inference.vesicles import segment_vesicles
+from synaptic_reconstruction.inference.mitochondria import segment_mitochondria
+
+
+def run_segmentation(image, model_path, model_key, tiling=None):
+    if model_key == "vesicles":
+        segmentation = segment_vesicles(image, model_path=model_path, tiling=tiling)
+    elif model_key == "mitochondria":
+        segmentation = segment_mitochondria(image, model_path=model_path, tiling=tiling)
+    return segmentation
+
+
+def load_model_weights(model, model_path):
+    # Load the entire checkpoint
+    checkpoint = torch.load(model_path, map_location=get_device())  # ["model_state"]
+    checkpoint_state_dict = checkpoint.get("state_dict", checkpoint)
+    # Filter out any non-model keys (e.g., training state, optimizer state)
+    model_state_dict = {k: v for k, v in checkpoint_state_dict.items() if k in model.state_dict()}
+    return model.load_state_dict(model_state_dict, strict=False).eval()
+    
+    # Extract the 'state_dict' from the checkpoint
+    state_dict = checkpoint.get('state_dict', checkpoint)
+    
+    # Filter out unnecessary keys from the state_dict
+    model_state_dict = model.state_dict()
+    filtered_state_dict = {k: v for k, v in state_dict.items() if k in model_state_dict}
+    
+    # Check for missing keys (optional)
+    missing_keys = set(model_state_dict.keys()) - set(filtered_state_dict.keys())
+    if missing_keys:
+        print(f"Warning: Missing keys in state_dict: {missing_keys}")
+    
+    # Check for unexpected keys (optional)
+    unexpected_keys = set(filtered_state_dict.keys()) - set(model_state_dict.keys())
+    if unexpected_keys:
+        print(f"Warning: Unexpected keys in state_dict: {unexpected_keys}")
+    
+    # Load the filtered state_dict into the model
+    model.load_state_dict(filtered_state_dict)
+    return model
 
 
 def download_and_organize_file(url, mode_type, app_name="synapse-net/models"):
@@ -30,23 +70,6 @@ def download_and_organize_file(url, mode_type, app_name="synapse-net/models"):
     os.rename(file_path, os.path.join(dir_name, mode_type, "best.pt"))
     return os.path.join(dir_name, mode_type, "best.pt")
 
-    
-    # # If file_path is a directory or has no extension, organize it into a folder with "best.pt"
-    # if os.path.isdir(file_path) or os.path.splitext(file_path)[1] == "":
-    #     # Create a directory for the file if needed
-    #     dir_path = file_path if os.path.isdir(file_path) else os.path.splitext(file_path)[0]
-    #     os.makedirs(dir_path, exist_ok=True)
-        
-    #     # Move the file into the directory with the name "best.pt"
-    #     new_file_path = os.path.join(dir_path, "best.pt")
-    #     if file_path != new_file_path:
-    #         os.rename(file_path, new_file_path)
-        
-    #     return dir_path
-    # else:
-    #     # Return the directory where the file resides
-    #     return os.path.dirname(file_path)
-
 
 def organize_file_path(path):
     # Check if path is a file or directory