Update pool visualization

Author: constantinpape

synapse_net/tools/cli.py

@@ -107,17 +107,23 @@ def pool_visualization_cli():
         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."
+        "--vesicle_paths", "-v", required=True, nargs="+",
+        help="The filepath(s) to the tif file(s) containing the vesicle segmentation."
     )
     parser.add_argument(
         "--table_paths", "-t", required=True, nargs="+",
-        help="The filepath to the table with the vesicle pool assignments."
+        help="The filepath(s) to the table(s) with the vesicle pool assignments."
     )
     parser.add_argument(
         "-s", "--segmentation_paths", nargs="+", help="Filepaths for additional segmentations."
     )
+    parser.add_argument(
+        "--split_pools", action="store_true", help="Whether to split the pools into individual layers.",
+    )
     args = parser.parse_args()
-    _visualize_vesicle_pools(args.input_path, args.vesicle_path, args.table_paths, args.segmentation_paths)
+    _visualize_vesicle_pools(
+        args.input_path, args.vesicle_paths, args.table_paths, args.segmentation_paths, args.split_pools
+    )
 
 
 # TODO: handle kwargs

synapse_net/tools/pool_visualization.py

@@ -2,19 +2,34 @@
 
 import imageio.v3 as imageio
 import napari
+import numpy as np
 import pandas as pd
 
+from elf.parallel import isin
 from ..file_utils import read_mrc
 
 
-def _create_pools(vesicles, table):
+def _create_pools(vesicles, table, split_pools):
     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
+    if split_pools:
+        unique_colors = np.unique(pool_colors)
+        pool_ret = {}
+        for this_color in unique_colors:
+            this_pool = pools.copy()
+            this_ids = [label_id for label_id, color in colormap.items() if color == this_color]
+            pool_mask = np.zeros(this_pool.shape, dtype="bool")
+            pool_mask = isin(this_pool, this_ids, out=pool_mask, block_shape=(32, 128, 128))
+            this_pool[~pool_mask] = 0
+            pool_ret[this_color] = this_pool
+    else:
+        pool_ret = {"pools": pools}
+
+    return pool_ret, colormap
 
 
 def _parse_tables(table_paths):
@@ -39,25 +54,36 @@ def load_table(path):
     return table
 
 
-def _visualize_vesicle_pools(input_path, vesicle_path, table_paths, segmentation_paths):
+def _visualize_vesicle_pools(input_path, vesicle_paths, table_paths, segmentation_paths, split_pools):
     # 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 vesicle layer, either from a single file with
+    if len(vesicle_paths) == 1:
+        vesicles = imageio.imread(vesicle_paths)
+    else:
+        vesicles = None
+        for path in vesicle_paths:
+            this_vesicles = imageio.imread(path)
+            if vesicles is None:
+                vesicles = this_vesicles.copy()
+            else:
+                ves_mask = this_vesicles != 0
+                vesicles[ves_mask] = this_vesicles[ves_mask]
 
     # Load the tables with the pool assignments.
     # Create and add the pool layer.
     table = _parse_tables(table_paths)
-    pools, colormap = _create_pools(vesicles, table)
+    pools, colormap = _create_pools(vesicles, table, split_pools)
 
     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)
+    for pool_name, pool in pools.items():
+        viewer.add_labels(pool, scale=scale, name=pool_name, colormap=colormap)
 
     # Add the additional segmentations.
     if segmentation_paths is not None: