refactor: optimized swc reader (#97)

Co-authored-by: anna-grim <[email protected]>

Author: anna-grim

src/segmentation_skeleton_metrics/skeleton_metric.py

@@ -108,6 +108,7 @@ def __init__(
         self.connections_path = connections_path
         self.output_dir = output_dir
         self.preexisting_merges = preexisting_merges
+        self.save_projections = save_projections
 
         # Label handler
         self.label_handler = gutil.LabelHandler(
@@ -116,16 +117,15 @@ def __init__(
 
         # Load data
         self.label_mask = pred_labels
-        self.load_groundtruth(gt_pointer, valid_labels)
+        self.load_groundtruth(gt_pointer)
         self.load_fragments(fragments_pointer)
 
         # Initialize writer
-        self.save_projections = save_projections
         if self.save_projections:
             self.init_zip_writer()
 
     # --- Load Data ---
-    def load_groundtruth(self, swc_pointer, valid_labels):
+    def load_groundtruth(self, swc_pointer):
         """
         Initializes "self.graphs" by iterating over "paths" which corresponds
         to neurons in the ground truth.
@@ -147,7 +147,6 @@ def load_groundtruth(self, swc_pointer, valid_labels):
             anisotropy=self.anisotropy,
             label_mask=self.label_mask,
             use_anisotropy=False,
-            valid_labels=valid_labels,
         )
         self.graphs = graph_builder.run(swc_pointer)
 
@@ -158,16 +157,18 @@ def load_groundtruth(self, swc_pointer, valid_labels):
     def load_fragments(self, swc_pointer):
         print("\n(2) Load Fragments")
         if swc_pointer:
+            coords_only = False #not self.save_projections
             graph_builder = gutil.GraphBuilder(
                 anisotropy=self.anisotropy,
+                coords_only=coords_only,
                 selected_ids=self.get_all_node_labels(),
                 use_anisotropy=True,
             )
             self.fragment_graphs = graph_builder.run(swc_pointer)
         else:
             self.fragment_graphs = None
 
-    def label_graphs(self, key, batch_size=128):
+    def label_graphs(self, key, batch_size=64):
         """
         Iterates over nodes in "graph" and stores the corresponding label from
         predicted segmentation mask (i.e. "self.label_mask") as a node-level

src/segmentation_skeleton_metrics/utils/graph_util.py

@@ -6,7 +6,7 @@
 @email: [email protected]
 
 """
-from concurrent.futures import ProcessPoolExecutor, as_completed
+from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
 from random import sample
 from scipy.spatial import distance
 from tqdm import tqdm
@@ -29,20 +29,20 @@ class GraphBuilder:
     def __init__(
         self,
         anisotropy=(1.0, 1.0, 1.0),
+        coords_only=False,
         label_mask=None,
         selected_ids=None,
         use_anisotropy=True,
-        valid_labels=None,
     ):
         # Instance attributes
         self.anisotropy = anisotropy
+        self.coords_only = coords_only
         self.label_mask = label_mask
         self.selected_ids = selected_ids
-        self.valid_labels = valid_labels
 
         # Reader
         anisotropy = anisotropy if use_anisotropy else (1.0, 1.0, 1.0)
-        self.swc_reader = swc_util.Reader(anisotropy)
+        self.swc_reader = swc_util.Reader(anisotropy, selected_ids=selected_ids)
 
     def run(self, swc_pointer):
         graphs = self._build_graphs_from_swcs(swc_pointer)
@@ -51,14 +51,14 @@ def run(self, swc_pointer):
 
     # --- Build Graphs ---
     def _build_graphs_from_swcs(self, swc_pointer):
-        with ProcessPoolExecutor() as executor:
+        with ThreadPoolExecutor() as executor:
             # Assign processes
             processes = list()
             swc_dicts = self.swc_reader.load(swc_pointer)
             while len(swc_dicts) > 0:
                 swc_dict = swc_dicts.pop()
-                if self._process_swc_dict(swc_dict["swc_id"]):
-                    processes.append(executor.submit(self.to_graph, swc_dict))
+                #if self._process_swc_dict(swc_dict["swc_id"]):
+                processes.append(executor.submit(self.to_graph, swc_dict))
 
             # Store results
             graphs = dict()
@@ -96,17 +96,19 @@ def to_graph(self, swc_dict):
         graph.set_voxels(swc_dict["voxel"])
 
         # Build graph
-        id_lookup = dict()
-        run_length = 0
-        for i, id_i in enumerate(swc_dict["id"]):
-            id_lookup[id_i] = i
-            if swc_dict["pid"][i] != -1:
-                parent = id_lookup[swc_dict["pid"][i]]
-                graph.add_edge(i, parent)
-                graph.run_length += graph.dist(i, parent)
-
-        # Set graph-level attributes
-        graph.graph["n_edges"] = graph.number_of_edges()
+        if not self.coords_only:
+            #graph.set_nodes()
+            id_lookup = dict()
+            run_length = 0
+            for i, id_i in enumerate(swc_dict["id"]):
+                id_lookup[id_i] = i
+                if swc_dict["pid"][i] != -1:
+                    parent = id_lookup[swc_dict["pid"][i]]
+                    graph.add_edge(i, parent)
+                    graph.run_length += graph.dist(i, parent)
+
+            # Set graph-level attributes
+            graph.graph["n_edges"] = graph.number_of_edges()
         return {swc_dict["swc_id"]: graph}
 
     # --- Label Graphs ---
@@ -127,6 +129,10 @@ def __init__(self, anisotropy=(1.0, 1.0, 1.0)):
         self.anisotropy = anisotropy
         self.run_length = 0
 
+    def set_nodes(self):
+        num_nodes = len(self.voxels)
+        self.add_nodes_from(np.arange(num_nodes))
+
     def set_voxels(self, voxels):
         self.voxels = np.array(voxels, dtype=np.int32)
 

src/segmentation_skeleton_metrics/utils/swc_util.py

@@ -48,7 +48,7 @@ class Reader:
 
     """
 
-    def __init__(self, anisotropy=(1.0, 1.0, 1.0)):
+    def __init__(self, anisotropy=(1.0, 1.0, 1.0), selected_ids=None):
         """
         Initializes a Reader object that loads swc files.
 
@@ -65,6 +65,7 @@ def __init__(self, anisotropy=(1.0, 1.0, 1.0)):
 
         """
         self.anisotropy = anisotropy
+        self.selected_ids = selected_ids or set()
 
     # --- Load Data ---
     def load(self, swc_pointer):
@@ -138,7 +139,10 @@ def load_from_local_path(self, path):
         """
         content = util.read_txt(path)
         filename = os.path.basename(path)
-        return self.parse(content, filename)
+        if self.confirm_load(filename):
+            return self.parse(content, filename)
+        else:
+            return None
 
     def load_from_local_paths(self, paths):
         """
@@ -161,9 +165,11 @@ def load_from_local_paths(self, paths):
             threads = list()
             pbar = tqdm(total=len(paths), desc="Read SWCs")
             for path in paths:
-                threads.append(
-                    executor.submit(self.load_from_local_path, path)
-                )
+                filename = os.path.basename(path)
+                if self.confirm_load(filename):
+                    threads.append(
+                        executor.submit(self.load_from_local_path, path)
+                    )
 
             # Store results
             swc_dicts = deque()
@@ -216,10 +222,14 @@ def load_from_local_zip(self, zip_path):
             # Assign threads
             threads = list()
             zipfile = ZipFile(zip_path, "r")
-            for f in [f for f in zipfile.namelist() if f.endswith(".swc")]:
-                threads.append(
-                    executor.submit(self.load_from_zipped_file, zipfile, f)
-                )
+            filesnames = [f for f in zipfile.namelist() if f.endswith(".swc")]
+            for filename in filesnames:
+                if self.confirm_load(filename):
+                    threads.append(
+                        executor.submit(
+                            self.load_from_zipped_file, zipfile, filename
+                        )
+                    )
 
             # Store results
             swc_dicts = deque()
@@ -249,6 +259,13 @@ def load_from_zipped_file(self, zipfile, path):
         filename = os.path.basename(path)
         return self.parse(content, filename)
 
+    def confirm_load(self, filename):
+        if len(self.selected_ids) > 0:
+            segment_id = get_segment_id(filename)
+            return True if segment_id in self.selected_ids else False
+        else:
+            return True
+
     # -- Process Text ---
     def parse(self, content, filename):
         """
@@ -335,6 +352,11 @@ def read_voxel(self, xyz_str, offset):
         xyz = [float(xyz_str[i]) + offset[i] for i in range(3)]
         return img_util.to_voxels(xyz, self.anisotropy)
 
+    
+# --- Helpers ---
+def get_segment_id(filename):
+    return int(filename.split(".")[0])
+
 
 def to_zipped_swc(zip_writer, graph, color=None):
     """