refactor: updated proposal trimming

Author: anna-grim

src/neuron_proofreader/merge_proofreading/merge_datasets.py

@@ -317,7 +317,7 @@ def __getitem__(self, idx):
         """
         # Get example
         brain_id, subgraph, label = self.get_site(idx)
-        voxel = subgraph.get_voxel(0)
+        voxel = subgraph.node_voxel(0)
 
         # Extract subgraph and image patches centered at site
         img_patch = self.get_img_patch(brain_id, voxel)
@@ -509,7 +509,7 @@ def get_segment_mask(self, brain_id, center, subgraph):
             segment_mask = np.zeros(self.patch_shape)
 
         # Annotate fragment
-        center = subgraph.get_voxel(0)
+        center = subgraph.node_voxel(0)
         offset = img_util.get_offset(center, self.patch_shape)
         for node1, node2 in subgraph.edges:
             # Get local voxel coordinates

src/neuron_proofreader/merge_proofreading/merge_inference.py

@@ -251,7 +251,7 @@ def __iter__(self):
 
         # Search graph
         visited_ids = set()
-        for u in self.graph.get_leafs():
+        for u in self.graph.leaf_nodes():
             component_id = self.graph.node_component_id[u]
             if component_id not in visited_ids and component_id in valid_ids:
                 visited_ids.add(component_id)
@@ -288,7 +288,7 @@ def find_fragments_to_search(self):
         return component_ids
 
     def get_patch_centers(self, nodes):
-        patch_centers = [self.graph.get_voxel(i) for i in nodes]
+        patch_centers = [self.graph.node_voxel(i) for i in nodes]
         return np.array(patch_centers, dtype=int)
 
     def get_label_mask(self, nodes, img_shape, offset):
@@ -304,8 +304,8 @@ def get_label_mask(self, nodes, img_shape, offset):
         # Annotate mask
         subgraph = self.get_contained_subgraph(nodes, img_shape, offset)
         for i, j in subgraph.edges:
-            voxel_i = self.graph.get_voxel(i) - offset
-            voxel_j = self.graph.get_voxel(j) - offset
+            voxel_i = self.graph.node_voxel(i) - offset
+            voxel_j = self.graph.node_voxel(j) - offset
             voxels = geometry_util.make_digital_line(voxel_i, voxel_j)
             img_util.annotate_voxels(segment_mask, voxels)
         return segment_mask
@@ -317,13 +317,13 @@ def get_contained_subgraph(self, nodes, img_shape, offset):
         while queue:
             # Visit node
             i = queue.pop()
-            voxel_i = self.graph.get_voxel(i) - offset
+            voxel_i = self.graph.node_voxel(i) - offset
             if not img_util.is_contained(voxel_i, img_shape, buffer=1):
                 continue
 
             # Update queue
             for j in self.graph.neighbors(i):
-                voxel_j = self.graph.get_voxel(j) - offset
+                voxel_j = self.graph.node_voxel(j) - offset
                 if img_util.is_contained(voxel_j, img_shape):
                     subgraph.add_edge(i, j)
                     if j not in visited:
@@ -332,7 +332,7 @@ def get_contained_subgraph(self, nodes, img_shape, offset):
         return subgraph
 
     def is_contained(self, node):
-        voxel = self.graph.get_voxel(node)
+        voxel = self.graph.node_voxel(node)
         shape = self.img_reader.shape()[2::]
         buffer = np.max(self.patch_shape) + 1
         return img_util.is_contained(voxel, shape, buffer=buffer)
@@ -601,7 +601,7 @@ def _generate_batch_nodes(self, root):
             if len(patch_centers) == 0 and self.graph.degree[i] > 2:
                 root = i
                 nodes.append(i)
-                patch_centers.append(self.graph.get_voxel(i))
+                patch_centers.append(self.graph.node_voxel(i))
 
             # Check whether to yield batch
             is_node_far = self.graph.dist(root, j) > 256
@@ -619,7 +619,7 @@ def _generate_batch_nodes(self, root):
             # Visit j
             if self.graph.degree[j] > 2:
                 nodes.append(j)
-                patch_centers.append(self.graph.get_voxel(j))
+                patch_centers.append(self.graph.node_voxel(j))
                 if len(patch_centers) == 1:
                     root = j
 

src/neuron_proofreader/new_proposal_graph.py

@@ -104,9 +104,8 @@ def __init__(
         self.merged_ids = set()
         self.n_merges_blocked = 0
         self.n_proposals_blocked = 0
-        self.node_proposals = defaultdict(set)
-        self.proposals = set()
 
+        self.reset_proposals()
         self.proposal_generator = ProposalGenerator(
             self,
             max_proposals_per_leaf=max_proposals_per_leaf,
@@ -163,6 +162,17 @@ def connect_soma_fragments(self, soma_centroids):
         results.append(f"# Soma Fragments Merged: {merge_cnt}")
         return "\n".join(results)
 
+    def relabel_nodes(self):
+        # Call parent class
+        old_proposals = self.list_proposals()
+        old_to_new = super().relabel_nodes()
+
+        # Update proposals
+        self.reset_proposals()
+        for proposal in old_proposals:
+            i, j = proposal
+            self.add_proposal(int(old_to_new[i]), int(old_to_new[j]))
+
     # --- Proposal Operations ---
     def add_proposal(self, i, j):
         """
@@ -198,15 +208,16 @@ def generate_proposals(self, search_radius):
         # Proposal generation
         proposals = self.proposal_generator(search_radius)
         self.store_proposals(proposals)
-        #self.trim_proposals() TEMP
+        self.trim_proposals()
+        self.relabel_nodes()
 
         # Set groundtruth
         if self.gt_path:
             gt_graph = SkeletonGraph(anisotropy=self.anisotropy)
             gt_graph.load(self.gt_path)
             self.gt_accepts = groundtruth_generation.run(gt_graph, self)
 
-    def get_sorted_proposals(self):
+    def sorted_proposals(self):
         """
         Return proposals sorted by physical length.
 
@@ -324,6 +335,10 @@ def remove_proposal(self, proposal):
         self.node_proposals[j].remove(i)
         self.proposals.remove(proposal)
 
+    def reset_proposals(self):
+        self.node_proposals = defaultdict(set)
+        self.proposals = set()
+
     def store_proposals(self, proposals):
         self.node_proposals = defaultdict(set)
         for proposal in proposals:

src/neuron_proofreader/skeleton_graph.py

@@ -225,7 +225,7 @@ def reassign_component_ids(self):
         component_id_to_swc_id = dict()
         for i, nodes in enumerate(nx.connected_components(self)):
             nodes = np.array(list(nodes), dtype=int)
-            component_id_to_swc_id[i + 1] = self.get_swc_id(nodes[0])
+            component_id_to_swc_id[i + 1] = self.node_swc_id(nodes[0])
             self.node_component_id[nodes] = i + 1
         self.component_id_to_swc_id = component_id_to_swc_id
 
@@ -240,12 +240,11 @@ def relabel_nodes(self):
         # Set edge ids
         old_to_new = dict(zip(old_node_ids, new_node_ids))
         old_edge_ids = list(self.edges)
-        edge_attrs = {(i, j): data for i, j, data in self.edges(data=True)}
 
         # Reset graph
         self.clear()
         for (i, j) in old_edge_ids:
-            self.add_edge(old_to_new[i], old_to_new[j], **edge_attrs[(i, j)])
+            self.add_edge(old_to_new[i], old_to_new[j])
 
         # Update attributes
         self.node_radius = self.node_radius[old_node_ids]
@@ -255,6 +254,7 @@ def relabel_nodes(self):
         self.reassign_component_ids()
         if self.kdtree:
             self.set_kdtree()
+        return old_to_new
 
     def remove_nodes(self, nodes, relabel_nodes=True):
         """
@@ -272,7 +272,7 @@ def remove_nodes(self, nodes, relabel_nodes=True):
             self.relabel_nodes()
 
     # --- Getters ---
-    def get_branchings(self):
+    def branching_nodes(self):
         """
         Gets all branching nodes in the graph.
 
@@ -283,7 +283,13 @@ def get_branchings(self):
         """
         return [i for i in self.nodes if self.degree[i] > 2]
 
-    def get_connected_nodes(self, root):
+    def component_id_from_swc_id(self, query_swc_id):
+        for component_id, swc_id in self.component_id_to_swc_id.items():
+            if query_swc_id == swc_id:
+                return component_id
+        raise ValueError(f"SWC ID={query_swc_id} not found")
+
+    def connected_nodes(self, root):
         """
         Gets all nodes connected to the given root node.
 
@@ -307,7 +313,7 @@ def get_connected_nodes(self, root):
                     visited.add(j)
         return visited
 
-    def get_leafs(self):
+    def leaf_nodes(self):
         """
         Gets all leaf nodes in the graph.
 
@@ -318,23 +324,7 @@ def get_leafs(self):
         """
         return [i for i in self.nodes if self.degree[i] == 1]
 
-    def get_voxel(self, i):
-        """
-        Gets the voxel coordinate of the given node.
-
-        Parameters
-        ----------
-        i : int
-            Node ID.
-
-        Returns
-        -------
-        float
-            Voxel coordinate of the given node.
-        """
-        return img_util.to_voxels(self.node_xyz[i], self.anisotropy)
-
-    def get_local_voxel(self, node, offset):
+    def node_local_voxel(self, node, offset):
         """
         Computes the local voxel coordinate of the given node within the image
         patch defined by "center" and "patch_shape".
@@ -352,10 +342,9 @@ def get_local_voxel(self, node, offset):
             Local voxel coordinate of the given node within the image patch
             defined by "center" and "patch_shape".
         """
-        voxel = self.get_voxel(node)
-        return tuple([v - o for v, o in zip(voxel, offset)])
+        return tuple([v - o for v, o in zip(self.node_voxel(node), offset)])
 
-    def get_node_segment_id(self, node):
+    def node_segment_id(self, node):
         """
         Gets the segment ID corresponding to the given node.
 
@@ -371,7 +360,40 @@ def get_node_segment_id(self, node):
         """
         return self.get_swc_id(node).split(".")[0]
 
-    def get_nodes_with_component_id(self, component_id):
+    def node_swc_id(self, i):
+        """
+        Gets the SWC ID of the given node.
+
+        Parameters
+        ----------
+        i : int
+            Node ID.
+
+        Returns
+        -------
+        str
+            SWC ID of the given node.
+        """
+        component_id = self.node_component_id[i]
+        return self.component_id_to_swc_id[component_id]
+
+    def node_voxel(self, i):
+        """
+        Gets the voxel coordinate of the given node.
+
+        Parameters
+        ----------
+        i : int
+            Node ID.
+
+        Returns
+        -------
+        float
+            Voxel coordinate of the given node.
+        """
+        return img_util.to_voxels(self.node_xyz[i], self.anisotropy)
+
+    def nodes_with_component_id(self, component_id):
         """
         Gets all nodes with the given componenet ID.
 
@@ -387,7 +409,7 @@ def get_nodes_with_component_id(self, component_id):
         """
         return set(np.where(self.node_component_id == component_id)[0])
 
-    def get_nodes_with_segment_id(self, segment_id):
+    def nodes_with_segment_id(self, segment_id):
         """
         Gets all nodes with the given segment ID.
 
@@ -406,19 +428,13 @@ def get_nodes_with_segment_id(self, segment_id):
         for swc_id in self.get_swc_ids():
             segment_id = int(swc_id.replace(".0", ""))
             if segment_id == query_id:
-                component_id = self.get_component_id_from_swc_id(swc_id)
+                component_id = self.component_id_from_swc_id(swc_id)
                 nodes = nodes.union(
-                    self.get_nodes_with_component_id(component_id)
+                    self.nodes_with_component_id(component_id)
                 )
         return nodes
 
-    def get_component_id_from_swc_id(self, query_swc_id):
-        for component_id, swc_id in self.component_id_to_swc_id.items():
-            if query_swc_id == swc_id:
-                return component_id
-        raise ValueError(f"SWC ID={query_swc_id} not found")
-
-    def nodes_within_distance(self, root, radius):
+    def nodes_within_distance(self, root, max_depth):
         queue = [(root, 0)]
         visited = {root}
         while queue:
@@ -428,11 +444,28 @@ def nodes_within_distance(self, root, radius):
             # Populate queue
             for j in self.neighbors(i):
                 dist_j = dist_i + self.dist(i, j)
-                if dist_j < radius and j not in visited:
+                if dist_j < max_depth and j not in visited:
                     queue.append((j, dist_j))
                     visited.add(j)
         return list(visited)
 
+    def path_from_leaf(self, leaf, max_depth=np.inf):
+        queue = [(leaf, 0)]
+        path = [leaf]
+        while queue:
+            # Visit node
+            i, dist_i = queue.pop()
+            if self.degree[i] != 2 and dist_i > 0:
+                return path
+
+            # 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 path:
+                    queue.append((j, dist_j))
+                    path.append(j)
+        return path
+
     def rooted_subgraph(self, root, radius):
         """
         Gets a rooted subgraph with the given radius (in microns).
@@ -478,23 +511,6 @@ def rooted_subgraph(self, root, radius):
         subgraph.node_xyz = self.node_xyz[idxs]
         return subgraph
 
-    def get_swc_id(self, i):
-        """
-        Gets the SWC ID of the given node.
-
-        Parameters
-        ----------
-        i : int
-            Node ID.
-
-        Returns
-        -------
-        str
-            SWC ID of the given node.
-        """
-        component_id = self.node_component_id[i]
-        return self.component_id_to_swc_id[component_id]
-
     def get_swc_ids(self):
         """
         Gets the set of all unique SWC IDs of nodes in the graph.
@@ -581,7 +597,7 @@ def write_entry(node, parent):
                 write_entry(j, i)
 
             # Finish
-            filename = self.get_swc_id(i)
+            filename = self.node_swc_id(i)
             filename = util.set_zip_path(zip_writer, filename, ".swc")
             zip_writer.writestr(filename, text_buffer.getvalue())
 
@@ -625,7 +641,7 @@ def clip_to_bbox(self, metadata_path):
             # Clip graph
             nodes = list()
             for i in self.nodes:
-                voxel = np.array(self.get_voxel(i))
+                voxel = np.array(self.node_voxel(i))
                 if not img_util.is_contained(voxel - origin, shape):
                     nodes.append(i)
             self.remove_nodes_from(nodes)