bug: fixed missing fragments in merge datasets (#561)

Co-authored-by: anna-grim <[email protected]>

Author: anna-grim

src/neuron_proofreader/merge_proofreading/merge_dataloading.py

@@ -14,7 +14,7 @@
 import numpy as np
 import pandas as pd
 
-TEST_BRAIN = "685221"
+TEST_BRAIN = "653159"
 
 
 # --- Load Skeletons ---

src/neuron_proofreader/merge_proofreading/merge_datasets.py

@@ -32,7 +32,6 @@
     geometry_util,
     img_util,
     ml_util,
-    swc_util,
     util,
 )
 
@@ -229,16 +228,19 @@ def remove_nonindexed_fragments(self, idxs):
             other sites are removed.
         """
         # Remove other fragments
-        for idx in [i for i in self.merge_sites_df.index if i not in idxs]:
+        visited = set()
+        for i in [i for i in self.merge_sites_df.index if i not in idxs]:
             # Extract site info
-            brain_id = self.merge_sites_df["brain_id"][idx]
-            xyz = self.merge_sites_df["xyz"][idx]
+            brain_id = self.merge_sites_df["brain_id"][i]
+            segment_id = self.merge_sites_df["segment_id"][i]
+            pair = (brain_id, segment_id)
 
             # Find fragment containing site
-            dist, node = self.graphs[brain_id].kdtree.query(xyz)
-            if dist < 20 and node in self.graphs[brain_id]:
-                nodes = self.graphs[brain_id].get_connected_nodes(node)
+            if pair not in visited:
+                nodes = self.graphs[brain_id].get_nodes_with_segment_id(segment_id)
                 self.graphs[brain_id].remove_nodes(nodes, False)
+                visited.add(pair)
+
         self.remove_empty_graphs()
 
         # Relabel nodes
@@ -254,7 +256,7 @@ def remove_empty_graphs(self):
         Removes graphs without any nodes.
         """
         for brain_id in list(self.graphs.keys()):
-            if len(self.graphs[brain_id]) == 0:
+            if len(self.graphs[brain_id].nodes) == 0:
                 del self.graphs[brain_id]
 
     # --- Getters ---

src/neuron_proofreader/skeleton_graph.py

@@ -376,6 +376,24 @@ 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):
+        nodes = set()
+        query_id = f"{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)
+                nodes = nodes.union(
+                    self.get_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 get_rooted_subgraph(self, root, radius):
         """
         Gets a rooted subgraph with the given radius (in microns).

src/neuron_proofreader/utils/geometry_util.py

@@ -9,7 +9,7 @@
 """
 
 from collections import defaultdict
-from scipy.interpolate import splprep, splev, UnivariateSpline
+from scipy.interpolate import UnivariateSpline
 from scipy.linalg import svd
 from scipy.spatial import distance
 from tqdm import tqdm