Convenience

setup.py

@@ -18,6 +18,7 @@
             "synapse_net.export_to_imod_objects = synapse_net.tools.cli:imod_object_cli",
             "synapse_net.run_supervised_training = synapse_net.training.supervised_training:main",
             "synapse_net.run_domain_adaptation = synapse_net.training.domain_adaptation:main",
+            "synapse_net.visualize_vesicle_pools = synapse_net.tools.cli:pool_visualization_cli",
         ],
         "napari.manifest": [
             "synapse_net = synapse_net:napari.yaml",

synapse_net/tools/cli.py

@@ -8,6 +8,7 @@
 from ..inference.inference import _get_model_registry, get_model, get_model_training_resolution, run_segmentation
 from ..inference.scalable_segmentation import scalable_segmentation
 from ..inference.util import inference_helper, parse_tiling
+from .pool_visualization import _visualize_vesicle_pools
 
 
 def imod_point_cli():
@@ -99,6 +100,25 @@ def imod_object_cli():
     )
 
 
+def pool_visualization_cli():
+    parser = argparse.ArgumentParser(description="Load tomogram data, vesicle pools and additional segmentations for viualization.")  # noqa
+    parser.add_argument(
+        "--input_path", "-i", required=True,
+        help="The filepath to the mrc file containing the tomogram data."
+    )
+    parser.add_argument(
+        "--vesicle_path", "-v", required=True, help="The filepath to the tif file containing the vesicle segmentation."
+    )
+    parser.add_argument(
+        "--table_path", "-t", required=True, help="The filepath to the table with the vesicle pool assignments."
+    )
+    parser.add_argument(
+        "-s", "--segmentation_paths", nargs="+", help="Filepaths for additional segmentations."
+    )
+    args = parser.parse_args()
+    _visualize_vesicle_pools(args.input_path, args.vesicle_path, args.table_path, args.segmentation_paths)
+
+
 # TODO: handle kwargs
 def segmentation_cli():
     parser = argparse.ArgumentParser(description="Run segmentation.")

synapse_net/tools/pool_visualization.py

@@ -0,0 +1,55 @@
+from pathlib import Path
+
+import imageio.v3 as imageio
+import napari
+import pandas as pd
+
+from ..file_utils import read_mrc
+
+
+def _create_pools(vesicles, table):
+    label_ids, pool_colors = table.label.values, table.color.values
+
+    pools = vesicles
+    colormap = {label_id: color for label_id, color in zip(label_ids, pool_colors)}
+    colormap[None] = [0, 0, 0, 0]
+
+    return pools, colormap
+
+
+def _visualize_vesicle_pools(input_path, vesicle_path, table_path, segmentation_paths):
+    # Load the tomogram data, including scale information.
+    data, voxel_size = read_mrc(input_path)
+    axes = "zyx" if data.ndim == 3 else "yx"
+    scale = tuple(float(voxel_size[ax]) for ax in axes)
+    print("Loading data with scale", scale, "nanometer")
+
+    # Load the vesicle layer.
+    vesicles = imageio.imread(vesicle_path)
+
+    # Load the table with the pool assignments.
+    # Create and add the pool layer.
+    table = pd.read_excel(table_path)
+    pools, colormap = _create_pools(vesicles, table)
+
+    viewer = napari.Viewer()
+    viewer.add_image(data, scale=scale)
+    viewer.add_labels(vesicles, scale=scale)
+    viewer.add_labels(pools, scale=scale, name="pools", colormap=colormap)
+
+    # Add the additional segmentations.
+    if segmentation_paths is not None:
+        for seg_path in segmentation_paths:
+            name = Path(seg_path).stem
+            seg = imageio.imread(seg_path)
+            viewer.add_labels(seg, name=name, scale=scale)
+
+    # FIXME something is wrong here.
+    # Add the scale bar.
+    # @magicgui(call_button="Add Scale Bar")
+    # def add_scale_bar(v: napari.Viewer):
+    #     v.scale_bar.visible = True
+    #     v.scale_bar.unit = "nm"
+    # viewer.window.add_dock_widget(add_scale_bar)
+
+    napari.run()

synapse_net/tools/size_filter_widget.py

@@ -3,6 +3,7 @@
 import napari.viewer
 
 import numpy as np
+import nifty.tools as nt
 
 from qtpy.QtWidgets import QHBoxLayout, QVBoxLayout, QPushButton, QSpinBox, QLabel
 from skimage.measure import regionprops, label
@@ -34,9 +35,9 @@ def __init__(self):
         threshold_layout.addWidget(self.size_threshold_widget)
         layout.addLayout(threshold_layout)
 
-        # Add a boolean field for processing binary masks.
-        self.is_mask = False
-        layout.addWidget(self._add_boolean_param("is_mask", self.is_mask, title="Binary Mask"))
+        # Add a boolean field for applying a label operation before size filtering.
+        self.apply_label = True
+        layout.addWidget(self._add_boolean_param("apply_label", self.apply_label, title="Remove Disconnected Pieces"))
 
         # Add an optional output layer name. If not given the segmentation will be over-written.
         self.output_layer_param, _ = self._add_string_param("output_layer", "", title="Output Layer", layout=layout)
@@ -48,19 +49,36 @@ def __init__(self):
         # Add the widgets to the layout.
         self.setLayout(layout)
 
+    def _filter_mask(self, segmentation, size_threshold):
+        segmentation = label(segmentation).astype(segmentation.dtype)
+        props = regionprops(segmentation)
+        filter_ids = [prop.label for prop in props if prop.area < size_threshold]
+        segmentation[np.isin(segmentation, filter_ids)] = 0
+        segmentation = (segmentation > 0).astype(segmentation.dtype)
+        return segmentation
+
+    def _filter_segmentation(self, segmentation, size_threshold, apply_label):
+        dtype = segmentation.dtype
+        if apply_label:
+            original_segmentation = segmentation.copy()
+            segmentation = label(segmentation)
+            props = regionprops(segmentation, original_segmentation)
+        else:
+            props = regionprops(segmentation)
+        filter_ids = [prop.label for prop in props if prop.area < size_threshold]
+        segmentation[np.isin(segmentation, filter_ids)] = 0
+        if apply_label:
+            mapping = {prop.label: int(prop.max_intensity) for prop in props if prop.label not in filter_ids}
+            mapping[0] = 0
+            segmentation = nt.takeDict(mapping, segmentation)
+        return segmentation.astype(dtype)
+
     def on_size_filter(self):
         size_threshold = self.size_threshold_widget.value()
         seg_layer = self._get_layer_selector_layer(self.segmentation_selector_name)
         segmentation = seg_layer.data.copy()
 
-        if self.is_mask:
-            segmentation = label(segmentation).astype(segmentation.dtype)
-
-        props = regionprops(segmentation)
-        filter_ids = [prop.label for prop in props if prop.area < size_threshold]
-        segmentation[np.isin(segmentation, filter_ids)] = 0
-        if self.is_mask:
-            segmentation = (segmentation > 0).astype(segmentation.dtype)
+        segmentation = self._filter_segmentation(segmentation, size_threshold, self.apply_label)
 
         # Write or overwrite segmentation layer.
         layer_name = self.output_layer_param.text()

synapse_net/tools/vesicle_pool_widget.py

@@ -108,6 +108,11 @@ def _update_pool_colors(self, pool_name, pool_color):
             next_color = pool_color
         self.pool_colors[pool_name] = next_color
 
+    def _add_colors(self, pool_properties, vesicle_colors):
+        colors = np.array([vesicle_colors[label_id] for label_id in pool_properties.label.values])
+        pool_properties["color"] = colors
+        return pool_properties
+
     def _compute_vesicle_pool(
         self,
         segmentation: np.ndarray,
@@ -136,9 +141,9 @@ def _compute_vesicle_pool(
             show_info("ERROR: Neither distances nor vesicle morphology were found.")
             return
         elif distances is None and morphology is not None:  # Only morphology props were found.
-            merged_df = pd.DataFrame(morphology).drop(columns=["index"])
+            merged_df = pd.DataFrame(morphology).drop(columns=["index"], errors="ignore")
         elif distances is not None and morphology is None:  # Only distances were found.
-            merged_df = pd.DataFrame(distances).drop(columns=["index"])
+            merged_df = pd.DataFrame(distances).drop(columns=["index"], errors="ignore")
         else:  # Both were found.
             distance_ids = distances.get("label", [])
             morphology_ids = morphology.get("label", [])
@@ -190,7 +195,7 @@ def _compute_vesicle_pool(
             # Overwrite the intersection of the two pool assignments with the new pool.
             pool_intersections = np.intersect1d(pool_vesicle_ids, old_pool_ids)
             old_pool_ids = [item for item in old_pool_ids if item not in pool_intersections]
-            pool_properties = pool_properties[~pool_properties['label'].isin(pool_intersections)]
+            pool_properties = pool_properties[~pool_properties["label"].isin(pool_intersections)]
 
             pool_assignments = sorted(pool_vesicle_ids + old_pool_ids)
 
@@ -242,6 +247,7 @@ def _compute_vesicle_pool(
         else:
             pool_layer = self.viewer.add_labels(vesicle_pools, name=pool_layer_name, colormap=vesicle_colors)
 
+        pool_properties = self._add_colors(pool_properties, vesicle_colors)
         self._add_properties_and_table(pool_layer, pool_properties, save_path=self.save_path.text())
         pool_layer.refresh()