Update distance measurement widget

constantinpape · constantinpape · commit d991be36e864 · 2024-11-16T17:41:57.000+01:00
diff --git a/setup.py b/setup.py
@@ -8,13 +8,12 @@
     name="synaptic_reconstruction",
     packages=find_packages(exclude=["test"]),
     version=__version__,
-    author="Constantin Pape; Sarah Muth",
+    author="Constantin Pape; Sarah Muth; Luca Freckmann",
     url="https://github.com/computational-cell-analytics/synaptic_reconstruction",
     license="MIT",
     entry_points={
         "console_scripts": [
-            "sr_tools.correct_segmentation = synaptic_reconstruction.tools.segmentation_correction:main",
-            "sr_tools.measure_distances = synaptic_reconstruction.tools.distance_measurement:main",
+            # TODO add segmentation CLI
         ],
         "napari.manifest": [
             "synaptic_reconstruction = synaptic_reconstruction:napari.yaml",
diff --git a/synaptic_reconstruction/distance_measurements.py b/synaptic_reconstruction/distance_measurements.py
@@ -1,6 +1,6 @@
 import os
 import multiprocessing as mp
-from typing import Dict, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
 import numpy as np
 
@@ -61,8 +61,9 @@ def _compute_boundary_distances(segmentation, resolution, n_threads):
     n = len(seg_ids)
 
     pairwise_distances = np.zeros((n, n))
-    end_points1 = np.zeros((n, n, 3), dtype="int")
-    end_points2 = np.zeros((n, n, 3), dtype="int")
+    ndim = segmentation.ndim
+    end_points1 = np.zeros((n, n, ndim), dtype="int")
+    end_points2 = np.zeros((n, n, ndim), dtype="int")
 
     properties = regionprops(segmentation)
     properties = {prop.label: prop for prop in properties}
@@ -80,8 +81,11 @@ def compute_distances_for_object(i):
             prop = properties[ngb_id]
 
             bb = prop.bbox
-            offset = np.array(bb[:3])
-            bb = np.s_[bb[0]:bb[3], bb[1]:bb[4], bb[2]:bb[5]]
+            offset = np.array(bb[:ndim])
+            if ndim == 2:
+                bb = np.s_[bb[0]:bb[2], bb[1]:bb[3]]
+            else:
+                bb = np.s_[bb[0]:bb[3], bb[1]:bb[4], bb[2]:bb[5]]
 
             mask = segmentation[bb] == ngb_id
             ngb_dist, ngb_index = distances[bb].copy(), indices[(slice(None),) + bb]
@@ -168,10 +172,9 @@ def _compute_seg_object_distances(segmentation, segmented_object, resolution, ve
     for prop in tqdm(props, disable=not verbose):
         bb = prop.bbox
         offset = np.array(bb[:ndim])
-        # bb = np.s_[bb[0]:bb[3], bb[1]:bb[4], bb[2]:bb[5]]
-        if len(bb) == 4:  # 2D bounding box
+        if ndim == 2:
             bb = np.s_[bb[0]:bb[2], bb[1]:bb[3]]
-        elif len(bb) == 6:  # 3D bounding box
+        else:
             bb = np.s_[bb[0]:bb[3], bb[1]:bb[4], bb[2]:bb[5]]
 
         label = prop.label
@@ -296,7 +299,7 @@ def create_pairwise_distance_lines(
     distances: np.ndarray,
     endpoints1: np.ndarray,
     endpoints2: np.ndarray,
-    seg_ids: np.ndarray,
+    seg_ids: List[List[int]],
     n_neighbors: Optional[int] = None,
     pairs: Optional[np.ndarray] = None,
     bb: Optional[Tuple[slice]] = None,
@@ -323,9 +326,7 @@ def create_pairwise_distance_lines(
     if pairs is None and n_neighbors is not None:
         pairs = _extract_nearest_neighbors(distances, seg_ids, n_neighbors, remove_duplicates=remove_duplicates)
     elif pairs is None:
-        pairs = [
-            [id1, id2] for id1 in seg_ids for id2 in seg_ids if id1 < id2
-        ]
+        pairs = [[id1, id2] for id1 in seg_ids for id2 in seg_ids if id1 < id2]
 
     assert pairs is not None
     pair_indices = (
diff --git a/synaptic_reconstruction/tools/distance_measure_widget.py b/synaptic_reconstruction/tools/distance_measure_widget.py
@@ -2,6 +2,7 @@
 
 import napari
 import napari.layers
+import numpy as np
 import pandas as pd
 
 from napari.utils.notifications import show_info
@@ -39,109 +40,103 @@ def __init__(self):
         self.viewer = napari.current_viewer()
         layout = QVBoxLayout()
 
-        self.image_selector_name1 = "Segmentation 1"
-        self.image_selector_name2 = "Segmentation 2"
-        # Create the image selection dropdown
-        # TODO: update the names to make it easier to distinguish what is what.
+        self.image_selector_name1 = "Segmentation"
+        self.image_selector_name2 = "Object"
+        # Create the image selection dropdown.
         self.segmentation1_selector_widget = self._create_layer_selector(self.image_selector_name1, layer_type="Labels")
         self.segmentation2_selector_widget = self._create_layer_selector(self.image_selector_name2, layer_type="Labels")
 
-        # create save path
+        # Create advanced settings.
         self.settings = self._create_settings_widget()
 
-        # create buttons
-        self.measure_pairwise_button = QPushButton("Measure Distance Pairwise")
-        self.measure_segmentation_to_object_button = QPushButton("Measure Distance Segmentation to Object")
+        # Create and connect buttons.
+        self.measure_button1 = QPushButton("Measure Distances")
+        self.measure_button1.clicked.connect(self.on_measure_seg_to_object)
 
-        # Connect buttons to functions
-        self.measure_pairwise_button.clicked.connect(self.on_measure_pairwise)
-        self.measure_segmentation_to_object_button.clicked.connect(self.on_measure_segmentation_to_object)
-        # self.load_model_button.clicked.connect(self.on_load_model)
+        self.measure_button2 = QPushButton("Measure Pairwise Distances")
+        self.measure_button2.clicked.connect(self.on_measure_pairwise)
 
-        # Add the widgets to the layout
+        # Add the widgets to the layout.
         layout.addWidget(self.segmentation1_selector_widget)
         layout.addWidget(self.segmentation2_selector_widget)
         layout.addWidget(self.settings)
-        # layout.addWidget(self.measure_pairwise_button)
-        layout.addWidget(self.measure_segmentation_to_object_button)
+        layout.addWidget(self.measure_button1)
+        layout.addWidget(self.measure_button2)
 
         self.setLayout(layout)
 
-    def on_measure_segmentation_to_object(self):
-        segmentation1_data = self._get_layer_selector_data(self.image_selector_name1)
-        segmentation2_data = self._get_layer_selector_data(self.image_selector_name2)
-        if segmentation1_data is None or segmentation2_data is None:
-            show_info("Please choose both segmentation layers.")
-            return
-
-        (distances,
-         endpoints1,
-         endpoints2,
-         seg_ids) = distance_measurements.measure_segmentation_to_object_distances(
-            segmentation=segmentation1_data,
-            segmented_object=segmentation2_data,
-            distance_type="boundary",
-        )
-
-        if self.save_path.text() != "":
-            data = {"label": seg_ids, "distance": distances}
+    def _to_table_data(self, distances, seg_ids, endpoints1=None, endpoints2=None):
+        assert len(distances) == len(seg_ids), f"{distances.shape}, {seg_ids.shape}"
+        if seg_ids.ndim == 2:
+            table_data = {"label1": seg_ids[:, 0], "label2": seg_ids[:, 1], "distance": distances}
+        else:
+            table_data = {"label": seg_ids, "distance": distances}
+        if endpoints1 is not None:
             axis_names = "zyx" if endpoints1.shape[1] == 3 else "yx"
-            data.update({f"begin-{ax}": endpoints1[:, i] for i, ax in enumerate(axis_names)})
-            data.update({f"end-{ax}": endpoints2[:, i] for i, ax in enumerate(axis_names)})
-            data = pd.DataFrame(data)
-            file_path = _save_distance_table(self.save_path.text(), data)
-
-        lines, properties = distance_measurements.create_object_distance_lines(
-            distances=distances,
-            endpoints1=endpoints1,
-            endpoints2=endpoints2,
-            seg_ids=seg_ids
-        )
+            table_data.update({f"begin-{ax}": endpoints1[:, i] for i, ax in enumerate(axis_names)})
+            table_data.update({f"end-{ax}": endpoints2[:, i] for i, ax in enumerate(axis_names)})
+        return pd.DataFrame(table_data)
 
-        # Add the lines layer
+    def _add_lines_and_table(self, lines, properties, table_data, name):
         line_layer = self.viewer.add_shapes(
             lines,
-            name="Distance Lines",
-            shape_type="line",  # Specify the shape type as 'line'
+            name=name,
+            shape_type="line",
             edge_width=2,
             edge_color="red",
-            blending="additive",  # Use 'additive' for blending if needed
+            blending="additive",
+            properties=properties,
         )
-
-        # FIXME: this doesn't work yet
         if add_table is not None:
             add_table(line_layer, self.viewer)
 
+        if self.save_path.text() != "":
+            file_path = _save_distance_table(self.save_path.text(), table_data)
+
         if self.save_path.text() != "":
             show_info(f"Added distance lines and saved file to {file_path}.")
         else:
             show_info("Added distance lines.")
 
+    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)
+
+        (distances,
+         endpoints1,
+         endpoints2,
+         seg_ids) = distance_measurements.measure_segmentation_to_object_distances(
+            segmentation=segmentation, segmented_object=object_data, distance_type="boundary",
+        )
+        lines, properties = distance_measurements.create_object_distance_lines(
+            distances=distances,
+            endpoints1=endpoints1,
+            endpoints2=endpoints2,
+            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):
-        if self.image is None:
-            show_info("Please choose a segmentation.")
-            return
-        if self.save_path.value() is None:
-            show_info("Please choose a save path.")
-            return
-        show_info("Not implemented yet.")
-        return
-        # distance_measurements.measure_pairwise_object_distances(
-        #     segmentation=segmentation, distance_type="boundary",
-        #     save_path=self.save_path
-        #     )
-        # lines, properties = distance_measurements.create_distance_lines(
-        #     measurement_path=self.save_path
-        # )
-
-        # # Add the lines layer
-        # self.viewer.add_lines(
-        #     lines, name="Distance Lines", visible=True, edge_width=2, edge_color="red", edge_blend="additive"
-        # )
+        segmentation = self._get_layer_selector_data(self.image_selector_name1)
+
+        (distances,
+         endpoints1,
+         endpoints2,
+         seg_ids) = distance_measurements.measure_pairwise_object_distances(
+            segmentation=segmentation, distance_type="boundary"
+        )
+        lines, properties = distance_measurements.create_pairwise_distance_lines(
+            distances=distances, endpoints1=endpoints1, endpoints2=endpoints2, seg_ids=seg_ids.tolist()
+        )
+        table_data = self._to_table_data(
+            distances=properties["distance"],
+            seg_ids=np.concatenate([properties["id_a"][:, None], properties["id_b"][:, None]], axis=1)
+        )
+        self._add_lines_and_table(lines, properties, table_data, name="pairwise-distances")
 
     def _create_settings_widget(self):
         setting_values = QWidget()
-        # setting_values.setToolTip(get_tooltip("embedding", "settings"))
         setting_values.setLayout(QVBoxLayout())
 
         self.save_path, layout = self._add_path_param(name="Save Table", select_type="file", value="")
