perf(perception): stabilize union-find tracks by importing fufpy

pyproject.toml

@@ -17,6 +17,7 @@ dependencies = [
     "jupyterlab",
     "codetiming",
     "numba",
+    "fufpy",
     "pygame",
     "cuda_python==12.6.0; platform_machine == 'aarch64'",
     "cuda-python==12.2.0; platform_machine == 'x86_64'",

tinynav/core/math_utils.py

@@ -4,6 +4,7 @@
 from geometry_msgs.msg import TransformStamped
 from nav_msgs.msg import Odometry
 import cv2
+import fufpy
 from tinynav.core.func import lru_cache_numpy
 
 @njit(cache=True)
@@ -237,49 +238,21 @@ def estimate_pose(kpts_prev, kpts_curr, depth, K, idx_valid=None):
     inlier_idx_original = idx_valid[inliers]
     return True, T, inliers_2d, inliers_3d, inlier_idx_original
 
-# Disjoint Set (Union-Find) implementation with path compression and union by rank
-@njit(cache=True)
+# Union–find via fufpy (https://github.com/LuisScoccola/fufpy)
 def uf_init(n):
-    parent = np.empty(n, np.int64)
-    rank = np.zeros(n, np.int64)
-    for i in range(n):
-        parent[i] = i
-    return parent, rank
+    return fufpy.dynamic_partition_create(int(n))
 
-@njit(cache=True)
-def uf_find(i, parent):
-    root = i
-    while parent[root] != root:
-        root = parent[root]
-    while parent[i] != i:
-        p = parent[i]
-        parent[i] = root
-        i = p
-    return root
 
-@njit(cache=True)
-def uf_union(a, b, parent, rank):
-    ra = uf_find(a, parent)
-    rb = uf_find(b, parent)
-    if ra == rb:
-        return ra
-    if rank[ra] < rank[rb]:
-        parent[ra] = rb
-        return rb
-    elif rank[ra] > rank[rb]:
-        parent[rb] = ra
-        return ra
-    else:
-        parent[rb] = ra
-        rank[ra] += 1
-        return ra
+def uf_union(a, b, uf, _rank=None):
+    return fufpy.dynamic_partition_union(uf, int(a), int(b))
+
 
-def uf_all_sets_list(parent):
-    root_to_members = {}
-    for i in range(len(parent)):
-        r = parent[i]
-        root_to_members.setdefault(r, []).append(i)
-    return list(root_to_members.values())
+def uf_all_sets_list(uf, min_component_size=1):
+    out = []
+    for part in fufpy.dynamic_partition_parts(uf):
+        if part.size >= int(min_component_size):
+            out.append(np.sort(part).tolist())
+    return out
 
 
 

tinynav/core/perception_node.py

@@ -338,7 +338,7 @@ async def process(self, left_msg, right_msg):
 
             with Timer(name="[init extract info]", text="[{name}] Elapsed time: {milliseconds:.0f} ms", logger=self.logger.debug):
                 extract_info = [await self.superpoint.infer(kf.image) for kf in self.keyframe_queue[-_N:]]
-            parent, rank = uf_init(len(self.keyframe_queue[-_N:]) * _M)
+                uf = uf_init(len(self.keyframe_queue[-_N:]) * _M)
 
             self.logger.debug(f"Processing {len(self.keyframe_queue)} keyframes for data association.")
 
@@ -408,12 +408,12 @@ async def process(self, left_msg, right_msg):
                             if match_idx != -1:
                                 idx_prev = i * _M + k
                                 idx_curr = j * _M + match_idx
-                                uf_union(idx_prev, idx_curr, parent, rank)
+                                uf_union(idx_prev, idx_curr, uf)
                                 count += 1
                         self.logger.debug(f"{i} match {j} after Pnp filter count: {count}")
 
             with Timer(name="[found track]", text="[{name}] Elapsed time: {milliseconds:.0f} ms", logger=self.logger.debug):
-                tracks = [track for track in uf_all_sets_list(parent) if len(track) >= 2]
+                tracks = uf_all_sets_list(uf, min_component_size=2)
                 self.logger.debug(f"Found {len(tracks)} tracks after data association.")
 
             with Timer(name="[add track]", text="[{name}] Elapsed time: {milliseconds:.0f} ms", logger=self.logger.debug):
@@ -450,7 +450,7 @@ async def process(self, left_msg, right_msg):
                         )
                         smart_factor.add(stereo_meas, X(pose_idx), calib)
                     graph.add(smart_factor)
-            
+
         with Timer(name="[Solver]", text="[{name}] Elapsed time: {milliseconds:.0f} ms", logger=self.logger.debug):
             params = gtsam.LevenbergMarquardtParams()
             # set iteration limit