Add cli to visualize vesicle pool assignments

Author: constantinpape

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_tools = 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()