bug: fixed memory issue

Author: anna-grim

src/segmentation_skeleton_metrics/skeleton_metric.py

@@ -7,7 +7,7 @@
 
 """
 
-
+from collections import deque
 from concurrent.futures import (
     as_completed,
     ProcessPoolExecutor,
@@ -31,7 +31,7 @@
     util
 )
 
-MERGE_DIST_THRESHOLD = 100
+MERGE_DIST_THRESHOLD = 200
 MIN_CNT = 40
 
 
@@ -112,15 +112,19 @@ def __init__(
         self.output_dir = output_dir
         self.preexisting_merges = preexisting_merges
 
-        # Load Data
-        print("\n(1) Load Data")
+        # Load ground truth
+        print("\n(1) Load Ground Truth")
         assert type(valid_labels) is set if valid_labels else True
-        self.label_mask = pred_labels
         self.valid_labels = valid_labels
         self.init_label_map(connections_path)
         self.init_graphs(gt_pointer)
+
+        print("\n(2) Load Prediction")
+        self.label_mask = pred_labels
         if fragments_pointer:
             self.load_fragments(fragments_pointer)
+        else:
+            self.fragment_graphs = None
 
         # Initialize writer
         self.save_projections = save_projections
@@ -160,7 +164,7 @@ def init_graphs(self, paths):
 
         Parameters
         ----------
-        paths : list[str]
+        paths : List[str]
             List of paths to swc files which correspond to neurons in the
             ground truth.
 
@@ -170,18 +174,35 @@ def init_graphs(self, paths):
 
         """
         # Build graphs
-        self.graphs = swc_util.Reader().load(paths)
-        self.fragment_graphs = None
+        swc_dicts = swc_util.Reader().load(paths)
+        self.graphs = self.build_graphs(swc_dicts)
 
         # Label nodes
         self.key_to_label_to_nodes = dict()  # {id: {label: nodes}}
         for key in tqdm(self.graphs, desc="Labeling Graphs"):
-            self.set_node_labels(key)
+            self.label_graphs(key)
             self.key_to_label_to_nodes[key] = gutil.init_label_to_nodes(
                 self.graphs[key]
             )
 
-    def set_node_labels(self, key, batch_size=128):
+    def build_graphs(self, swc_dicts):
+        graphs = dict()
+        with ProcessPoolExecutor() as executor:
+            # Assign processes
+            processes = list()
+            for swc_dict in swc_dicts:
+                processes.append(
+                    executor.submit(gutil.to_graph, swc_dict)
+                )
+
+            # Store results
+            pbar = tqdm(total=len(processes), desc="Build Graphs")
+            for process in as_completed(processes):
+                graphs.update(process.result())
+                pbar.update(1)
+        return graphs
+            
+    def label_graphs(self, key, batch_size=128):
         """
         Iterates over nodes in "graph" and stores the corresponding label from
         predicted segmentation mask (i.e. "self.label_mask") as a node-level
@@ -238,7 +259,7 @@ def get_patch_labels(self, key, nodes):
         # Get bounding box
         bbox = {"min": [np.inf, np.inf, np.inf], "max": [0, 0, 0]}
         for i in nodes:
-            voxel = deepcopy(self.graphs[key].nodes[i]["voxel"])
+            voxel = deepcopy(self.graphs[key].graph["voxel"][i])
             for idx in range(3):
                 if voxel[idx] < bbox["min"][idx]:
                     bbox["min"][idx] = voxel[idx]
@@ -359,20 +380,20 @@ def load_fragments(self, fragments_pointer):
             Dictionary that maps an swc id to the fragment graph.
 
         """
-        # Read fragments
-        reader = swc_util.Reader(anisotropy=self.anisotropy, min_size=40)
-        fragment_graphs = reader.load(fragments_pointer)
-        self.fragment_ids = set(fragment_graphs.keys())
-
-        # Filter fragments
-        self.fragment_graphs = dict()
-        for label in self.get_all_node_labels():
-            if label in fragment_graphs:
-                self.fragment_graphs[label] = fragment_graphs[label]
-            else:
-                self.fragment_graphs[label] = nx.Graph(
-                    filename=f"{label}.swc", run_length=0, n_edges=1
-                )
+        # Read SWC files
+        reader = swc_util.Reader(anisotropy=self.anisotropy)
+        swc_dicts = deque(reader.load(fragments_pointer))
+
+        # Filter SWC files
+        filtered_swc_dicts = list()
+        labels = self.get_all_node_labels()
+        while len(swc_dicts) > 0:
+            swc_dict = swc_dicts.popleft()
+            swc_id = int(swc_dict["swc_id"])
+            if swc_id in labels:
+                swc_dict["swc_id"] = swc_id
+                filtered_swc_dicts.append(swc_dict)
+        self.fragment_graphs = self.build_graphs(filtered_swc_dicts)
         print("# Fragments:", len(self.fragment_graphs))
 
     def init_zip_writer(self):
@@ -416,7 +437,7 @@ def run(self):
             ...
 
         """
-        print("\n(2) Evaluation")
+        print("\n(3) Evaluation")
 
         # Split evaluation
         self.detect_splits()
@@ -564,16 +585,14 @@ def detect_merges(self):
             pbar = tqdm(total=len(self.graphs), desc="Count Merges:")
             for key, graph in self.graphs.items():
                 if graph.number_of_nodes() > 0:
-                    kdtree = KDTree(gutil.to_array(graph))
+                    kdtree = KDTree(graph.graph["voxel"])
                     self.count_merges(key, kdtree)
                 pbar.update(1)
 
         # Process merges
-        pbar = tqdm(total=len(self.graphs), desc="Compute Percent Merged:")
         for (key_1, key_2), label in self.find_label_intersections():
             self.process_merge(key_1, label, -1)
             self.process_merge(key_2, label, -1)
-            pbar.update(1)
 
         for key, label, xyz in self.merged_labels:
             self.process_merge(key, label, xyz, update_merged_labels=False)
@@ -610,8 +629,8 @@ def count_merges(self, key, kdtree):
 
                 # Check if fragment is a merge mistake
                 for label in labels:
-                    rl = self.fragment_graphs[label].graph["run_length"]
-                    self.is_fragment_merge(key, label, kdtree)
+                    if label in self.fragment_graphs:
+                        self.is_fragment_merge(key, label, kdtree)
 
     def is_fragment_merge(self, key, label, kdtree):
         """
@@ -634,7 +653,7 @@ def is_fragment_merge(self, key, label, kdtree):
         None
 
         """
-        for voxel in gutil.to_array(self.fragment_graphs[label])[::2]:
+        for voxel in self.fragment_graphs[label].graph["voxel"]:
             if kdtree.query(voxel)[0] > MERGE_DIST_THRESHOLD:
                 # Check whether to get inverse of label
                 if self.inverse_label_map:
@@ -643,10 +662,10 @@ def is_fragment_merge(self, key, label, kdtree):
                     equivalent_label = label
 
                 # Record merge mistake
-                xyz = img_util.to_physical(voxel)
+                xyz = img_util.to_physical(voxel, self.anisotropy)
                 self.merge_cnt[key] += 1
                 self.merged_labels.add((key, equivalent_label, tuple(xyz)))
-                if self.save_projections:
+                if self.save_projections and label in self.fragment_graphs:
                     swc_util.to_zipped_swc(
                         self.zip_writer[key], self.fragment_graphs[label]
                     )
@@ -768,13 +787,13 @@ def get_merged_label(self, label):
         Returns:
         -------
         str or list
-            The first matching label found in "self.fragment_ids" or the
-            original associated labels from "inverse_label_map" if no matches
-            are found.
+            The first matching label found in "self.fragment_graphs.keys()" or
+            the original associated labels from "inverse_label_map" if no 
+            matches are found.
 
         """
         for l in self.inverse_label_map[label]:
-            if l in self.fragment_ids:
+            if l in self.fragment_graphs.keys():
                 return l
         return self.inverse_label_map[label]
 
@@ -988,8 +1007,8 @@ def list_metrics(self):
 
     # -- util --
     def dist(self, key, i, j):
-        xyz_i = self.graphs[key].nodes[i]["voxel"]
-        xyz_j = self.graphs[key].nodes[j]["voxel"]
+        xyz_i = self.graphs[key].graph["voxel"][i]
+        xyz_j = self.graphs[key].graph["voxel"][j]
         return distance.euclidean(xyz_i, xyz_j)
 
     def init_counter(self):
@@ -1010,7 +1029,7 @@ def init_counter(self):
         return {key: 0 for key in self.graphs}
 
     def to_local_voxels(self, key, i, offset):
-        voxel = np.array(self.graphs[key].nodes[i]["voxel"])
+        voxel = np.array(self.graphs[key].graph["voxel"][i])
         offset = np.array(offset)
         return tuple(voxel - offset)
 

src/segmentation_skeleton_metrics/utils/graph_util.py

@@ -12,11 +12,58 @@
 
 import networkx as nx
 import numpy as np
-from scipy.spatial.distance import euclidean as get_dist
+from scipy.spatial import distance
 
 from segmentation_skeleton_metrics.utils import img_util
 
-ANISOTROPY = (0.748, 0.748, 1.0)
+ANISOTROPY = np.array([0.748, 0.748, 1.0])
+
+
+def to_graph(swc_dict):
+    """
+    Builds a graph from a dictionary that contains the contents of an SWC
+    file.
+
+    Parameters
+    ----------
+    swc_dict : dict
+        ...
+
+    Returns
+    -------
+    networkx.Graph
+        Graph built from an SWC file.
+
+    """
+    # Initializations
+    old_to_new = dict()
+    run_length = 0
+    voxels = np.zeros((len(swc_dict["id"]), 3), dtype=np.int32)
+
+    # Build graph
+    graph = nx.Graph()
+    for i in range(len(swc_dict["id"])):            
+        # Get node id
+        old_id = swc_dict["id"][i]
+        old_to_new[old_id] = i
+
+        # Update graph
+        voxels[i] = swc_dict["voxel"][i]
+        if swc_dict["pid"][i] != -1:
+            # Add edge
+            parent = old_to_new[swc_dict["pid"][i]]
+            graph.add_edge(i, parent)
+
+            # Update run length
+            xyz_i = voxels[i] * ANISOTROPY
+            xyz_p = voxels[parent] * ANISOTROPY
+            run_length += distance.euclidean(xyz_i, xyz_p)
+
+    # Set graph-level attributes
+    graph.graph["n_edges"] = graph.number_of_edges()
+    graph.graph["run_length"] = run_length
+    graph.graph["voxel"] = voxels
+    return {swc_dict["swc_id"]: graph}
 
 
 # --- Update graph ---
@@ -153,32 +200,13 @@ def compute_run_length(graph):
     """
     path_length = 0
     for i, j in nx.dfs_edges(graph):
-        xyz_1 = img_util.to_physical(graph.nodes[i]["voxel"], ANISOTROPY)
-        xyz_2 = img_util.to_physical(graph.nodes[j]["voxel"], ANISOTROPY)
-        path_length += get_dist(xyz_1, xyz_2)
+        xyz_i = img_util.to_physical(graph.graph["voxel"][i], ANISOTROPY)
+        xyz_j = img_util.to_physical(graph.graph["voxel"][j], ANISOTROPY)
+        path_length += distance.euclidean(xyz_i, xyz_j)
     return path_length
 
 
 # -- miscellaneous --
-def to_array(graph):
-    """
-    Converts node coordinates from a graph into a NumPy array.
-
-    Parameters
-    ----------
-    graph : networkx.Graph
-        Graph that contains nodes with "voxel" attributes.
-
-    Returns
-    -------
-    numpy.ndarray
-        Array where each row represents the 3D coordinates of a node.
-
-    """
-    voxels = nx.get_node_attributes(graph, "voxel")
-    return np.array([voxels[i] for i in graph.nodes])
-
-
 def sample_leaf(graph):
     """
     Samples leaf node from "graph".