synapse_net.tools.pool_visualization

  1from pathlib import Path
  2
  3import imageio.v3 as imageio
  4import napari
  5import numpy as np
  6import pandas as pd
  7
  8from elf.parallel import isin
  9from ..file_utils import read_mrc
 10
 11
 12def _create_pools(vesicles, table, split_pools):
 13    label_ids, pool_colors = table.label.values, table.color.values
 14
 15    pools = vesicles
 16    colormap = {label_id: color for label_id, color in zip(label_ids, pool_colors)}
 17    colormap[None] = [0, 0, 0, 0]
 18
 19    if split_pools:
 20        unique_colors = np.unique(pool_colors)
 21        pool_ret = {}
 22        for this_color in unique_colors:
 23            this_pool = pools.copy()
 24            this_ids = [label_id for label_id, color in colormap.items() if color == this_color]
 25            pool_mask = np.zeros(this_pool.shape, dtype="bool")
 26            pool_mask = isin(this_pool, this_ids, out=pool_mask, block_shape=(32, 128, 128))
 27            this_pool[~pool_mask] = 0
 28            pool_ret[this_color] = this_pool
 29    else:
 30        pool_ret = {"pools": pools}
 31
 32    return pool_ret, colormap
 33
 34
 35def _parse_tables(table_paths):
 36    def load_table(path):
 37        if path.endswith(".csv"):
 38            return pd.read_csv(path)
 39        elif path.endswith(".xlsx"):
 40            return pd.read_excel(path)
 41        else:
 42            raise RuntimeError("Unknown file ending.")
 43
 44    if len(table_paths) == 1:
 45        table = load_table(table_paths[0])
 46    else:
 47        table = []
 48        for table_path in table_paths:
 49            this_table = load_table(table_path)
 50            pool_name = Path(table_path).stem
 51            this_table["pool"] = [pool_name] * len(this_table)
 52            table.append(this_table)
 53        table = pd.concat(table)
 54    return table
 55
 56
 57def _visualize_vesicle_pools(input_path, vesicle_paths, table_paths, segmentation_paths, split_pools):
 58    # Load the tomogram data, including scale information.
 59    data, voxel_size = read_mrc(input_path)
 60    axes = "zyx" if data.ndim == 3 else "yx"
 61    scale = tuple(float(voxel_size[ax]) for ax in axes)
 62    print("Loading data with scale", scale, "nanometer")
 63
 64    # Load the vesicle layer, either from a single file with
 65    if len(vesicle_paths) == 1:
 66        vesicles = imageio.imread(vesicle_paths)
 67    else:
 68        vesicles = None
 69        for path in vesicle_paths:
 70            this_vesicles = imageio.imread(path)
 71            if vesicles is None:
 72                vesicles = this_vesicles.copy()
 73            else:
 74                ves_mask = this_vesicles != 0
 75                vesicles[ves_mask] = this_vesicles[ves_mask]
 76
 77    # Load the tables with the pool assignments.
 78    # Create and add the pool layer.
 79    table = _parse_tables(table_paths)
 80    pools, colormap = _create_pools(vesicles, table, split_pools)
 81
 82    viewer = napari.Viewer()
 83    viewer.add_image(data, scale=scale)
 84    viewer.add_labels(vesicles, scale=scale)
 85    for pool_name, pool in pools.items():
 86        viewer.add_labels(pool, scale=scale, name=pool_name, colormap=colormap)
 87
 88    # Add the additional segmentations.
 89    if segmentation_paths is not None:
 90        for seg_path in segmentation_paths:
 91            name = Path(seg_path).stem
 92            seg = imageio.imread(seg_path)
 93            viewer.add_labels(seg, name=name, scale=scale)
 94
 95    # FIXME something is wrong here.
 96    # Add the scale bar.
 97    # @magicgui(call_button="Add Scale Bar")
 98    # def add_scale_bar(v: napari.Viewer):
 99    #     v.scale_bar.visible = True
100    #     v.scale_bar.unit = "nm"
101    # viewer.window.add_dock_widget(add_scale_bar)
102
103    napari.run()