Refactor save mips

src/segmentation_skeleton_metrics/evaluate.py

@@ -9,6 +9,7 @@
 
 """
 
+from tqdm import tqdm
 from zipfile import ZipFile
 
 import numpy as np
@@ -386,14 +387,15 @@ def save_mips(self, gt_graphs, fragment_graphs):
         """
         output_dir = os.path.join(self.output_dir, f"{self.prefix}mips")
         util.mkdir(output_dir, delete=True)
-        for key, gt_graph in gt_graphs.items():
+        for key, gt_graph in tqdm(gt_graphs.items(), desc="Save MIPs"):
             # Save GT mips
             viz.save_mips([gt_graph], output_dir, gt_graph.name)
 
             # Save fragment mips
             filename = f"{gt_graph.name}-fragments"
             fragments = get_intersecting_fragments(gt_graph, fragment_graphs)
             viz.save_mips(fragments, output_dir, filename)
+            del fragments
 
     def save_skeletons_with_merge(self, gt_graphs, fragment_graphs, zf):
         """

src/segmentation_skeleton_metrics/visualization.py

@@ -32,7 +32,6 @@ def save_mips(graph_list, output_dir, filename, dilation=16):
     dilation : int, optional
         Dilation radius applied during graph rasterization. Default is 16.
     """
-    _, shape = _get_combined_bbox(graph_list)
     mip_xy, mip_xz, mip_yz = _rasterize_graphs(graph_list, dilation)
     _plot_and_save_mips(mip_xy, mip_xz, mip_yz, output_dir, filename)
 
@@ -42,15 +41,16 @@ def _rasterize_graphs(graph_list, dilation):
     struct2d = np.ones((dilation,) * 2, dtype=bool)
     min_voxel, shape = _get_combined_bbox(graph_list)
 
-    mip_xy = np.ones((shape[1], shape[2], 3), dtype=float)
-    mip_xz = np.ones((shape[0], shape[2], 3), dtype=float)
-    mip_yz = np.ones((shape[0], shape[1], 3), dtype=float)
-
+    mip_xy = np.full((shape[1], shape[2], 3), 255, dtype=np.uint8)
+    mip_xz = np.full((shape[0], shape[2], 3), 255, dtype=np.uint8)
+    mip_yz = np.full((shape[0], shape[1], 3), 255, dtype=np.uint8)
     cc_idx = 0
     for graph in graph_list:
         shifted_voxels = graph.node_voxel - min_voxel
         for cc_nodes in nx.connected_components(graph):
             color = np.array(colors[cc_idx % len(colors)])
+            if color.dtype != np.uint8:
+                color = (color * 255).astype(np.uint8)
             cc_voxels = shifted_voxels[list(cc_nodes)]
             z, y, x = cc_voxels[:, 0], cc_voxels[:, 1], cc_voxels[:, 2]
             _paint_projections(
@@ -101,9 +101,13 @@ def _make_dilated_local(a, b, mip_shape, struct2d, pad):
 
 
 def _get_combined_bbox(graph_list):
-    all_voxels = np.vstack([graph.node_voxel for graph in graph_list])
-    min_voxel = all_voxels.min(axis=0)
-    shape = tuple((all_voxels.max(axis=0) - min_voxel) + 1)
+    min_voxel = np.full(3, np.inf)
+    max_voxel = np.full(3, -np.inf)
+    for graph in graph_list:
+        np.minimum(min_voxel, graph.node_voxel.min(axis=0), out=min_voxel)
+        np.maximum(max_voxel, graph.node_voxel.max(axis=0), out=max_voxel)
+    min_voxel = min_voxel.astype(int)
+    shape = tuple((max_voxel - min_voxel).astype(int) + 1)
     return min_voxel, shape