minor upds

src/deep_neurographs/groundtruth_generation.py

@@ -18,7 +18,7 @@
 from deep_neurographs.geometry import dist as get_dist
 from deep_neurographs.utils import util
 
-ALIGNED_THRESHOLD = 4
+ALIGNED_THRESHOLD = 4.5
 MIN_INTERSECTION = 10
 
 

src/deep_neurographs/inference.py

@@ -234,7 +234,7 @@ def build_graph(self, fragments_pointer):
         # Save valid labels and current graph
         swcs_path = os.path.join(self.output_dir, "processed-swcs.zip")
         labels_path = os.path.join(self.output_dir, "valid_labels.txt")
-        n_saved = self.graph.to_zipped_swcs(swcs_path)
+        n_saved = self.graph.to_zipped_swcs(swcs_path, min_size=100)
         self.graph.save_labels(labels_path)
         self.report(f"# SWCs Saved: {n_saved}")
 

src/deep_neurographs/machine_learning/heterograph_models.py

@@ -255,7 +255,6 @@ def forward(self, x_dict, edge_index_dict, edge_attr_dict):
 
 # --- Utils ---
 def reformat_edge_key(key):
-    print(key)
     return tuple([rm_non_alphanumeric(s) for s in key.split(",")])
 
 

src/deep_neurographs/neurograph.py

@@ -954,7 +954,7 @@ def component_path_length(self, root):
         return path_length
 
     # --- write graph to swcs ---
-    def to_zipped_swcs(self, zip_path, color=None, min_size=100):
+    def to_zipped_swcs(self, zip_path, color=None, min_size=0):
         with zipfile.ZipFile(zip_path, "w") as zip_writer:
             cnt = 0
             for nodes in nx.connected_components(self):

src/deep_neurographs/utils/graph_util.py

@@ -152,7 +152,7 @@ def schedule_processes(self, swc_dicts):
             a connected component of the graph corresponding to "swc_dicts".
 
         """
-        with ProcessPoolExecutor() as executor:
+        with ProcessPoolExecutor(max_workers=1) as executor:
             # Assign Processes
             i = 0
             processes = [None] * len(swc_dicts)
@@ -195,7 +195,7 @@ def get_irreducibles(self, swc_dict):
         # Build dense graph
         swc_dict["idx"] = dict(zip(swc_dict["id"], range(len(swc_dict["id"]))))
         graph, _ = swc_util.to_graph(swc_dict, set_attrs=True)
-        self.clip_branches(graph)
+        self.clip_branches(graph, swc_dict["swc_id"])
         self.prune_branches(graph)
 
         # Extract irreducibles
@@ -210,7 +210,7 @@ def get_irreducibles(self, swc_dict):
                         irreducibles.append(result)
         return irreducibles
 
-    def clip_branches(self, graph):
+    def clip_branches(self, graph, swc_id):
         """
         Deletes all nodes from "graph" that are not contained in "img_bbox".
 
@@ -232,6 +232,49 @@ def clip_branches(self, graph):
                     delete_nodes.add(i)
             graph.remove_nodes_from(delete_nodes)
 
+    def prune_branches(self, graph):
+        """
+        Prunes all short branches from "graph". A short branch is a path
+        between a leaf and branching node where the path length is less than
+        "self.prune_depth".
+
+        Parameters
+        ----------
+        graph : networkx.Graph
+            Graph to be pruned.
+
+        Returns
+        -------
+        networkx.Graph
+            Graph with short branches pruned.
+
+        """
+        first_pass = True
+        deleted_nodes = list()
+        n_passes = 0
+        while len(deleted_nodes) > 0 or first_pass:
+            # Visit leafs
+            n_passes += 1
+            deleted_nodes = list()
+            for leaf in get_leafs(graph):
+                branch = [leaf]
+                length = 0
+                for (i, j) in nx.dfs_edges(graph, source=leaf):
+                    # Visit edge
+                    length += compute_dist(graph, i, j)
+                    if graph.degree(j) == 2:
+                        branch.append(j)
+                    elif graph.degree(j) > 2:
+                        deleted_nodes.extend(branch)
+                        graph.remove_nodes_from(branch)
+                        break
+
+                    # Check whether to stop
+                    if length > self.prune_depth or first_pass:
+                        graph.remove_nodes_from(branch[0:min(3, len(branch))])
+                        break
+            first_pass = False
+
     def get_component_irreducibles(self, graph, swc_dict):
         """
         Gets the irreducible components of "graph".
@@ -301,50 +344,6 @@ def get_component_irreducibles(self, graph, swc_dict):
         else:
             return False
 
-    def prune_branches(self, graph):
-        """
-        Prunes all short branches from "graph". A short branch is a path
-        between a leaf and branching node where the path length is less than
-        "self.prune_depth".
-
-        Parameters
-        ----------
-        graph : networkx.Graph
-            Graph to be pruned.
-
-        Returns
-        -------
-        networkx.Graph
-            Graph with short branches pruned.
-
-        """
-        first_pass = True
-        deleted_nodes = list()
-        n_passes = 0
-        while len(deleted_nodes) > 0 or first_pass:
-            # Visit leafs
-            n_passes += 1
-            deleted_nodes = list()
-            for leaf in get_leafs(graph):
-                branch = [leaf]
-                length = 0
-                for (i, j) in nx.dfs_edges(graph, source=leaf):
-                    # Visit edge
-                    length += compute_dist(graph, i, j)
-                    if graph.degree(j) == 2:
-                        branch.append(j)
-                    elif graph.degree(j) > 2:
-                        deleted_nodes.extend(branch)
-                        graph.remove_nodes_from(branch)
-                        break
-
-                    # Check whether to stop
-                    if length > self.prune_depth or first_pass:
-                        graph.remove_nodes_from(branch[0:min(5, len(branch))])
-                        break
-
-            first_pass = False
-
 
 # --- Utils ---
 def get_irreducible_nodes(graph):

src/deep_neurographs/utils/swc_util.py

@@ -642,8 +642,8 @@ def __add_attributes(swc_dict, graph):
 
     """
     attrs = dict()
-    for idx, node_id in enumerate(swc_dict["id"]):
-        attrs[node_id] = {
+    for idx, node in enumerate(swc_dict["id"]):
+        attrs[node] = {
             "xyz": swc_dict["xyz"][idx],
             "radius": swc_dict["radius"][idx],
         }