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