Merge pull request #55 from computational-cell-analytics/53-next-steps-plugin-implementation

53 next steps plugin implementation

Author: constantinpape

synaptic_reconstruction/tools/base_widget.py

@@ -61,7 +61,7 @@ def _update_selector(self, selector, layer_filter):
         image_layers = [layer.name for layer in self.viewer.layers if isinstance(layer, layer_filter)]  # if isinstance(layer, napari.layers.Image)
         selector.addItems(image_layers)
 
-    def _get_layer_selector_data(self, selector_name):
+    def _get_layer_selector_data(self, selector_name, return_metadata=False):
         """Return the data for the layer currently selected in a given selector."""
         if selector_name in self.layer_selectors:
             selector_widget = self.layer_selectors[selector_name]
@@ -72,7 +72,10 @@ def _get_layer_selector_data(self, selector_name):
             if isinstance(image_selector, QComboBox):
                 selected_layer_name = image_selector.currentText()
                 if selected_layer_name in self.viewer.layers:
-                    return self.viewer.layers[selected_layer_name].data
+                    if return_metadata:
+                        return self.viewer.layers[selected_layer_name].metadata
+                    else:
+                        return self.viewer.layers[selected_layer_name].data
         return None  # Return None if layer not found
 
     def _add_string_param(self, name, value, title=None, placeholder=None, layout=None, tooltip=None):
@@ -169,6 +172,15 @@ def _add_shape_param(self, names, values, min_val, max_val, step=1, title=None,
             title=title[1] if title is not None else title, tooltip=tooltip
         )
         layout.addLayout(y_layout)
+        
+        if len(names) == 3:
+            z_layout = QVBoxLayout()
+            z_param, _ = self._add_int_param(
+                names[2], values[2], min_val=min_val, max_val=max_val, layout=z_layout, step=step,
+                title=title[2] if title is not None else title, tooltip=tooltip
+            )
+            layout.addLayout(z_layout)
+            return x_param, y_param, z_param, layout
 
         return x_param, y_param, layout
 

synaptic_reconstruction/tools/distance_measure_widget.py

@@ -101,33 +101,53 @@ def _add_lines_and_table(self, lines, properties, table_data, name):
     def on_measure_seg_to_object(self):
         segmentation = self._get_layer_selector_data(self.image_selector_name1)
         object_data = self._get_layer_selector_data(self.image_selector_name2)
+        # get metadata from layer if available
+        metadata = self._get_layer_selector_data(self.image_selector_name1, return_metadata=True)
+        resolution = metadata.get("voxel_size", None)
+        if resolution is not None:
+            resolution = [v for v in resolution.values()]
+        # if user input is present override metadata
+        if self.voxel_size_param.value() != 0.0:  # changed from default
+            resolution = segmentation.ndim * [self.voxel_size_param.value()]
 
         (distances,
          endpoints1,
          endpoints2,
          seg_ids) = distance_measurements.measure_segmentation_to_object_distances(
             segmentation=segmentation, segmented_object=object_data, distance_type="boundary",
+            resolution=resolution
         )
         lines, properties = distance_measurements.create_object_distance_lines(
             distances=distances,
             endpoints1=endpoints1,
             endpoints2=endpoints2,
-            seg_ids=seg_ids
+            seg_ids=seg_ids,
         )
         table_data = self._to_table_data(distances, seg_ids, endpoints1, endpoints2)
         self._add_lines_and_table(lines, properties, table_data, name="distances")
 
     def on_measure_pairwise(self):
         segmentation = self._get_layer_selector_data(self.image_selector_name1)
+        if segmentation is None:
+            show_info("Please choose a segmentation.")
+            return
+        # get metadata from layer if available
+        metadata = self._get_layer_selector_data(self.image_selector_name1, return_metadata=True)
+        resolution = metadata.get("voxel_size", None)
+        if resolution is not None:
+            resolution = [v for v in resolution.values()]
+        # if user input is present override metadata
+        if self.voxel_size_param.value() != 0.0:  # changed from default
+            resolution = segmentation.ndim * [self.voxel_size_param.value()]
 
         (distances,
          endpoints1,
          endpoints2,
          seg_ids) = distance_measurements.measure_pairwise_object_distances(
-            segmentation=segmentation, distance_type="boundary"
+            segmentation=segmentation, distance_type="boundary", resolution=resolution
         )
         lines, properties = distance_measurements.create_pairwise_distance_lines(
-            distances=distances, endpoints1=endpoints1, endpoints2=endpoints2, seg_ids=seg_ids.tolist()
+            distances=distances, endpoints1=endpoints1, endpoints2=endpoints2, seg_ids=seg_ids.tolist(),
         )
         table_data = self._to_table_data(
             distances=properties["distance"],
@@ -142,5 +162,10 @@ def _create_settings_widget(self):
         self.save_path, layout = self._add_path_param(name="Save Table", select_type="file", value="")
         setting_values.layout().addLayout(layout)
 
+        self.voxel_size_param, layout = self._add_float_param(
+            "voxel_size", 0.0, min_val=0.0, max_val=100.0,
+        )
+        setting_values.layout().addLayout(layout)
+
         settings = self._make_collapsible(widget=setting_values, title="Advanced Settings")
         return settings

synaptic_reconstruction/tools/segmentation_widget.py

@@ -1,10 +1,12 @@
 import napari
-import napari.layers
 from napari.utils.notifications import show_info
 from qtpy.QtWidgets import QWidget, QVBoxLayout, QPushButton, QLabel, QComboBox
 
 from .base_widget import BaseWidget
-from .util import run_segmentation, get_model, get_model_registry, _available_devices
+from .util import (run_segmentation, get_model, get_model_registry, _available_devices, get_device,
+                   get_current_tiling, compute_scale_from_voxel_size)
+from synaptic_reconstruction.inference.util import get_default_tiling
+import copy
 
 
 class SegmentationWidget(BaseWidget):
@@ -13,6 +15,7 @@ def __init__(self):
 
         self.viewer = napari.current_viewer()
         layout = QVBoxLayout()
+        self.tiling = {}
 
         # Create the image selection dropdown.
         self.image_selector_name = "Image data"
@@ -56,36 +59,47 @@ def on_predict(self):
             return
 
         # Load the model.
-        model = get_model(model_type, self.device)
+        device = get_device(self.device_dropdown.currentText())
+        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.")
             return
 
-        # get tile shape and halo from the viewer
-        tiling = {
-            "tile": {
-                "x": self.tile_x_param.value(),
-                "y": self.tile_y_param.value(),
-                "z": 1
-            },
-            "halo": {
-                "x": self.halo_x_param.value(),
-                "y": self.halo_y_param.value(),
-                "z": 1
-            }
-        }
+        # load current tiling
+        self.tiling = get_current_tiling(self.tiling, self.default_tiling, model_type)
 
         # TODO: Use scale derived from the image resolution.
-        scale = [self.scale_param.value()]
+        # 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
+
+        if self.voxel_size_param.value() != 0.0:  # changed from default
+            voxel_size = {}
+            # override voxel size with user input
+            if len(image.shape) == 3:
+                voxel_size["x"] = self.voxel_size_param.value()
+                voxel_size["y"] = self.voxel_size_param.value()
+                voxel_size["z"] = self.voxel_size_param.value()
+            else:
+                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.")
+
         segmentation = run_segmentation(
-            image, model=model, model_type=model_type, tiling=tiling, scale=scale
+            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")
+        self.viewer.add_labels(segmentation, name=f"{model_type}-segmentation", metadata=metadata)
         show_info(f"Segmentation of {model_type} added to layers.")
 
     def _create_settings_widget(self):
@@ -94,31 +108,36 @@ def _create_settings_widget(self):
         setting_values.setLayout(QVBoxLayout())
 
         # Create UI for the device.
-        self.device = "auto"
+        device = "auto"
         device_options = ["auto"] + _available_devices()
 
-        self.device_dropdown, layout = self._add_choice_param("device", self.device, device_options)
+        self.device_dropdown, layout = self._add_choice_param("device", device, device_options)
         setting_values.layout().addLayout(layout)
 
         # Create UI for the tile shape.
-        # TODO: make the tiling 3d and get the default values from 'inference'
-        self.tile_x, self.tile_y = 512, 512  # 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,
+        self.default_tiling = get_default_tiling()
+        self.tiling = copy.deepcopy(self.default_tiling)
+        self.tiling["tile"]["x"], self.tiling["tile"]["y"], self.tiling["tile"]["z"], layout = self._add_shape_param(
+            ("tile_x", "tile_y", "tile_z"),
+            (self.default_tiling["tile"]["x"], self.default_tiling["tile"]["y"], self.default_tiling["tile"]["z"]),
+            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 = 64, 64  # 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,
+
+        self.tiling["halo"]["x"], self.tiling["halo"]["y"], self.tiling["halo"]["z"], layout = self._add_shape_param(
+            ("halo_x", "halo_y", "halo_z"),
+            (self.default_tiling["halo"]["x"], self.default_tiling["halo"]["y"], self.default_tiling["halo"]["z"]),
+            min_val=0, max_val=512,
             # tooltip=get_tooltip("embedding", "halo")
         )
         setting_values.layout().addLayout(layout)
 
-        self.scale_param, layout = self._add_float_param(
-            "scale", 0.5, min_val=0.0, max_val=8.0,
+        # read voxel size from layer metadata
+        self.voxel_size_param, layout = self._add_float_param(
+            "voxel_size", 0.0, min_val=0.0, max_val=100.0,
         )
         setting_values.layout().addLayout(layout)
 

synaptic_reconstruction/tools/util.py

@@ -4,6 +4,7 @@
 import torch
 import numpy as np
 import pooch
+import warnings
 
 from ..inference.vesicles import segment_vesicles
 from ..inference.mitochondria import segment_mitochondria
@@ -23,6 +24,11 @@ def get_model(model_type: str, device: Optional[Union[str, torch.device]] = None
     device = get_device(device)
     model_registry = get_model_registry()
     model_path = model_registry.fetch(model_type)
+    warnings.filterwarnings(
+        "ignore",
+        message="You are using `torch.load` with `weights_only=False`",
+        category=FutureWarning
+    )
     model = torch.load(model_path)
     model.to(device)
     return model
@@ -73,12 +79,12 @@ def get_cache_dir():
 
 def get_model_training_resolution(model_type):
     resolutions = {
-        "active_zone": 1.44,
-        "compartments": 3.47,
-        "mitochondria": 1.0,  # FIXME: this is a dummy value, we need to determine the real one
-        "vesicles_2d": 1.35,
-        "vesicles_3d": 1.35,
-        "vesicles_cryo": 0.88,
+        "active_zone": {"x": 1.44, "y": 1.44, "z": 1.44},
+        "compartments": {"x": 3.47, "y": 3.47, "z": 3.47},
+        "mitochondria": {"x": 1.0, "y": 1.0, "z": 1.0},  # FIXME: this is a dummy value, we need to determine the real one
+        "vesicles_2d": {"x": 1.35, "y": 1.35},
+        "vesicles_3d": {"x": 1.35, "y": 1.35, "z": 1.35},
+        "vesicles_cryo": {"x": 1.35, "y": 1.35, "z": 0.88},
     }
     return resolutions[model_type]
 
@@ -168,3 +174,50 @@ def _available_devices():
         else:
             available_devices.append(device)
     return available_devices
+
+
+def get_current_tiling(tiling: dict, default_tiling: dict, model_type: str):
+    # get tiling values from qt objects
+    for k, v in tiling.items():
+        for k2, v2 in v.items():
+            if isinstance(v2, int):
+                continue
+            tiling[k][k2] = v2.value()
+    # check if user inputs tiling/halo or not
+    if default_tiling == tiling:
+        if "2d" in model_type:
+            # if its a 2d model expand x,y and set z to 1
+            tiling = {
+                "tile": {
+                    "x": 512,
+                    "y": 512,
+                    "z": 1
+                },
+                "halo": {
+                    "x": 64,
+                    "y": 64,
+                    "z": 1
+                }
+            }
+    elif "2d" in model_type:
+        # if its a 2d model set z to 1
+        tiling["tile"]["z"] = 1
+        tiling["halo"]["z"] = 1
+
+    return tiling
+
+
+def compute_scale_from_voxel_size(
+    voxel_size: dict,
+    model_type: str
+) -> List[float]:
+    training_voxel_size = get_model_training_resolution(model_type)
+    scale = [
+        voxel_size["x"] / training_voxel_size["x"],
+        voxel_size["y"] / training_voxel_size["y"],
+    ]
+    if len(voxel_size) == 3 and len(training_voxel_size) == 3:
+        scale.append(
+            voxel_size["z"] / training_voxel_size["z"]
+        )
+    return scale

synaptic_reconstruction/tools/volume_reader.py

@@ -1,6 +1,7 @@
 import os
 
 from typing import Callable, List, Optional, Sequence, Union
+import mrcfile
 from napari.types import LayerData
 
 from elf.io import open_file, is_dataset
@@ -22,7 +23,17 @@ def get_reader(path: PathOrPaths) -> Optional[ReaderFunction]:
 def _read_mrc(path, fname):
     with open_file(path, mode="r") as f:
         data = f["data"][:]
-    layer_attributes = {"name": fname, "colormap": "gray"}
+    voxel_size = read_voxel_size(path)
+    metadata = {
+        "file_path": path,
+        "voxel_size": voxel_size
+    }
+    layer_attributes = {
+        "name": fname,
+        "colormap": "gray",
+        "metadata": metadata
+    }
+
     return [(data, layer_attributes)]
 
 
@@ -61,3 +72,26 @@ def read_image_volume(path: PathOrPaths) -> List[LayerData]:
     except Exception as e:
         print(f"Failed to read file: {e}")
         return
+
+
+def read_voxel_size(input_path: str) -> dict | None:
+    """Read voxel size from mrc/rec file and store it in layer_attributes.
+    The original unit of voxel size is Angstrom and we convert it to nanometers
+    by dividing it by ten.
+
+    Args:
+        input_path (str): path to mrc/rec file
+        layer_attributes (dict): napari layer attributes to store voxel size to
+    """
+    new_voxel_size = None
+    with mrcfile.open(input_path, permissive=True) as mrc:
+        try:
+            voxel_size = mrc.voxel_size
+            new_voxel_size = {
+                "x": voxel_size.x / 10,
+                "y": voxel_size.y / 10,
+                "z": voxel_size.z / 10,
+            }
+        except Exception as e:
+            print(f"Failed to read voxel size: {e}")
+    return new_voxel_size