diff --git a/pyproject.toml b/pyproject.toml index efea388..46d351c 100644 --- a/pyproject.toml +++ b/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'", diff --git a/tinynav/core/math_utils.py b/tinynav/core/math_utils.py index e5bf697..39193ae 100644 --- a/tinynav/core/math_utils.py +++ b/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 diff --git a/tinynav/core/perception_node.py b/tinynav/core/perception_node.py index 1ead4e2..b77cdad 100644 --- a/tinynav/core/perception_node.py +++ b/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 diff --git a/uv.lock b/uv.lock index 5457508..943f1dd 100644 --- a/uv.lock +++ b/uv.lock @@ -1008,6 +1008,19 @@ http = [ { name = "aiohttp" }, ] +[[package]] +name = "fufpy" +version = "0.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numba" }, + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5f/26/d90fcf23c5cd5108403bf0f45c9b48df2b1574d7e44cc53fb2c132e3eb1e/fufpy-0.1.1.tar.gz", hash = "sha256:c05d336b3f2484170b08d83971d901d6a38d3f5e351f55ecff3061dd3bab99d3", size = 4631, upload-time = "2025-02-28T19:35:03.955Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/27/2f50a21216f50ebc84a77e286002e54bd1c515a79404d44810de4f95ac03/fufpy-0.1.1-py3-none-any.whl", hash = "sha256:e256b5426860b30b7ad592f7377abd1e63a999cf9642e4a80761bc6aebe1ecd5", size = 5116, upload-time = "2025-02-28T19:35:02.165Z" }, +] + [[package]] name = "gdown" version = "5.2.0" @@ -3996,6 +4009,7 @@ dependencies = [ { name = "cuda-python", version = "12.2.0", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'x86_64'" }, { name = "cuda-python", version = "12.6.0", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'aarch64'" }, { name = "einops" }, + { name = "fufpy" }, { name = "huggingface-hub" }, { name = "jupyterlab" }, { name = "matplotlib" }, @@ -4040,6 +4054,7 @@ requires-dist = [ { name = "cuda-python", marker = "platform_machine == 'aarch64'", specifier = "==12.6.0" }, { name = "cuda-python", marker = "platform_machine == 'x86_64'", specifier = "==12.2.0" }, { name = "einops" }, + { name = "fufpy", specifier = ">=0.1.1" }, { name = "huggingface-hub" }, { name = "jupyterlab" }, { name = "lerobot", extras = ["lekiwi"], marker = "extra == 'lekiwi'", specifier = "==0.3.3" },