refactor: optimized memory consumption

Author: anna-grim

demo/evaluation_results.xls

@@ -1 +1,8 @@
  Label   -   xyz
+ 0   -   -1
+ 0   -   -1
+ 0   -   -1
+ 0   -   -1
+ 0   -   -1
+ 0   -   -1
+ 0   -   -1

demo/merged_ids-segmentation.txt

@@ -0,0 +1,15 @@
+import networkx as nx
+import numpy as np
+
+
+class SkeletonGraph(nx.Graph):
+
+    def __init__(self):
+        # Call parent class
+        super(SkeletonGraph, self).__init__()
+
+    def get_labels(self):
+        return np.unique(self.graph["label"])
+
+    def nodes_with_label(self, label):
+        return np.where(self.graph["label"] == label)[0]

demo/pred_labels-old.tif

@@ -12,9 +12,7 @@
     ProcessPoolExecutor,
     ThreadPoolExecutor,
 )
-from copy import deepcopy
 from scipy.spatial import distance, KDTree
-from time import time
 from tqdm import tqdm
 from zipfile import ZipFile
 
@@ -118,7 +116,7 @@ def __init__(
 
         # Load data
         self.label_mask = pred_labels
-        self.load_groundtruth(gt_pointer)
+        self.load_groundtruth(gt_pointer, valid_labels)
         self.load_fragments(fragments_pointer)
 
         # Initialize writer
@@ -127,7 +125,7 @@ def __init__(
             self.init_zip_writer()
 
     # --- Load Data ---
-    def load_groundtruth(self, swc_pointer):
+    def load_groundtruth(self, swc_pointer, valid_labels):
         """
         Initializes "self.graphs" by iterating over "paths" which corresponds
         to neurons in the ground truth.
@@ -149,16 +147,13 @@ def load_groundtruth(self, swc_pointer):
             anisotropy=self.anisotropy,
             label_mask=self.label_mask,
             use_anisotropy=False,
+            valid_labels=valid_labels,
         )
         self.graphs = graph_builder.run(swc_pointer)
 
         # Label nodes
-        self.key_to_label_to_nodes = dict()  # {id: {label: nodes}}
         for key in tqdm(self.graphs, desc="Labeling Graphs"):
             self.label_graphs(key)
-            self.key_to_label_to_nodes[key] = gutil.init_label_to_nodes(
-                self.graphs[key]
-            )
 
     def load_fragments(self, swc_pointer):
         print("\n(2) Load Fragments")
@@ -220,10 +215,12 @@ def label_graphs(self, key, batch_size=128):
             threads.append(executor.submit(self.get_patch_labels, key, batch))
 
             # Process results
+            n_nodes = self.graphs[key].number_of_nodes()
+            self.graphs[key].graph["label"] = np.zeros((n_nodes), dtype=int)
             for thread in as_completed(threads):
                 node_to_label = thread.result()
                 for i, label in node_to_label.items():
-                    self.graphs[key].nodes[i].update({"label": label})
+                    self.graphs[key].graph["label"][i] = label
 
     def get_patch_labels(self, key, nodes):
         # Get bounding box
@@ -287,11 +284,11 @@ def get_node_labels(self, key, inverse_bool=False):
         """
         if inverse_bool:
             output = set()
-            for l in self.key_to_label_to_nodes[key].keys():
+            for l in self.graphs[key].get_labels():
                 output = output.union(self.label_handler.inverse_mapping[l])
             return output
         else:
-            return set(self.key_to_label_to_nodes[key].keys())
+            return self.graphs[key].get_labels()
 
     def init_zip_writer(self):
         """
@@ -372,9 +369,9 @@ def adjust_metrics(self, key):
         """
         for label in self.preexisting_merges:
             label = self.label_map[label] if self.label_map else label
-            if label in self.key_to_label_to_nodes[key].keys():
+            if label in self.graphs[key].get_labels():
                 # Extract subgraph
-                nodes = deepcopy(self.key_to_label_to_nodes[key][label])
+                nodes = self.graphs[key].nodes_with_label(label)
                 subgraph = self.graphs[key].subgraph(nodes)
 
                 # Adjust metrics
@@ -385,7 +382,6 @@ def adjust_metrics(self, key):
 
                 # Update graph
                 self.graphs[key].remove_nodes_from(nodes)
-                del self.key_to_label_to_nodes[key][label]
 
     # -- Split Detection --
     def detect_splits(self):
@@ -402,7 +398,6 @@ def detect_splits(self):
         None
 
         """
-        t0 = time()
         pbar = tqdm(total=len(self.graphs), desc="Split Detection")
         with ProcessPoolExecutor() as executor:
             # Assign processes
@@ -420,17 +415,10 @@ def detect_splits(self):
             self.split_percent = dict()
             for process in as_completed(processes):
                 key, graph, split_percent = process.result()
-                self.graphs[key] = gutil.delete_nodes(graph, 0)
-                self.key_to_label_to_nodes[key] = gutil.init_label_to_nodes(
-                    self.graphs[key]
-                )
+                self.graphs[key] = gutil.remove_nodes(graph, 0)
                 self.split_percent[key] = split_percent
                 pbar.update(1)
 
-        # Report runtime
-        t, unit = util.time_writer(time() - t0)
-        print(f"Runtime: {round(t, 2)} {unit}\n")
-
     def quantify_splits(self):
         """
         Counts the number of splits, number of omit edges, and percent of omit
@@ -449,9 +437,11 @@ def quantify_splits(self):
         self.omit_cnts = dict()
         self.omit_percent = dict()
         for key in self.graphs:
+            # Get counts
             n_pred_edges = self.graphs[key].number_of_edges()
             n_target_edges = self.graphs[key].graph["n_edges"]
 
+            # Compute stats
             self.split_cnt[key] = gutil.count_splits(self.graphs[key])
             self.omit_cnts[key] = n_target_edges - n_pred_edges
             self.omit_percent[key] = 1 - n_pred_edges / n_target_edges
@@ -517,7 +507,8 @@ def count_merges(self, key, kdtree):
 
         """
         for label in self.get_node_labels(key):
-            if len(self.key_to_label_to_nodes[key][label]) > MIN_CNT:
+            nodes = self.graphs[key].nodes_with_label(label)
+            if len(nodes) > MIN_CNT:
                 for label in self.label_handler.get_class(label):
                     if label in self.fragment_graphs:
                         self.is_fragment_merge(key, label, kdtree)
@@ -581,8 +572,8 @@ def find_label_intersections(self):
                 keys = frozenset((key_1, key_2))
                 if key_1 != key_2 and keys not in visited:
                     visited.add(keys)
-                    labels_1 = self.get_node_labels(key_1)
-                    labels_2 = self.get_node_labels(key_2)
+                    labels_1 = set(self.graphs[key_1].get_labels())
+                    labels_2 = set(self.graphs[key_2].get_labels())
                     for label in labels_1.intersection(labels_2):
                         label_intersections.add((keys, label))
         return label_intersections
@@ -605,15 +596,14 @@ def process_merge(self, key, label, xyz, update_merged_labels=True):
         None
 
         """
-        if label in self.key_to_label_to_nodes[key]:
+        if label in self.graphs[key].get_labels():
             # Compute metrics
-            nodes = list(self.key_to_label_to_nodes[key][label])
+            nodes = self.graphs[key].nodes_with_label(label)
             subgraph = self.graphs[key].subgraph(nodes)
             self.merged_edges_cnt[key] += subgraph.number_of_edges()
 
             # Update self
             self.graphs[key].remove_nodes_from(nodes)
-            del self.key_to_label_to_nodes[key][label]
             if update_merged_labels:
                 self.merged_labels.add((key, label, -1))
 

src/segmentation_skeleton_metrics/skeleton_graph.py

@@ -45,8 +45,8 @@ def run(process_id, graph):
             continue
 
         # Visit edge
-        label_i = graph.nodes[i]["label"]
-        label_j = graph.nodes[j]["label"]
+        label_i = int(graph.graph["label"][i])
+        label_j = int(graph.graph["label"][j])
         if is_split(label_i, label_j):
             graph.remove_edge(i, j)
             split_cnt += 1
@@ -91,12 +91,13 @@ def check_misalignment(graph, visited_edges, nb, root):
     while len(queue) > 0:
         # Visit node
         j = queue.popleft()
-        if graph.nodes[j]["label"] != 0:
-            label_collisions.add(graph.nodes[j]["label"])
+        label_j = int(graph.graph["label"][j])
+        if label_j != 0:
+            label_collisions.add(label_j)
         visited.add(j)
 
         # Update queue
-        if graph.nodes[j]["label"] == 0:
+        if label_j == 0:
             for k in graph.neighbors(j):
                 if k not in visited:
                     if frozenset({j, k}) not in visited_edges or k == nb:
@@ -106,7 +107,7 @@ def check_misalignment(graph, visited_edges, nb, root):
     # Upd zero nodes
     if len(label_collisions) == 1:
         label = label_collisions.pop()
-        graph = gutil.upd_labels(graph, visited, label)
+        upd_labels(graph, visited, label)
 
 
 # -- Helpers --
@@ -128,3 +129,26 @@ def is_split(a, b):
 
     """
     return (a > 0 and b > 0) and (a != b)
+
+
+def upd_labels(graph, nodes, label):
+    """
+    Updates the label of each node in "nodes" with "label".
+
+    Parameters
+    ----------
+    graph : networkx.Graph
+        Graph to be updated.
+    nodes : list
+        List of nodes to be updated.
+    label : int
+        New label of each node in "nodes".
+
+    Returns
+    -------
+    networkx.Graph
+        Updated graph.
+
+    """
+    for i in nodes:
+        graph.graph["label"][i] = label

src/segmentation_skeleton_metrics/skeleton_metric.py

@@ -15,6 +15,7 @@
 import numpy as np
 from scipy.spatial import distance
 
+from segmentation_skeleton_metrics.skeleton_graph import SkeletonGraph
 from segmentation_skeleton_metrics.utils import img_util, swc_util, util
 
 ANISOTROPY = np.array([0.748, 0.748, 1.0])
@@ -28,11 +29,13 @@ def __init__(
         label_mask=None,
         selected_ids=None,
         use_anisotropy=True,
+        valid_labels=None,
     ):
         # Instance attributes
         self.anisotropy = anisotropy
         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)
@@ -62,9 +65,9 @@ def _build_graphs_from_swcs(self, swc_pointer):
     def _process_swc_dict(self, swc_id):
         if self.selected_ids:
             segment_id = get_segment_id(swc_id)
-            if segment_id not in self.selected_ids:
-                return False
-        return True
+            return True if segment_id in self.selected_ids else False
+        else:
+            return True
 
     def to_graph(self, swc_dict):
         """
@@ -87,7 +90,7 @@ def to_graph(self, swc_dict):
         voxels = np.array(swc_dict["voxel"], dtype=np.int32)
 
         # Build graph
-        graph = nx.Graph()
+        graph = SkeletonGraph()
         id_lookup = dict()
         run_length = 0
         for i in range(len(swc_dict["id"])):
@@ -202,23 +205,7 @@ def build_labels_graph(self, connections_path):
 
     # --- Main ---
     def get(self, label):
-        """
-        Gets label of voxel in "self.label_mask".
-
-        Parameters
-        ----------
-        i : int
-            Node ID.
-        voxel : numpy.ndarray
-            Image coordinate of voxel to be read.
-
-        Returns
-        -------
-        int
-           Label of voxel.
-
-        """
-        if len(self.mapping) > 0:
+        if self.use_mapping():
             return self.mapping.get(label, 0)
         elif self.valid_labels:
             return 0 if label not in self.valid_labels else label
@@ -232,14 +219,14 @@ def use_mapping(self):
 
 
 # --- Update graph ---
-def delete_nodes(graph, target_label):
+def remove_nodes(graph, target_label):
     """
-    Deletes nodes in "graph" whose label is "target_label".
+    Deletes nodes in the given graph whose label is "target_label".
 
     Parameters
     ----------
     graph : networkx.Graph
-        Graph to be searched and edited.
+        Graph with a graph-level attribute called "label".
     target_label : int
         Label to be deleted from graph.
 
@@ -249,36 +236,8 @@ def delete_nodes(graph, target_label):
         Updated graph.
 
     """
-    delete_nodes = []
-    for i in graph.nodes:
-        label = graph.nodes[i]["label"]
-        if label == target_label:
-            delete_nodes.append(i)
-    graph.remove_nodes_from(delete_nodes)
-    return graph
-
-
-def upd_labels(graph, nodes, label):
-    """
-    Updates the label of each node in "nodes" with "label".
-
-    Parameters
-    ----------
-    graph : networkx.Graph
-        Graph to be updated.
-    nodes : list
-        List of nodes to be updated.
-    label : int
-        New label of each node in "nodes".
-
-    Returns
-    -------
-    networkx.Graph
-        Updated graph.
-
-    """
-    for i in nodes:
-        graph.nodes[i].update({"label": label})
+    nodes = np.where(graph.graph["label"] == target_label)[0]
+    graph.remove_nodes_from(nodes)
     return graph
 
 

src/segmentation_skeleton_metrics/split_detection.py

@@ -181,7 +181,7 @@ def load_from_local_zips(self, zip_dir):
         with ProcessPoolExecutor() as executor:
             # Assign threads
             processes = list()
-            for f in zip_names:
+            for f in zip_names[0:250]:  # TEMP
                 zip_path = os.path.join(zip_dir, f)
                 processes.append(
                     executor.submit(self.load_from_local_zip, zip_path)