refactor: updated feature generation

Author: anna-grim

src/neuron_proofreader/proposal_graph.py

@@ -388,25 +388,24 @@ def proposal_avg_radii(self, proposal):
     def proposal_directionals(self, proposal, depth):
         # Extract points along branches
         i, j = tuple(proposal)
-        xyz_list_i = self.truncated_edge_attr_xyz(i, depth)
-        xyz_list_j = self.truncated_edge_attr_xyz(j, depth)
-        origin = self.proposal_midpoint(proposal)
+        path_i = self.path_thru_node(i, depth)
+        path_j = self.path_thru_node(j, depth)
+        path_xyz_i = self.node_xyz[np.array(path_i)]
+        path_xyz_j = self.node_xyz[np.array(path_j)]
+        print("Path Lengths:", len(path_i), len(path_j))
 
         # Compute tangent vectors - branches
-        direction_i = geometry.get_directional(xyz_list_i, origin, depth)
-        direction_j = geometry.get_directional(xyz_list_j, origin, depth)
-        direction = geometry.tangent(self.proposal_attr(proposal, "xyz"))
-        if np.isnan(direction).any():
-            direction[0] = 0
-            direction[1] = 0
+        dir_i = geometry.tangent(path_xyz_i)
+        dir_j = geometry.tangent(path_xyz_j)
+        dir_proposal = geometry.tangent(self.proposal_attr(proposal, "xyz"))
 
         # Compute features
-        dot_i = abs(np.dot(direction, direction_i))
-        dot_j = abs(np.dot(direction, direction_j))
+        dot_i = abs(np.dot(dir_proposal, dir_i))
+        dot_j = abs(np.dot(dir_proposal, dir_j))
         if self.is_simple(proposal):
-            dot_ij = np.dot(direction_i, direction_j)
+            dot_ij = np.dot(dir_i, dir_j)
         else:
-            dot_ij = np.dot(direction_i, direction_j)
+            dot_ij = np.dot(dir_i, dir_j)
             if not self.is_simple(proposal):
                 dot_ij = max(dot_ij, -dot_ij)
         return np.array([dot_i, dot_j, dot_ij])
@@ -418,67 +417,7 @@ def proposal_midpoint(self, proposal):
         i, j = tuple(proposal)
         return geometry.midpoint(self.node_xyz[i], self.node_xyz[j])
 
-    def truncated_edge_attr_xyz(self, i, depth):
-        xyz_path_list = self.edge_attr(i, "xyz")
-        return [geometry.truncate_path(path, depth) for path in xyz_path_list]
-
     # --- Helpers ---
-    def node_attr(self, i, key):
-        if key == "xyz":
-            return self.node_xyz[i]
-        elif key == "radius":
-            return self.node_radius[i]
-        else:
-            return self.nodes[i][key]
-
-    def edge_attr(self, i, key="xyz", ignore=False):
-        """
-        Gets the edge attribute specified by "key" for all edges connected to
-        the given node.
-
-        Parameters
-        ----------
-        i : int
-            Node for which the edge attributes are to be retrieved.
-        key : str, optional
-            Key specifying the type of edge attribute to retrieve. The default
-            is "xyz".
-        ignore : bool, optional
-            If True, it will only consider direct neighbors of node "i". If
-            False, the method will follow add the edge attributes along the
-            path of chain-like connections from node "i" to its neighbors,
-            provided that the neighbor nodes have degree 2.
-
-        Returns
-        -------
-        List[numpy.ndarray]
-            Edge attribute specified by "key" for all edges connected to the
-            given node.
-        """
-        attrs = list()
-        for j in self.neighbors(i):
-            attr_ij = self.orient_edge_attr((i, j), i, key=key)
-            if not ignore:
-                root = i
-                while self.degree[j] == 2:
-                    k = [k for k in self.neighbors(j) if k != root][0]
-                    attr_jk = self.orient_edge_attr((j, k), j, key=key)
-                    if key == "xyz":
-                        attr_ij = np.vstack([attr_ij, attr_jk])
-                    else:
-                        attr_ij = np.concatenate((attr_ij, attr_jk))
-                    root = j
-                    j = k
-            attrs.append(attr_ij)
-        return attrs
-
-    def edge_length(self, edge):
-        xyz = self.edges[edge]["xyz"]
-        if len(xyz) < 2:
-            return 0.0
-        else:
-            return np.linalg.norm(xyz[1:] - xyz[:-1], axis=1).sum()
-
     def find_fragments_near_xyz(self, query_xyz, max_dist):
         hits = dict()
         xyz_list = geometry.query_ball(self.kdtree, query_xyz, max_dist)
@@ -506,13 +445,6 @@ def is_soma(self, i):
         """
         return self.node_component_id[i] in self.soma_ids
 
-    def orient_edge_attr(self, edge, i, key="xyz"):
-        node_attr = self.node_attr(i, key)
-        if (self.edges[edge][key][0] == node_attr).all():
-            return self.edges[edge][key]
-        else:
-            return np.flip(self.edges[edge][key], axis=0)
-
     def update_component_ids(self, component_id, root):
         """
         Updates the component_id of all nodes connected to "root".

src/neuron_proofreader/skeleton_graph.py

@@ -769,6 +769,36 @@ def path_length(self, path):
         else:
             return 0
 
+    def path_thru_node(self, i, max_depth=np.inf):
+        if self.degree[i] == 1:
+            return self.path_from_leaf(i, max_depth)
+        else:
+            assert self.degree[i] == 2
+            j, k = self.neighbors(i)
+            path_ij = self.directed_path(i, j, max_depth=max_depth)
+            path_ik = self.directed_path(i, k, max_depth=max_depth)
+            return path_ij[::-1] + path_ik[1:]
+
+    def directed_path(self, start_node, next_node, max_depth=np.inf):
+        queue = [(next_node, 0)]
+        visited = [start_node, next_node]
+        path = list()
+        while queue:
+            # Visit node
+            i, dist_i = queue.pop()
+            if self.degree[i] != 2:
+                return path
+            else:
+                path.append(i)
+
+            # Update queue
+            for j in self.neighbors(i):
+                dist_j = dist_i + self.dist(i, j)
+                if dist_j < max_depth and j not in visited:
+                    queue.append((j, dist_j))
+                    visited.append(j)
+        return visited
+
     def rooted_subgraph(self, root, radius):
         """
         Gets a rooted subgraph with the given radius (in microns).

src/neuron_proofreader/split_proofreading/split_feature_extraction.py

@@ -9,7 +9,8 @@
 
 """
 
-from concurrent.futures import ThreadPoolExecutor, as_completed
+from concurrent.futures import as_completed, ThreadPoolExecutor
+from scipy.spatial import KDTree
 from skimage.transform import resize
 from torch_geometric.data import HeteroData
 
@@ -98,9 +99,7 @@ def __init__(self, graph):
         """
         # Instance attributes
         self.graph = graph
-
-        # Build KD-tree from leaf nodes
-        self.graph.set_kdtree(node_type="leaf")
+        self.kdtree = KDTree(graph.node_xyz[np.array(graph.leaf_nodes())])
 
     def __call__(self, subgraph, features):
         """
@@ -398,9 +397,10 @@ def __init__(
         self.patch_shape = patch_shape
 
         # Annotate mask
-        node1, node2 = tuple(self.proposal)
-        self.annotate_edge(node1)
-        self.annotate_edge(node2)
+        i, j = self.proposal
+        self.voxels = {u: self.get_branch_voxels(u) for u in [i, j]}
+        self.annotate_edge(i)
+        self.annotate_edge(j)
         self.annotate_proposal()
 
     # --- Core Routines ---
@@ -462,25 +462,8 @@ def get_branch_profile(self, node):
         profile : numpy.ndarray
             Intensity profile along the branch containing the given node.
         """
-        def check_emptiness():
-            """
-            Checks if voxels is empty.
-            """
-            if len(voxels) < 2:
-                voxels.append(self.graph.get_local_voxel(node, self.offset))
-
-        # Get branch voxel coordinates
-        voxels = self.get_branch_voxels(node)
-        voxels = geometry_util.make_voxels_connected(voxels)
-        voxels = img_util.get_contained_voxels(voxels, self.mask.shape)
-        check_emptiness()
-
-        # Resample voxels
-        voxels = np.array(voxels)
-        voxels = geometry_util.resample_curve_3d(voxels, 16).astype(int)
-        voxels = img_util.get_contained_voxels(voxels, self.mask.shape)
-        check_emptiness()
-        return self._extract_profile(voxels)
+        profile = self._extract_profile(self.voxels[node])
+        return geometry_util.resample_curve_1d(profile, 16)
 
     def _extract_profile(self, voxels):
         """
@@ -513,9 +496,7 @@ def annotate_edge(self, node):
         node : int
             Node ID used to get branch to be annotated.
         """
-        voxels = self.get_branch_voxels(node)
-        voxels = geometry_util.make_voxels_connected(voxels)
-        img_util.annotate_voxels(self.mask, voxels, val=0.5)
+        img_util.annotate_voxels(self.mask, self.voxels[node], val=0.5)
 
     def annotate_proposal(self):
         """
@@ -540,8 +521,8 @@ def get_profile_line(self, n_pts=None):
             Voxel line between the two nodes of a proposal.
         """
         node1, node2 = self.proposal
-        voxel1 = self.graph.get_local_voxel(node1, self.offset)
-        voxel2 = self.graph.get_local_voxel(node2, self.offset)
+        voxel1 = self.graph.node_local_voxel(node1, self.offset)
+        voxel2 = self.graph.node_local_voxel(node2, self.offset)
         if n_pts:
             return geometry_util.make_line(voxel1, voxel2, n_pts)
         else:
@@ -563,10 +544,22 @@ def get_branch_voxels(self, node):
             Voxel coordinates representing the edge path in local patch
             coordinates.
         """
-        pts = np.vstack(self.graph.edge_attr(node, "xyz"))
-        anisotropy = self.graph.anisotropy
-        voxels = [img_util.to_voxels(xyz, anisotropy) for xyz in pts]
-        return geometry_util.shift_path(voxels, self.offset)
+        queue = [(node, self.graph.node_local_voxel(node, self.offset))]
+        visited = {node}
+        voxels = list()
+        while queue:
+            # Visit node
+            i, voxel_i = queue.pop()
+            voxels.append(voxel_i)
+
+            # Update queue
+            for j in self.graph.neighbors(i):
+                voxel_j = self.graph.node_local_voxel(j, self.offset)
+                contained_j = img_util.is_contained(voxel_j, self.patch_shape)
+                if contained_j and j not in visited:
+                    queue.append((j, voxel_j))
+                    visited.add(j)
+        return geometry_util.make_voxels_connected(voxels)
 
 
 # --- Feature Data Structures ---
@@ -929,19 +922,22 @@ def __init__(self, object_ids):
 
 
 # --- Helpers ---
-def check_list_length(my_list, min_length=2):
+def check_list_length(arr, min_length=2):
     """
-    Checks that the list contains at least "min_length" items.
+    Checks that the array contains at least "min_length" items.
 
     Parameters
     ----------
-    my_list : list
-        List to be checked.
+    arr : list
+        Array to be checked.
     min_length : int
-        Minimum items that must be contained in the list
+        Minimum length of the array.
     """
-    while len(my_list) < min_length:
-        my_list.append(my_list[-1])
+    if arr.shape[0] < min_length:
+        pad_size = min_length - arr.shape[0]
+        padding = np.repeat(arr[-1:], pad_size, axis=0)
+        arr = np.concatenate([arr, padding], axis=0)
+    return arr
 
 
 def get_feature_dict():