From f8c9eca374ea30f1c1e263eb50950973721ac0ee Mon Sep 17 00:00:00 2001 From: anna-grim Date: Mon, 8 Jun 2026 00:22:53 +0000 Subject: [PATCH] refactor: more memory efficient mips --- src/segmentation_skeleton_metrics/evaluate.py | 3 ++- src/segmentation_skeleton_metrics/visualization.py | 9 +++++---- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/src/segmentation_skeleton_metrics/evaluate.py b/src/segmentation_skeleton_metrics/evaluate.py index 7fdf517..e067aea 100644 --- a/src/segmentation_skeleton_metrics/evaluate.py +++ b/src/segmentation_skeleton_metrics/evaluate.py @@ -9,6 +9,7 @@ """ +from tqdm import tqdm from zipfile import ZipFile import numpy as np @@ -386,7 +387,7 @@ 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(desc="Save MIPs", gt_graphs.items()): # Save GT mips viz.save_mips([gt_graph], output_dir, gt_graph.name) diff --git a/src/segmentation_skeleton_metrics/visualization.py b/src/segmentation_skeleton_metrics/visualization.py index 1806339..df5075e 100644 --- a/src/segmentation_skeleton_metrics/visualization.py +++ b/src/segmentation_skeleton_metrics/visualization.py @@ -42,15 +42,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(