refactor: updated doubles detection alg

src/deep_neurographs/doubles_removal.py

@@ -9,7 +9,9 @@
 """
 
 import networkx as nx
+import numpy as np
 
+from deep_neurographs import geometry
 from deep_neurographs import utils
 
 
@@ -46,11 +48,11 @@ def run(neurograph, min_size, max_size, node_spacing, output_dir=None):
     doubles_cnt = 0
     neurograph.init_kdtree()
     not_doubles = set()
-    for i, nodes in enumerate(components):
+    for i, idx in enumerate(np.argsort([len(c) for c in components])):
         # Determine whether to inspect fragment
-        swc_id = get_swc_id(neurograph, nodes)
+        swc_id = get_swc_id(neurograph, components[idx])
         if swc_id not in not_doubles:
-            xyz_arr = inspect_component(neurograph, nodes)
+            xyz_arr = inspect_component(neurograph, components[idx])
             upper = len(xyz_arr) * node_spacing < max_size
             lower = len(xyz_arr) * node_spacing > min_size
             if upper and lower:
@@ -59,13 +61,16 @@ def run(neurograph, min_size, max_size, node_spacing, output_dir=None):
                     doubles_cnt += 1
                     if output_dir:
                         neurograph.to_swc(
-                            output_dir, nodes, color="1.0 0.0 0.0"
+                            output_dir, components[idx], color="1.0 0.0 0.0"
                         )
-                    neurograph = remove_component(neurograph, nodes, swc_id)
                     not_doubles.add(not_double_id)
+                    neurograph = remove_component(
+                        neurograph, components[idx], swc_id
+                    )
+
 
         # Update progress bar
-        if i > cnt * chunk_size:
+        if i >= cnt * chunk_size:
             cnt, t1 = utils.report_progress(
                 i + 1, n_components, chunk_size, cnt, t0, t1
             )
@@ -99,21 +104,32 @@ def is_double(neurograph, fragment, swc_id_i):
     # Compute projections
     hits = dict()
     for xyz_i in fragment:
-        for xyz_j in neurograph.query_kdtree(xyz_i, 5):
+        hits_i = dict()
+        for xyz_j in neurograph.query_kdtree(xyz_i, 25):
             try:
                 swc_id_j = neurograph.xyz_to_swc(xyz_j)
                 if swc_id_i != swc_id_j:
-                    hits = utils.append_dict_value(hits, swc_id_j, 1)
+                    d = geometry.dist(xyz_i, xyz_j)
+                    hits_i = check_hits(hits_i, swc_id_j, d)
             except:
                 pass
+        if len(hits_i) > 0:
+            best_swc_id = utils.find_best(hits_i)
+            best_dist = hits_i[best_swc_id]
+            hits = utils.append_dict_value(hits, best_swc_id, best_dist)
 
     # Check criteria
-    if len(hits) > 0:
-        swc_id_j = utils.find_best(hits)
-        percent_hit = len(hits[swc_id_j]) / len(fragment)
-    else:
-        percent_hit = 0
-    return swc_id_j if swc_id_j is not None and percent_hit > 0.5 else None
+    for swc_id_j, dists in hits.items():
+        median = np.median(dists)
+        percent_hit = len(dists) / len(fragment)
+        std = np.std(dists)
+        if percent_hit > 0.4 and (median < 20 and std < 2):
+            return swc_id_j
+        elif percent_hit > 0.3 and (median < 15 and std < 1.5):
+            return swc_id_j
+        elif percent_hit > 0.6 and std < 1:
+            return swc_id_j
+    return False
 
 
 # --- utils ---
@@ -214,3 +230,11 @@ def remove_xyz_entries(neurograph, i, j):
     for xyz in neurograph.edges[i, j]["xyz"]:
         del neurograph.xyz_to_edge[tuple(xyz)]
     return neurograph
+
+def check_hits(hits, key, value):
+    if key in hits.keys():
+        if value < hits[key]:
+            hits[key] = value
+    else:
+        hits[key] = value
+    return hits

src/deep_neurographs/machine_learning/inference.py

@@ -83,7 +83,7 @@ def run(
     progress_cnt = 1
     t0, t1 = utils.init_timers()
     chunk_size = max(int(n_batches * 0.02), 1)
-    for i, batch in enumerate(batches):
+    for i, batch in enumerate(batches):        
         # Prediction
         proposals_i = [proposals[j] for j in batch]
         accepts_i, graph = predict(
@@ -103,7 +103,7 @@ def run(
         accepts.extend(accepts_i)
 
         # Report progress
-        if i > progress_cnt * chunk_size:
+        if i >= progress_cnt * chunk_size:
             progress_cnt, t1 = utils.report_progress(
                 i + 1, n_batches, chunk_size, progress_cnt, t0, t1
             )
@@ -142,7 +142,7 @@ def predict(
         proposal_probs,
         idx_to_edge,
         search_radius,
-        high_threshold=0.95,
+        high_threshold=0.9,
         threshold=confidence_threshold,
     )
     return accepts, graph

src/deep_neurographs/neurograph.py

@@ -606,10 +606,11 @@ def upd_ids(self, swc_id, r):
 
         """
         queue = [r]
-        visited = []
+        visited = set()
         while len(queue) > 0:
             i = queue.pop()
             self.nodes[i]["swc_id"] = swc_id
+            visited.add(i)
             for j in [j for j in self.neighbors(i) if j not in visited]:
                 queue.append(j)
 

src/deep_neurographs/utils.py

@@ -163,7 +163,7 @@ def find_best(my_dict):
     best_vote_cnt = 0
     for key in my_dict.keys():
         val_type = type(my_dict[key])
-        vote_cnt = my_dict[key] if val_type == int else len(my_dict[key])
+        vote_cnt = my_dict[key] if val_type == float else len(my_dict[key])
         if vote_cnt > best_vote_cnt:
             best_key = key
             best_vote_cnt = vote_cnt