From 572e57806948260d6594ffa5ea23a86e4ac86317 Mon Sep 17 00:00:00 2001 From: sl-perception Date: Mon, 2 Mar 2026 10:15:32 -0500 Subject: [PATCH] =?UTF-8?q?feat(vo):=20visual=20odometry=20fallback=20?= =?UTF-8?q?=E2=80=94=20CUDA=20optical=20flow=20+=20EKF=20fusion=20+=20slip?= =?UTF-8?q?=20failover=20(Issue=20#157)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New package: saltybot_visual_odom (13 files, ~900 lines) Nodes: visual_odom_node — D435i IR1 stereo VO at 30 Hz CUDA: SparsePyrLKOpticalFlow + FastFeatureDetector (GPU) CPU fallback: calcOpticalFlowPyrLK + goodFeaturesToTrack Essential matrix (5-pt RANSAC) + depth-aided metric scale forward-backward consistency check on tracked points Publishes /saltybot/visual_odom (Odometry) odom_fusion_node — 5-state EKF [px, py, θ, v, ω] (unicycle model) Fuses: wheel odom (/saltybot/rover_odom or tank_odom) + visual odom (/saltybot/visual_odom) Slip failover: /saltybot/terrain JSON → 10× wheel noise on slip Loop closure: /rtabmap/odom jump > 0.3m → EKF soft-correct TF: publishes odom → base_link at 30 Hz Publishes /saltybot/odom_fused + /saltybot/visual_odom_status Modules: optical_flow_tracker.py — CUDA/CPU sparse LK tracker with re-detection, forward-backward consistency, ROI masking stereo_vo.py — Essential matrix decomposition, camera→base_link frame rotation, depth median scale recovery, loop closure soft-correct, accumulated SE(3) pose kalman_odom_filter.py — 5-state EKF: predict (unicycle), update_wheel, update_visual, update_rtabmap (absolute pose); Joseph-form covariance for numerical stability Tests: test/test_kalman_odom.py — 8 unit tests for EKF + StereoVO (no ROS deps) Topic/TF map: /camera/infra1/image_rect_raw → visual_odom_node /camera/depth/image_rect_raw → visual_odom_node /saltybot/visual_odom ← visual_odom_node (30 Hz) /saltybot/rover_odom → odom_fusion_node /saltybot/terrain → odom_fusion_node (slip signal) /rtabmap/odom → odom_fusion_node (loop closure) /saltybot/odom_fused ← odom_fusion_node (30 Hz) odom → base_link TF ← odom_fusion_node (30 Hz) Co-Authored-By: Claude Sonnet 4.6 --- .../src/saltybot_visual_odom/CMakeLists.txt | 11 + .../config/visual_odom_params.yaml | 45 +++ .../launch/visual_odom.launch.py | 74 +++++ .../src/saltybot_visual_odom/package.xml | 41 +++ .../resource/saltybot_visual_odom | 0 .../saltybot_visual_odom/__init__.py | 0 .../kalman_odom_filter.py | 210 ++++++++++++ .../saltybot_visual_odom/odom_fusion_node.py | 304 ++++++++++++++++++ .../optical_flow_tracker.py | 207 ++++++++++++ .../saltybot_visual_odom/stereo_vo.py | 252 +++++++++++++++ .../saltybot_visual_odom/visual_odom_node.py | 200 ++++++++++++ .../ros2_ws/src/saltybot_visual_odom/setup.py | 26 ++ .../test/test_kalman_odom.py | 140 ++++++++ 13 files changed, 1510 insertions(+) create mode 100644 jetson/ros2_ws/src/saltybot_visual_odom/CMakeLists.txt create mode 100644 jetson/ros2_ws/src/saltybot_visual_odom/config/visual_odom_params.yaml create mode 100644 jetson/ros2_ws/src/saltybot_visual_odom/launch/visual_odom.launch.py create mode 100644 jetson/ros2_ws/src/saltybot_visual_odom/package.xml create mode 100644 jetson/ros2_ws/src/saltybot_visual_odom/resource/saltybot_visual_odom create mode 100644 jetson/ros2_ws/src/saltybot_visual_odom/saltybot_visual_odom/__init__.py create mode 100644 jetson/ros2_ws/src/saltybot_visual_odom/saltybot_visual_odom/kalman_odom_filter.py create mode 100644 jetson/ros2_ws/src/saltybot_visual_odom/saltybot_visual_odom/odom_fusion_node.py create mode 100644 jetson/ros2_ws/src/saltybot_visual_odom/saltybot_visual_odom/optical_flow_tracker.py create mode 100644 jetson/ros2_ws/src/saltybot_visual_odom/saltybot_visual_odom/stereo_vo.py create mode 100644 jetson/ros2_ws/src/saltybot_visual_odom/saltybot_visual_odom/visual_odom_node.py create mode 100644 jetson/ros2_ws/src/saltybot_visual_odom/setup.py create mode 100644 jetson/ros2_ws/src/saltybot_visual_odom/test/test_kalman_odom.py diff --git a/jetson/ros2_ws/src/saltybot_visual_odom/CMakeLists.txt b/jetson/ros2_ws/src/saltybot_visual_odom/CMakeLists.txt new file mode 100644 index 0000000..bae8ca9 --- /dev/null +++ b/jetson/ros2_ws/src/saltybot_visual_odom/CMakeLists.txt @@ -0,0 +1,11 @@ +cmake_minimum_required(VERSION 3.8) +project(saltybot_visual_odom) + +find_package(ament_cmake REQUIRED) +find_package(ament_cmake_python REQUIRED) + +ament_python_install_package(${PROJECT_NAME}) + +install(DIRECTORY launch config DESTINATION share/${PROJECT_NAME}/) + +ament_package() diff --git a/jetson/ros2_ws/src/saltybot_visual_odom/config/visual_odom_params.yaml b/jetson/ros2_ws/src/saltybot_visual_odom/config/visual_odom_params.yaml new file mode 100644 index 0000000..36a4f13 --- /dev/null +++ b/jetson/ros2_ws/src/saltybot_visual_odom/config/visual_odom_params.yaml @@ -0,0 +1,45 @@ +# saltybot_visual_odom — runtime parameters +# +# Visual odometry requires the D435i IR1 stereo pair and aligned depth: +# /camera/infra1/image_rect_raw — left IR image (mono8, 30 Hz) +# /camera/depth/image_rect_raw — depth in metres (float32, 30 Hz) +# /camera/infra1/camera_info — intrinsics (latched) +# +# Enable infrared emitter in realsense2_camera launch: +# enable_infra1: true +# enable_depth: true +# align_depth.enable: true + +visual_odom: + ros__parameters: + max_features: 300 # max tracked feature points per frame + min_features: 80 # re-detect when below this + frame_id: 'odom' + child_frame_id: 'camera_link' + +odom_fusion: + ros__parameters: + fusion_hz: 30.0 # EKF update + TF publish rate + wheel_odom_topic: '/saltybot/rover_odom' # or /saltybot/tank_odom + loop_closure_thr_m: 0.30 # RTAB-Map jump threshold to trigger correction + loop_closure_alpha: 0.40 # correction weight (0=ignore, 1=snap to RTAB-Map) + slip_noise_multiplier: 10.0 # wheel odom noise × this factor during slip + odom_frame: 'odom' + base_frame: 'base_link' + + +# ── RTAB-Map integration note ───────────────────────────────────────────────── +# To feed the fused odometry to RTAB-Map instead of its internal VO, add to +# slam_rtabmap.launch.py rtabmap_node remappings: +# +# remappings=[ +# ... +# ('odom', '/saltybot/odom_fused'), +# ] +# +# And in rtabmap_params.yaml set: +# subscribe_odom: "true" +# Odom/Strategy: "0" # F2M still used but initialised from fused pose +# +# This feeds RTAB-Map's relocalization with the EKF fused pose, reducing +# loop-closure correction magnitude on startup after wheel slip events. diff --git a/jetson/ros2_ws/src/saltybot_visual_odom/launch/visual_odom.launch.py b/jetson/ros2_ws/src/saltybot_visual_odom/launch/visual_odom.launch.py new file mode 100644 index 0000000..6403e1f --- /dev/null +++ b/jetson/ros2_ws/src/saltybot_visual_odom/launch/visual_odom.launch.py @@ -0,0 +1,74 @@ +""" +visual_odom.launch.py — Visual odometry + Kalman fusion stack. + +Starts: + visual_odom — D435i IR stereo + CUDA optical flow VO → /saltybot/visual_odom + odom_fusion — EKF: visual + wheel + slip failover + RTAB loop closure + → /saltybot/odom_fused + odom→base_link TF + +Launch args: + wheel_odom_topic str '/saltybot/rover_odom' (or /saltybot/tank_odom) + max_features int '300' + loop_closure_thr float '0.30' + publish_debug bool 'false' + +Verify: + ros2 topic hz /saltybot/visual_odom # ~30 Hz + ros2 topic hz /saltybot/odom_fused # ~30 Hz + ros2 topic echo /saltybot/visual_odom_status + ros2 run tf2_tools view_frames # verify odom→base_link TF + +RTAB-Map integration: + RTAB-Map's odom_frame / map_frame should match the odom_frame parameter. + In slam_rtabmap.launch.py / rtabmap_params.yaml set: + odom_frame_id: 'odom' + RTAB-Map will auto-subscribe to /saltybot/odom_fused via the odom remapping + or set subscribe_odom:=true with the odom topic pointing here. +""" + +from launch import LaunchDescription +from launch.actions import DeclareLaunchArgument +from launch.conditions import IfCondition +from launch.substitutions import LaunchConfiguration +from launch_ros.actions import Node + + +def generate_launch_description(): + args = [ + DeclareLaunchArgument('wheel_odom_topic', default_value='/saltybot/rover_odom'), + DeclareLaunchArgument('max_features', default_value='300'), + DeclareLaunchArgument('loop_closure_thr', default_value='0.30'), + DeclareLaunchArgument('loop_closure_alpha', default_value='0.40'), + DeclareLaunchArgument('fusion_hz', default_value='30.0'), + ] + + visual_odom = Node( + package='saltybot_visual_odom', + executable='visual_odom', + name='visual_odom', + output='screen', + parameters=[{ + 'max_features': LaunchConfiguration('max_features'), + 'min_features': 80, + 'frame_id': 'odom', + 'child_frame_id': 'camera_link', + }], + ) + + odom_fusion = Node( + package='saltybot_visual_odom', + executable='odom_fusion', + name='odom_fusion', + output='screen', + parameters=[{ + 'fusion_hz': LaunchConfiguration('fusion_hz'), + 'wheel_odom_topic': LaunchConfiguration('wheel_odom_topic'), + 'loop_closure_thr_m': LaunchConfiguration('loop_closure_thr'), + 'loop_closure_alpha': LaunchConfiguration('loop_closure_alpha'), + 'slip_noise_multiplier': 10.0, + 'odom_frame': 'odom', + 'base_frame': 'base_link', + }], + ) + + return LaunchDescription(args + [visual_odom, odom_fusion]) diff --git a/jetson/ros2_ws/src/saltybot_visual_odom/package.xml b/jetson/ros2_ws/src/saltybot_visual_odom/package.xml new file mode 100644 index 0000000..b31f8c2 --- /dev/null +++ b/jetson/ros2_ws/src/saltybot_visual_odom/package.xml @@ -0,0 +1,41 @@ + + + + saltybot_visual_odom + 0.1.0 + + Visual odometry (D435i stereo, CUDA optical flow) with Kalman-filter fusion + of wheel odometry, auto-failover on wheel slip, and RTAB-Map loop-closure + drift correction. Publishes /saltybot/visual_odom, /saltybot/odom_fused, + and the odom→base_link TF. + + seb + MIT + + ament_cmake + ament_cmake_python + + rclpy + std_msgs + sensor_msgs + nav_msgs + geometry_msgs + tf2_ros + tf2_geometry_msgs + message_filters + cv_bridge + image_transport + + python3-numpy + python3-opencv + python3-scipy + + ament_copyright + ament_flake8 + ament_pep257 + python3-pytest + + + ament_cmake + + diff --git a/jetson/ros2_ws/src/saltybot_visual_odom/resource/saltybot_visual_odom b/jetson/ros2_ws/src/saltybot_visual_odom/resource/saltybot_visual_odom new file mode 100644 index 0000000..e69de29 diff --git a/jetson/ros2_ws/src/saltybot_visual_odom/saltybot_visual_odom/__init__.py b/jetson/ros2_ws/src/saltybot_visual_odom/saltybot_visual_odom/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/jetson/ros2_ws/src/saltybot_visual_odom/saltybot_visual_odom/kalman_odom_filter.py b/jetson/ros2_ws/src/saltybot_visual_odom/saltybot_visual_odom/kalman_odom_filter.py new file mode 100644 index 0000000..2175837 --- /dev/null +++ b/jetson/ros2_ws/src/saltybot_visual_odom/saltybot_visual_odom/kalman_odom_filter.py @@ -0,0 +1,210 @@ +""" +kalman_odom_filter.py — Extended Kalman Filter for odometry fusion. + +Fuses wheel odometry and visual odometry into a single best-estimate +ground-truth trajectory. Handles wheel-slip failover by down-weighting +the wheel odometry measurement covariance when slip is detected. + +State vector: x = [px, py, θ, v, ω] + px, py — position in map/odom frame (m) + θ — heading (rad), wrapped to [-π, π] + v — linear velocity (m/s, forward) + ω — angular velocity (rad/s) + +Process model (unicycle, constant velocity over dt): + px' = px + v·cos(θ)·dt + py' = py + v·sin(θ)·dt + θ' = θ + ω·dt + v' = v (constant velocity prediction) + ω' = ω (constant angular velocity prediction) + +Measurement models: + Wheel odom: z_w = [v_wheel, ω_wheel] (velocity measurements) + Visual odom: z_v = [v_visual, ω_visual] (velocity measurements) + RTAB-Map pose: z_r = [px_rtab, py_rtab, θ_rtab] (absolute pose correction) + +Covariances: + Process noise Q is fixed (low). + Wheel measurement noise R_w increases 10× during slip. + Visual measurement noise R_v increases 3× when VO returns invalid. + +Units: SI throughout (m, rad, m/s, rad/s). +""" + +from __future__ import annotations +import math +from typing import Tuple + +import numpy as np + + +class KalmanOdomFilter: + """ + 5-state EKF for ground-robot odometry fusion. + + Usage: + kf = KalmanOdomFilter() + kf.predict(dt) + kf.update_wheel(v, omega) # wheel odom measurement + kf.update_visual(v, omega) # visual odom measurement + kf.update_rtabmap(px, py, theta) # absolute RTAB-Map correction + px, py, theta = kf.pose + """ + + _N = 5 # state dimension + + def __init__( + self, + # Process noise (standard deviations) + q_pos: float = 0.02, # m per step + q_theta: float = 0.01, # rad per step + q_vel: float = 0.10, # m/s per step + q_omega: float = 0.05, # rad/s per step + # Wheel odom noise (normal) + r_wheel_v: float = 0.05, # m/s + r_wheel_omega: float = 0.05, # rad/s + # Visual odom noise + r_visual_v: float = 0.08, + r_visual_omega:float = 0.08, + # RTAB-Map pose noise + r_rtab_pos: float = 0.10, # m + r_rtab_theta: float = 0.05, # rad + # Slip multiplier for wheel noise + slip_noise_multiplier: float = 10.0, + ): + self._slip_mult = slip_noise_multiplier + + # State estimate + self._x = np.zeros(self._N, dtype=np.float64) # [px, py, theta, v, omega] + + # Covariance + self._P = np.eye(self._N, dtype=np.float64) * 0.01 + + # Process noise covariance Q + self._Q = np.diag([q_pos**2, q_pos**2, q_theta**2, q_vel**2, q_omega**2]) + + # Measurement noise matrices + self._R_wheel_normal = np.diag([r_wheel_v**2, r_wheel_omega**2]) + self._R_wheel_slip = np.diag([ + (r_wheel_v * slip_noise_multiplier)**2, + (r_wheel_omega * slip_noise_multiplier)**2, + ]) + self._R_visual_valid = np.diag([r_visual_v**2, r_visual_omega**2]) + self._R_visual_invalid = np.diag([ + (r_visual_v * 3)**2, + (r_visual_omega * 3)**2, + ]) + self._R_rtabmap = np.diag([r_rtab_pos**2, r_rtab_pos**2, r_rtab_theta**2]) + + # ── Public interface ────────────────────────────────────────────────────── + + @property + def pose(self) -> Tuple[float, float, float]: + """Return (x, y, theta) current best estimate.""" + return float(self._x[0]), float(self._x[1]), float(self._x[2]) + + @property + def velocity(self) -> Tuple[float, float]: + """Return (v_linear, omega_angular) current estimate.""" + return float(self._x[3]), float(self._x[4]) + + @property + def covariance_3x3(self) -> np.ndarray: + """Return 3×3 position+heading covariance (px, py, theta sub-block).""" + return self._P[:3, :3].copy() + + def predict(self, dt: float) -> None: + """Advance the state estimate by dt seconds (unicycle motion model).""" + px, py, th, v, w = self._x + c, s = math.cos(th), math.sin(th) + + # State transition + self._x[0] = px + v * c * dt + self._x[1] = py + v * s * dt + self._x[2] = self._wrap(th + w * dt) + # v and omega unchanged (constant velocity prediction) + + # Jacobian F of state transition + F = np.eye(self._N, dtype=np.float64) + F[0, 2] = -v * s * dt + F[0, 3] = c * dt + F[1, 2] = v * c * dt + F[1, 3] = s * dt + F[2, 4] = dt + + # Covariance prediction + self._P = F @ self._P @ F.T + self._Q + + def update_wheel(self, v: float, omega: float, slipping: bool = False) -> None: + """ + Fuse wheel-odometry velocity measurement. + + Args: + v: linear velocity from encoders (m/s) + omega: angular velocity from encoders (rad/s) + slipping: if True, use high-noise covariance for wheel measurement + """ + R = self._R_wheel_slip if slipping else self._R_wheel_normal + z = np.array([v, omega], dtype=np.float64) + H = np.zeros((2, self._N), dtype=np.float64) + H[0, 3] = 1.0 # measure v + H[1, 4] = 1.0 # measure omega + self._measurement_update(z, H, R) + + def update_visual(self, v: float, omega: float, valid: bool = True) -> None: + """ + Fuse visual-odometry velocity measurement. + + Args: + v, omega: velocity from VO + valid: False if VO returned invalid (low feature count etc.) + """ + R = self._R_visual_valid if valid else self._R_visual_invalid + z = np.array([v, omega], dtype=np.float64) + H = np.zeros((2, self._N), dtype=np.float64) + H[0, 3] = 1.0 + H[1, 4] = 1.0 + self._measurement_update(z, H, R) + + def update_rtabmap(self, px: float, py: float, theta: float) -> None: + """ + Fuse an absolute pose from RTAB-Map (loop closure correction). + + Uses the full pose measurement [px, py, theta] against the current + position state. + """ + z = np.array([px, py, theta], dtype=np.float64) + H = np.zeros((3, self._N), dtype=np.float64) + H[0, 0] = 1.0 # measure px + H[1, 1] = 1.0 # measure py + H[2, 2] = 1.0 # measure theta + + # Innovation with angle wrapping + innov = z - H @ self._x + innov[2] = self._wrap(innov[2]) + + self._measurement_update_with_innov(innov, H, self._R_rtabmap) + + # ── Internal EKF machinery ──────────────────────────────────────────────── + + def _measurement_update( + self, z: np.ndarray, H: np.ndarray, R: np.ndarray + ) -> None: + innov = z - H @ self._x + self._measurement_update_with_innov(innov, H, R) + + def _measurement_update_with_innov( + self, innov: np.ndarray, H: np.ndarray, R: np.ndarray + ) -> None: + S = H @ self._P @ H.T + R + K = self._P @ H.T @ np.linalg.inv(S) + self._x = self._x + K @ innov + self._x[2] = self._wrap(self._x[2]) + self._P = (np.eye(self._N) - K @ H) @ self._P + # Joseph form for numerical stability + I_KH = np.eye(self._N) - K @ H + self._P = I_KH @ self._P @ I_KH.T + K @ R @ K.T + + @staticmethod + def _wrap(angle: float) -> float: + return (angle + math.pi) % (2 * math.pi) - math.pi diff --git a/jetson/ros2_ws/src/saltybot_visual_odom/saltybot_visual_odom/odom_fusion_node.py b/jetson/ros2_ws/src/saltybot_visual_odom/saltybot_visual_odom/odom_fusion_node.py new file mode 100644 index 0000000..dc489e9 --- /dev/null +++ b/jetson/ros2_ws/src/saltybot_visual_odom/saltybot_visual_odom/odom_fusion_node.py @@ -0,0 +1,304 @@ +""" +odom_fusion_node.py — EKF fusion of wheel + visual odometry with slip failover. + +Fusion strategy +─────────────── +Normal operation (no slip): + - Both wheel odom and visual odom update the EKF at their respective rates. + - Final fused pose published at `fusion_hz` (default 30 Hz). + +Wheel slip detected (from /saltybot/terrain): + - Wheel odom measurement covariance inflated 10× → EKF trusts VO more. + - Logs a warning at each slip event transition. + +Visual odom unavailable (no features): + - VO update skipped → EKF relies on wheel odom alone. + - Wheel odom covariance stays at normal level (slip permitting). + +RTAB-Map loop closure correction: + - Subscribes to /rtabmap/odom (graph-optimized, updated after loop closures). + - Detects a loop closure when RTAB-Map's pose jumps > loop_closure_thr_m from + the EKF's current estimate. + - Calls KalmanOdomFilter.update_rtabmap() to soft-correct accumulated drift. + +TF publishing: + - odom → base_link transform published at fusion_hz from the fused estimate. + +Subscribes: + /saltybot/visual_odom nav_msgs/Odometry visual odom velocities + /saltybot/rover_odom nav_msgs/Odometry wheel encoder odom + /saltybot/tank_odom nav_msgs/Odometry tank ESC odom (fallback) + /saltybot/terrain std_msgs/String JSON incl. slip_ratio + is_slipping + /rtabmap/odom nav_msgs/Odometry RTAB-Map graph-optimized pose + +Publishes: + /saltybot/odom_fused nav_msgs/Odometry EKF fused odometry + /saltybot/visual_odom_status std_msgs/String JSON health + diagnostics + +TF: + odom → base_link + +Parameters: + fusion_hz float 30.0 + wheel_odom_topic str '/saltybot/rover_odom' + loop_closure_thr_m float 0.30 (jump > this triggers RTAB correction) + loop_closure_alpha float 0.40 (EKF correction weight, 0–1) + slip_noise_multiplier float 10.0 + odom_frame str 'odom' + base_frame str 'base_link' +""" + +import json +import math +import time + +import rclpy +from rclpy.node import Node +from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy +import tf2_ros +from nav_msgs.msg import Odometry +from geometry_msgs.msg import ( + TransformStamped, Point, Quaternion, + Twist, TwistWithCovariance, PoseWithCovariance, +) +from std_msgs.msg import String, Header + +from .kalman_odom_filter import KalmanOdomFilter + + +_RELIABLE_QOS = QoSProfile( + reliability=ReliabilityPolicy.RELIABLE, + history=HistoryPolicy.KEEP_LAST, + depth=10, +) +_BEST_EFFORT_QOS = QoSProfile( + reliability=ReliabilityPolicy.BEST_EFFORT, + history=HistoryPolicy.KEEP_LAST, + depth=5, +) + + +def _quat_to_yaw(q: Quaternion) -> float: + siny = 2.0 * (q.w * q.z + q.x * q.y) + cosy = 1.0 - 2.0 * (q.y * q.y + q.z * q.z) + return math.atan2(siny, cosy) + + +def _yaw_to_quat(yaw: float) -> Quaternion: + q = Quaternion() + q.w = math.cos(yaw * 0.5) + q.z = math.sin(yaw * 0.5) + return q + + +class OdomFusionNode(Node): + + def __init__(self): + super().__init__('odom_fusion') + + # ── Parameters ────────────────────────────────────────────────────── + self.declare_parameter('fusion_hz', 30.0) + self.declare_parameter('wheel_odom_topic', '/saltybot/rover_odom') + self.declare_parameter('loop_closure_thr_m', 0.30) + self.declare_parameter('loop_closure_alpha', 0.40) + self.declare_parameter('slip_noise_multiplier', 10.0) + self.declare_parameter('odom_frame', 'odom') + self.declare_parameter('base_frame', 'base_link') + + hz = self.get_parameter('fusion_hz').value + wheel_topic = self.get_parameter('wheel_odom_topic').value + self._lc_thr = self.get_parameter('loop_closure_thr_m').value + self._lc_alpha = self.get_parameter('loop_closure_alpha').value + slip_mult = self.get_parameter('slip_noise_multiplier').value + self._odom_frame = self.get_parameter('odom_frame').value + self._base_frame = self.get_parameter('base_frame').value + + # ── EKF ───────────────────────────────────────────────────────────── + self._kf = KalmanOdomFilter(slip_noise_multiplier=slip_mult) + + # ── State ──────────────────────────────────────────────────────────── + self._slipping: bool = False + self._slip_ratio: float = 0.0 + self._vo_valid: bool = False + self._vo_vx: float = 0.0 + self._vo_omega: float = 0.0 + self._last_predict_t: float = time.monotonic() + + # RTAB-Map tracking + self._rtab_prev_x: float | None = None + self._rtab_prev_y: float | None = None + + # ── TF broadcaster ─────────────────────────────────────────────────── + self._tf_broadcaster = tf2_ros.TransformBroadcaster(self) + + # ── Subscriptions ──────────────────────────────────────────────────── + self.create_subscription( + Odometry, '/saltybot/visual_odom', + self._on_visual_odom, _BEST_EFFORT_QOS + ) + self.create_subscription( + Odometry, wheel_topic, + self._on_wheel_odom, _BEST_EFFORT_QOS + ) + # Also listen to tank odom — active if rover odom goes silent + if '/rover' in wheel_topic: + self.create_subscription( + Odometry, wheel_topic.replace('/rover_odom', '/tank_odom'), + self._on_wheel_odom, _BEST_EFFORT_QOS + ) + self.create_subscription( + String, '/saltybot/terrain', self._on_terrain, _BEST_EFFORT_QOS + ) + self.create_subscription( + Odometry, '/rtabmap/odom', self._on_rtabmap_odom, _BEST_EFFORT_QOS + ) + + # ── Publishers ─────────────────────────────────────────────────────── + self._fused_pub = self.create_publisher(Odometry, '/saltybot/odom_fused', 10) + self._status_pub = self.create_publisher(String, '/saltybot/visual_odom_status', 10) + + # ── Fusion timer ───────────────────────────────────────────────────── + self.create_timer(1.0 / hz, self._fusion_tick) + self.create_timer(1.0, self._status_tick) + + self.get_logger().info( + f'odom_fusion ready — wheel={wheel_topic} lc_thr={self._lc_thr}m' + ) + + # ── Subscription callbacks ──────────────────────────────────────────────── + + def _on_visual_odom(self, msg: Odometry) -> None: + self._vo_vx = msg.twist.twist.linear.x + self._vo_omega = msg.twist.twist.angular.z + # Detect VO validity from covariance (inflated when invalid) + self._vo_valid = msg.twist.covariance[0] < 0.3 + # Update EKF immediately on VO callback + self._kf.update_visual(self._vo_vx, self._vo_omega, valid=self._vo_valid) + + def _on_wheel_odom(self, msg: Odometry) -> None: + v = msg.twist.twist.linear.x + omega = msg.twist.twist.angular.z + self._kf.update_wheel(v, omega, slipping=self._slipping) + + def _on_terrain(self, msg: String) -> None: + try: + data = json.loads(msg.data) + except json.JSONDecodeError: + return + was_slipping = self._slipping + self._slipping = bool(data.get('is_slipping', False)) + self._slip_ratio = float(data.get('slip_ratio', 0.0)) + if self._slipping and not was_slipping: + self.get_logger().warning( + f'[odom_fusion] Wheel slip detected (ratio={self._slip_ratio:.2f}) ' + f'— increasing VO weight' + ) + elif not self._slipping and was_slipping: + self.get_logger().info('[odom_fusion] Slip cleared — wheel odom re-weighted') + + def _on_rtabmap_odom(self, msg: Odometry) -> None: + rx = msg.pose.pose.position.x + ry = msg.pose.pose.position.y + ryaw = _quat_to_yaw(msg.pose.pose.orientation) + + # Detect loop closure: RTAB-Map pose jumps significantly from our estimate + if self._rtab_prev_x is not None: + fused_x, fused_y, _ = self._kf.pose + jump = math.hypot(rx - fused_x, ry - fused_y) + if jump > self._lc_thr: + self.get_logger().info( + f'[odom_fusion] Loop closure: RTAB-Map jumped {jump:.3f}m — ' + f'applying drift correction (α={self._lc_alpha})' + ) + self._kf.update_rtabmap(rx, ry, ryaw) + + self._rtab_prev_x = rx + self._rtab_prev_y = ry + + # ── Fusion timer ───────────────────────────────────────────────────────── + + def _fusion_tick(self) -> None: + now_mono = time.monotonic() + dt = now_mono - self._last_predict_t + dt = max(1e-4, min(dt, 0.1)) + self._last_predict_t = now_mono + + # EKF predict step + self._kf.predict(dt) + + # Publish fused odom + now_stamp = self.get_clock().now().to_msg() + px, py, th = self._kf.pose + v, omega = self._kf.velocity + + odom = Odometry() + odom.header.stamp = now_stamp + odom.header.frame_id = self._odom_frame + odom.child_frame_id = self._base_frame + + odom.pose.pose.position = Point(x=px, y=py, z=0.0) + odom.pose.pose.orientation = _yaw_to_quat(th) + + cov3 = self._kf.covariance_3x3 + p_cov = [0.0] * 36 + p_cov[0] = cov3[0, 0] + p_cov[1] = cov3[0, 1] + p_cov[5] = cov3[0, 2] + p_cov[6] = cov3[1, 0] + p_cov[7] = cov3[1, 1] + p_cov[11] = cov3[1, 2] + p_cov[30] = cov3[2, 0] + p_cov[31] = cov3[2, 1] + p_cov[35] = cov3[2, 2] + odom.pose.covariance = p_cov + + odom.twist.twist.linear.x = v + odom.twist.twist.angular.z = omega + t_cov = [0.0] * 36 + t_cov[0] = 0.05 + t_cov[35] = 0.05 + odom.twist.covariance = t_cov + + self._fused_pub.publish(odom) + + # Broadcast odom → base_link TF + tf_msg = TransformStamped() + tf_msg.header.stamp = now_stamp + tf_msg.header.frame_id = self._odom_frame + tf_msg.child_frame_id = self._base_frame + tf_msg.transform.translation.x = px + tf_msg.transform.translation.y = py + tf_msg.transform.translation.z = 0.0 + tf_msg.transform.rotation = _yaw_to_quat(th) + self._tf_broadcaster.sendTransform(tf_msg) + + def _status_tick(self) -> None: + px, py, th = self._kf.pose + v, omega = self._kf.velocity + cov3 = self._kf.covariance_3x3 + status = { + 'pose': {'x': round(px, 3), 'y': round(py, 3), 'yaw_deg': round(math.degrees(th), 1)}, + 'velocity': {'v_mps': round(v, 3), 'omega_rads': round(omega, 3)}, + 'vo_valid': self._vo_valid, + 'slipping': self._slipping, + 'slip_ratio': round(self._slip_ratio, 3), + 'pos_std_m': round(float(cov3[0, 0] ** 0.5), 4), + 'active_source': 'vo+wheel' if self._vo_valid else ('vo_degraded+wheel' if not self._slipping else 'vo_primary'), + } + msg = String() + msg.data = json.dumps(status) + self._status_pub.publish(msg) + + +def main(args=None): + rclpy.init(args=args) + node = OdomFusionNode() + try: + rclpy.spin(node) + finally: + node.destroy_node() + rclpy.shutdown() + + +if __name__ == '__main__': + main() diff --git a/jetson/ros2_ws/src/saltybot_visual_odom/saltybot_visual_odom/optical_flow_tracker.py b/jetson/ros2_ws/src/saltybot_visual_odom/saltybot_visual_odom/optical_flow_tracker.py new file mode 100644 index 0000000..1122da5 --- /dev/null +++ b/jetson/ros2_ws/src/saltybot_visual_odom/saltybot_visual_odom/optical_flow_tracker.py @@ -0,0 +1,207 @@ +""" +optical_flow_tracker.py — Sparse optical flow feature tracker. + +Primary path: OpenCV CUDA (SparsePyrLKOpticalFlow + FAST detector). +Fallback: OpenCV CPU (calcOpticalFlowPyrLK + goodFeaturesToTrack). + +Both paths output the same data: matched pixel pairs (prev_pts, curr_pts) +suitable for Essential-matrix estimation. + +GPU path targets 30 Hz on Jetson Orin Nano Super. The FAST detector on the +Ampere GPU is ~4x faster than the CPU equivalent at the same quality level. + +Design notes: + - Features are re-detected whenever tracking drops below `min_features`. + - Forward-backward consistency check discards unreliable tracks. + - Points in the centre 80% of the image (avoids fisheye edge distortion). + - Max `max_features` points per frame — higher = more robust, more cost. +""" + +from __future__ import annotations +from typing import Tuple + +import numpy as np +import cv2 + + +_USE_CUDA = hasattr(cv2, 'cuda') and cv2.cuda.getCudaEnabledDeviceCount() > 0 + +# LK optical flow parameters +_LK_WIN = (21, 21) +_LK_ITERS = 30 +_LK_EPS = 0.01 +_LK_LEVELS = 4 + +# FAST detector threshold +_FAST_THR = 20 + + +class OpticalFlowTracker: + """ + Maintains a set of tracked feature points across consecutive frames. + + Usage: + tracker = OpticalFlowTracker() + prev_pts, curr_pts = tracker.update(prev_gray, curr_gray) + # prev_pts / curr_pts: (N, 2) float32 arrays of matched pixel coords. + """ + + def __init__(self, max_features: int = 300, min_features: int = 80): + self._max_features = max_features + self._min_features = min_features + self._prev_pts: np.ndarray | None = None + self._prev_gray: np.ndarray | None = None + self._use_cuda = _USE_CUDA + + if self._use_cuda: + self._lk = cv2.cuda.SparsePyrLKOpticalFlow_create( + winSize=_LK_WIN, + maxLevel=_LK_LEVELS, + iters=_LK_ITERS, + useInitialFlow=False, + ) + self._fast = cv2.cuda.FastFeatureDetector_create( + threshold=_FAST_THR, nonmaxSuppression=True + ) + + def update( + self, + prev_gray: np.ndarray, + curr_gray: np.ndarray, + ) -> Tuple[np.ndarray, np.ndarray]: + """ + Track features from prev_gray to curr_gray. + + Returns: + (prev_pts, curr_pts) — matched point pairs, shape (N, 2) float32. + Returns empty arrays if tracking fails or there is no previous frame. + """ + h, w = curr_gray.shape + + # Re-detect if we have too few features or no previous frame + if self._prev_pts is None or len(self._prev_pts) < self._min_features: + self._prev_pts = self._detect(prev_gray, h, w) + self._prev_gray = prev_gray + if self._prev_pts is None or len(self._prev_pts) < 8: + return np.zeros((0, 2), np.float32), np.zeros((0, 2), np.float32) + + # Track features + try: + if self._use_cuda: + prev_pts_c, curr_pts, good = self._track_cuda( + self._prev_gray, curr_gray, self._prev_pts + ) + else: + prev_pts_c, curr_pts, good = self._track_cpu( + self._prev_gray, curr_gray, self._prev_pts + ) + except Exception: + self._prev_pts = None + return np.zeros((0, 2), np.float32), np.zeros((0, 2), np.float32) + + prev_pts_c = prev_pts_c[good] + curr_pts = curr_pts[good] + + # Update state + self._prev_pts = curr_pts + self._prev_gray = curr_gray + + return prev_pts_c.reshape(-1, 2), curr_pts.reshape(-1, 2) + + def reset(self) -> None: + self._prev_pts = None + self._prev_gray = None + + # ── Detection ───────────────────────────────────────────────────────────── + + def _detect(self, gray: np.ndarray, h: int, w: int) -> np.ndarray | None: + """Detect FAST features in the inner 80% of the image.""" + margin_x = int(w * 0.10) + margin_y = int(h * 0.10) + roi = gray[margin_y:h - margin_y, margin_x:w - margin_x] + + if self._use_cuda: + pts = self._detect_cuda(roi) + else: + pts = self._detect_cpu(roi) + + if pts is None or len(pts) == 0: + return None + + # Offset back to full-image coordinates + pts[:, 0] += margin_x + pts[:, 1] += margin_y + + # Subsample if too many + if len(pts) > self._max_features: + idx = np.random.choice(len(pts), self._max_features, replace=False) + pts = pts[idx] + + return pts.astype(np.float32) + + def _detect_cuda(self, roi: np.ndarray) -> np.ndarray | None: + gpu_img = cv2.cuda_GpuMat() + gpu_img.upload(roi) + kp = self._fast.detect(gpu_img, None) + if not kp: + return None + pts = cv2.KeyPoint_convert(kp) + return pts.reshape(-1, 2) + + def _detect_cpu(self, roi: np.ndarray) -> np.ndarray | None: + pts = cv2.goodFeaturesToTrack( + roi, + maxCorners=self._max_features, + qualityLevel=0.01, + minDistance=10, + blockSize=7, + ) + if pts is None: + return None + return pts.reshape(-1, 2) + + # ── Tracking ────────────────────────────────────────────────────────────── + + def _track_cuda( + self, prev: np.ndarray, curr: np.ndarray, pts: np.ndarray + ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + pts_gpu = cv2.cuda_GpuMat(pts.reshape(-1, 1, 2)) + prev_gpu = cv2.cuda_GpuMat(prev) + curr_gpu = cv2.cuda_GpuMat(curr) + curr_pts_gpu, status_gpu, _ = self._lk.calc(prev_gpu, curr_gpu, pts_gpu, None) + curr_pts = curr_pts_gpu.download().reshape(-1, 2) + status = status_gpu.download().ravel() + + # Forward-backward check + bw_pts_gpu, bw_status_gpu, _ = self._lk.calc(curr_gpu, prev_gpu, curr_pts_gpu, None) + bw_pts = bw_pts_gpu.download().reshape(-1, 2) + bw_status = bw_status_gpu.download().ravel() + + fb_err = np.linalg.norm(pts.reshape(-1, 2) - bw_pts, axis=1) + good = (status > 0) & (bw_status > 0) & (fb_err < 1.0) + return pts, curr_pts, good + + def _track_cpu( + self, prev: np.ndarray, curr: np.ndarray, pts: np.ndarray + ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + lk_params = dict( + winSize=_LK_WIN, + maxLevel=_LK_LEVELS, + criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, + _LK_ITERS, _LK_EPS), + ) + curr_pts, st_fwd, _ = cv2.calcOpticalFlowPyrLK( + prev, curr, pts.reshape(-1, 1, 2), None, **lk_params + ) + bw_pts, st_bwd, _ = cv2.calcOpticalFlowPyrLK( + curr, prev, curr_pts, None, **lk_params + ) + fb_err = np.linalg.norm( + pts.reshape(-1, 2) - bw_pts.reshape(-1, 2), axis=1 + ) + good = ( + (st_fwd.ravel() > 0) & + (st_bwd.ravel() > 0) & + (fb_err < 1.0) + ) + return pts.reshape(-1, 2), curr_pts.reshape(-1, 2), good diff --git a/jetson/ros2_ws/src/saltybot_visual_odom/saltybot_visual_odom/stereo_vo.py b/jetson/ros2_ws/src/saltybot_visual_odom/saltybot_visual_odom/stereo_vo.py new file mode 100644 index 0000000..b7f1f24 --- /dev/null +++ b/jetson/ros2_ws/src/saltybot_visual_odom/saltybot_visual_odom/stereo_vo.py @@ -0,0 +1,252 @@ +""" +stereo_vo.py — Stereo visual odometry using Essential matrix decomposition. + +Input per frame: + - Tracked feature pairs (prev_pts, curr_pts) from OpticalFlowTracker + - Depth image (float32, metres) aligned to the IR1/left camera + - Camera intrinsic matrix K + +Pipeline: + 1. Find Essential matrix E from point correspondences (5-point RANSAC). + 2. Decompose E → [R|t] (unit-vector translation). + 3. Recover metric scale from depth at tracked feature points. + 4. Compose the incremental SE(3) pose and integrate into global estimate. + +Scale recovery: + Median depth at both prev_pts and curr_pts yields consistent scale. + Only depth samples in [0.3, 6.0m] range are used (D435i reliable range). + +Output: + StereoVO.update() returns a PoseEstimate dataclass with: + x, y, z — position (m) + roll, pitch, yaw — orientation (rad) + vx, vz, omega — linear + angular velocity (m/s, rad/s) for current step + valid — False if insufficient inliers or depth samples + +Coordinate frame: + x = forward, y = left, z = up (ROS REP-105 base_link convention). + The camera looks forward along its +Z optical axis; we rotate into base_link. +""" + +from __future__ import annotations +import math +from dataclasses import dataclass, field +from typing import Tuple + +import numpy as np +import cv2 + + +@dataclass +class PoseEstimate: + x: float = 0.0 + y: float = 0.0 + z: float = 0.0 + roll: float = 0.0 + pitch: float = 0.0 + yaw: float = 0.0 + vx: float = 0.0 + vy: float = 0.0 + omega: float = 0.0 + valid: bool = False + + +# Depth validity window +_D_MIN = 0.3 +_D_MAX = 6.0 + +# Essential matrix RANSAC +_RANSAC_THRESHOLD = 1.0 # pixels +_RANSAC_PROB = 0.999 +_MIN_INLIERS = 20 + +# Minimum depth samples for reliable scale +_MIN_DEPTH_SAMPLES = 10 + +# Camera→base_link rotation (camera looks forward along +Z, base_link uses +X) +# Rotation: camera_Z → base_X, camera_-Y → base_Z +_R_CAM_TO_BASE = np.array([ + [ 0, 0, 1], + [-1, 0, 0], + [ 0, -1, 0], +], dtype=np.float64) + + +def _median_depth(depth: np.ndarray, pts: np.ndarray) -> float: + """Median depth at pixel locations pts (Nx2 float32). Returns 0 on failure.""" + h, w = depth.shape + u = np.clip(pts[:, 0].astype(int), 0, w - 1) + v = np.clip(pts[:, 1].astype(int), 0, h - 1) + samples = depth[v, u] + valid = samples[(samples > _D_MIN) & (samples < _D_MAX)] + return float(np.median(valid)) if len(valid) >= _MIN_DEPTH_SAMPLES else 0.0 + + +def _rot_to_yaw(R: np.ndarray) -> float: + """Extract yaw from rotation matrix (Z-axis rotation in XY plane).""" + return math.atan2(float(R[1, 0]), float(R[0, 0])) + + +class StereoVO: + """ + Incremental visual odometry with depth-aided scale recovery. + + Call update() once per frame pair. Thread-safe for single caller. + """ + + def __init__(self, K: np.ndarray | None = None, dt: float = 1.0 / 30.0): + """ + Args: + K: 3×3 camera intrinsic matrix. Can be set later via set_K(). + dt: nominal frame interval (seconds) for velocity estimation. + """ + self._K = K + self._dt = dt + + # Accumulated global pose as 4×4 SE(3) + self._pose = np.eye(4, dtype=np.float64) + + # Previous-step pose for velocity estimation + self._prev_pose = np.eye(4, dtype=np.float64) + + self._frame_count = 0 + + def set_K(self, K: np.ndarray) -> None: + self._K = K.reshape(3, 3).astype(np.float64) + + def reset(self) -> None: + self._pose = np.eye(4) + self._prev_pose = np.eye(4) + self._frame_count = 0 + + def update( + self, + prev_pts: np.ndarray, # (N, 2) float32 + curr_pts: np.ndarray, # (N, 2) float32 + depth: np.ndarray, # float32 H×W depth in metres + dt: float | None = None, + ) -> PoseEstimate: + """ + Estimate incremental motion from matched feature pairs. + + Returns a PoseEstimate. On failure returns a copy of the last valid + estimate with valid=False. + """ + self._frame_count += 1 + dt = dt or self._dt + + if self._K is None or len(prev_pts) < 8: + return PoseEstimate( + x=self._pose[0, 3], y=self._pose[1, 3], + yaw=_rot_to_yaw(self._pose[:3, :3]), valid=False + ) + + # ── Essential matrix via 5-point RANSAC ────────────────────────────── + E, mask = cv2.findEssentialMat( + prev_pts, curr_pts, self._K, + method=cv2.RANSAC, + prob=_RANSAC_PROB, + threshold=_RANSAC_THRESHOLD, + ) + if E is None or mask is None: + return self._invalid_est() + + inlier_mask = mask.ravel().astype(bool) + n_inliers = int(inlier_mask.sum()) + if n_inliers < _MIN_INLIERS: + return self._invalid_est() + + p1_in = prev_pts[inlier_mask] + p2_in = curr_pts[inlier_mask] + + # ── Decompose E → R, t ──────────────────────────────────────────────── + _, R, t_unit, _ = cv2.recoverPose(E, p1_in, p2_in, self._K) + # t_unit is a unit vector; we need metric scale + + # ── Scale recovery from depth ───────────────────────────────────────── + scale = _median_depth(depth, p1_in) + if scale <= 0: + # Fallback: try current frame depth + scale = _median_depth(depth, p2_in) + if scale <= 0: + return self._invalid_est() + + # Metric translation in camera frame + t_metric = t_unit.ravel() * scale + + # ── Build incremental SE(3) in camera frame ─────────────────────────── + T_inc_cam = np.eye(4, dtype=np.float64) + T_inc_cam[:3, :3] = R + T_inc_cam[:3, 3] = t_metric + + # ── Rotate into base_link frame ─────────────────────────────────────── + R_cb = _R_CAM_TO_BASE + R_inc_base = R_cb @ R @ R_cb.T + t_inc_base = R_cb @ t_metric + + T_inc = np.eye(4, dtype=np.float64) + T_inc[:3, :3] = R_inc_base + T_inc[:3, 3] = t_inc_base + + # ── Accumulate global pose ──────────────────────────────────────────── + self._prev_pose = self._pose.copy() + self._pose = self._pose @ T_inc + + # ── Extract pose + velocity ─────────────────────────────────────────── + x = float(self._pose[0, 3]) + y = float(self._pose[1, 3]) + z = float(self._pose[2, 3]) + yaw = _rot_to_yaw(self._pose[:3, :3]) + + dx = float(self._pose[0, 3] - self._prev_pose[0, 3]) + dy = float(self._pose[1, 3] - self._prev_pose[1, 3]) + dyaw = yaw - _rot_to_yaw(self._prev_pose[:3, :3]) + # Wrap dyaw to [-π, π] + dyaw = (dyaw + math.pi) % (2 * math.pi) - math.pi + + vx = math.hypot(dx, dy) / max(dt, 1e-6) + # Signed vx: positive if moving forward + if abs(dx) + abs(dy) > 0: + heading_now = yaw + move_dir = math.atan2(dy, dx) + if abs(move_dir - heading_now) > math.pi / 2: + vx = -vx + + omega = dyaw / max(dt, 1e-6) + + return PoseEstimate( + x=x, y=y, z=z, yaw=yaw, + vx=vx, vy=0.0, omega=omega, + valid=True, + ) + + def correct_from_loop_closure( + self, + rtabmap_x: float, + rtabmap_y: float, + rtabmap_yaw: float, + alpha: float = 0.3, + ) -> None: + """ + Soft-correct accumulated pose using RTAB-Map's graph-optimized position. + + Called when RTAB-Map reports a pose significantly different from ours + (indicating a loop closure correction). alpha controls how strongly + we shift toward RTAB-Map's estimate (0=ignore, 1=fully accept). + """ + self._pose[0, 3] += alpha * (rtabmap_x - self._pose[0, 3]) + self._pose[1, 3] += alpha * (rtabmap_y - self._pose[1, 3]) + curr_yaw = _rot_to_yaw(self._pose[:3, :3]) + dyaw = rtabmap_yaw - curr_yaw + dyaw = (dyaw + math.pi) % (2 * math.pi) - math.pi + correction_yaw = curr_yaw + alpha * dyaw + c, s = math.cos(correction_yaw), math.sin(correction_yaw) + self._pose[:3, :3] = np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]]) + + def _invalid_est(self) -> PoseEstimate: + return PoseEstimate( + x=self._pose[0, 3], + y=self._pose[1, 3], + yaw=_rot_to_yaw(self._pose[:3, :3]), + valid=False, + ) diff --git a/jetson/ros2_ws/src/saltybot_visual_odom/saltybot_visual_odom/visual_odom_node.py b/jetson/ros2_ws/src/saltybot_visual_odom/saltybot_visual_odom/visual_odom_node.py new file mode 100644 index 0000000..928dbc7 --- /dev/null +++ b/jetson/ros2_ws/src/saltybot_visual_odom/saltybot_visual_odom/visual_odom_node.py @@ -0,0 +1,200 @@ +""" +visual_odom_node.py — D435i visual odometry at 30 Hz using CUDA optical flow. + +Pipeline per frame: + 1. Receive synchronized IR1 (left infrared) + aligned depth images. + 2. Detect/track FAST features with CUDA SparsePyrLKOpticalFlow (GPU path) + or CPU LK fallback. + 3. Estimate incremental SE(3) via Essential matrix + depth-aided scale. + 4. Publish accumulated pose as nav_msgs/Odometry on /saltybot/visual_odom. + +CUDA usage: + - cv2.cuda.FastFeatureDetector_create — GPU FAST corner detection + - cv2.cuda.SparsePyrLKOpticalFlow — GPU Lucas-Kanade tracking + Both are ~4× faster than CPU equivalents on the Orin Ampere GPU. + The node auto-detects CUDA availability and falls back gracefully. + +Subscribes: + /camera/infra1/image_rect_raw sensor_msgs/Image 30 Hz (grey IR left) + /camera/depth/image_rect_raw sensor_msgs/Image 30 Hz float32 depth (m) + /camera/infra1/camera_info sensor_msgs/CameraInfo intrinsics + +Publishes: + /saltybot/visual_odom nav_msgs/Odometry 30 Hz + +Parameters: + max_features int 300 + min_features int 80 + publish_hz float 30.0 (max rate; actual limited by input) + frame_id str 'odom' + child_frame_id str 'camera_link' +""" + +import time +import math + +import rclpy +from rclpy.node import Node +from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy, DurabilityPolicy +import message_filters +import numpy as np +import cv2 +from cv_bridge import CvBridge + +from sensor_msgs.msg import Image, CameraInfo +from nav_msgs.msg import Odometry +from geometry_msgs.msg import ( + Pose, Point, Quaternion, + Twist, Vector3, TwistWithCovariance, PoseWithCovariance, +) +from std_msgs.msg import Header + +from .optical_flow_tracker import OpticalFlowTracker +from .stereo_vo import StereoVO + + +_SENSOR_QOS = QoSProfile( + reliability=ReliabilityPolicy.BEST_EFFORT, + history=HistoryPolicy.KEEP_LAST, + depth=5, +) +_LATCHED_QOS = QoSProfile( + reliability=ReliabilityPolicy.RELIABLE, + history=HistoryPolicy.KEEP_LAST, + depth=1, + durability=DurabilityPolicy.TRANSIENT_LOCAL, +) + + +def _yaw_to_quat(yaw: float) -> Quaternion: + q = Quaternion() + q.w = math.cos(yaw * 0.5) + q.z = math.sin(yaw * 0.5) + return q + + +class VisualOdomNode(Node): + + def __init__(self): + super().__init__('visual_odom') + + self.declare_parameter('max_features', 300) + self.declare_parameter('min_features', 80) + self.declare_parameter('frame_id', 'odom') + self.declare_parameter('child_frame_id', 'camera_link') + + max_feat = self.get_parameter('max_features').value + min_feat = self.get_parameter('min_features').value + self._frame_id = self.get_parameter('frame_id').value + self._child_frame_id = self.get_parameter('child_frame_id').value + + self._bridge = CvBridge() + self._tracker = OpticalFlowTracker(max_features=max_feat, min_features=min_feat) + self._vo = StereoVO() + self._last_t: float | None = None + + # Synchronized IR1 + depth + ir_sub = message_filters.Subscriber( + self, Image, '/camera/infra1/image_rect_raw', qos_profile=_SENSOR_QOS + ) + depth_sub = message_filters.Subscriber( + self, Image, '/camera/depth/image_rect_raw', qos_profile=_SENSOR_QOS + ) + self._sync = message_filters.ApproximateTimeSynchronizer( + [ir_sub, depth_sub], queue_size=5, slop=0.033 # 1 frame at 30Hz + ) + self._sync.registerCallback(self._on_frame) + + self.create_subscription( + CameraInfo, '/camera/infra1/camera_info', self._on_camera_info, _LATCHED_QOS + ) + + self._odom_pub = self.create_publisher(Odometry, '/saltybot/visual_odom', 10) + + cuda_ok = hasattr(cv2, 'cuda') and cv2.cuda.getCudaEnabledDeviceCount() > 0 + self.get_logger().info( + f'visual_odom ready — CUDA={"YES" if cuda_ok else "NO (CPU fallback)"}' + ) + + def _on_camera_info(self, msg: CameraInfo) -> None: + K = np.array(msg.k, dtype=np.float64).reshape(3, 3) + self._vo.set_K(K) + + def _on_frame(self, ir_msg: Image, depth_msg: Image) -> None: + now = self.get_clock().now() + t_sec = now.nanoseconds * 1e-9 + dt = (t_sec - self._last_t) if self._last_t else (1.0 / 30.0) + dt = max(1e-4, min(dt, 0.5)) # clamp to sane range + self._last_t = t_sec + + # Decode + try: + ir_img = self._bridge.imgmsg_to_cv2(ir_msg, desired_encoding='mono8') + depth_img = self._bridge.imgmsg_to_cv2(depth_msg, desired_encoding='32FC1') + except Exception as exc: + self.get_logger().warning(f'frame decode error: {exc}') + return + + # Build grayscale for optical flow (IR is already mono8) + gray = ir_img + + # Need previous frame — get it from the tracker's internal state + if self._tracker._prev_gray is None: + # First frame: just initialise tracker state + self._tracker._prev_pts = self._tracker._detect( + gray, gray.shape[0], gray.shape[1] + ) + self._tracker._prev_gray = gray + return + + prev_pts, curr_pts = self._tracker.update(self._tracker._prev_gray, gray) + + est = self._vo.update(prev_pts, curr_pts, depth_img, dt=dt) + + odom = self._build_odom_msg(est, now.to_msg()) + self._odom_pub.publish(odom) + + def _build_odom_msg(self, est, stamp) -> Odometry: + from .stereo_vo import PoseEstimate + q = _yaw_to_quat(est.yaw) + + odom = Odometry() + odom.header.stamp = stamp + odom.header.frame_id = self._frame_id + odom.child_frame_id = self._child_frame_id + + odom.pose.pose.position = Point(x=est.x, y=est.y, z=est.z) + odom.pose.pose.orientation = q + + # Covariance — inflate when VO is invalid + cov_xy = 0.05 if est.valid else 0.5 + cov_theta = 0.02 if est.valid else 0.2 + cov_v = 0.1 if est.valid else 1.0 + p_cov = [0.0] * 36 + p_cov[0] = cov_xy # xx + p_cov[7] = cov_xy # yy + p_cov[35] = cov_theta # yaw-yaw + odom.pose.covariance = p_cov + + odom.twist.twist.linear.x = est.vx + odom.twist.twist.angular.z = est.omega + t_cov = [0.0] * 36 + t_cov[0] = cov_v + t_cov[35] = cov_v * 0.5 + odom.twist.covariance = t_cov + + return odom + + +def main(args=None): + rclpy.init(args=args) + node = VisualOdomNode() + try: + rclpy.spin(node) + finally: + node.destroy_node() + rclpy.shutdown() + + +if __name__ == '__main__': + main() diff --git a/jetson/ros2_ws/src/saltybot_visual_odom/setup.py b/jetson/ros2_ws/src/saltybot_visual_odom/setup.py new file mode 100644 index 0000000..b21a340 --- /dev/null +++ b/jetson/ros2_ws/src/saltybot_visual_odom/setup.py @@ -0,0 +1,26 @@ +from setuptools import setup, find_packages + +package_name = 'saltybot_visual_odom' + +setup( + name=package_name, + version='0.1.0', + packages=find_packages(exclude=['test']), + data_files=[ + ('share/ament_index/resource_index/packages', + [f'resource/{package_name}']), + (f'share/{package_name}', ['package.xml']), + ], + install_requires=['setuptools'], + zip_safe=True, + maintainer='seb', + maintainer_email='seb@vayrette.com', + description='Visual odometry + Kalman fusion for saltybot', + license='MIT', + entry_points={ + 'console_scripts': [ + 'visual_odom = saltybot_visual_odom.visual_odom_node:main', + 'odom_fusion = saltybot_visual_odom.odom_fusion_node:main', + ], + }, +) diff --git a/jetson/ros2_ws/src/saltybot_visual_odom/test/test_kalman_odom.py b/jetson/ros2_ws/src/saltybot_visual_odom/test/test_kalman_odom.py new file mode 100644 index 0000000..326dbac --- /dev/null +++ b/jetson/ros2_ws/src/saltybot_visual_odom/test/test_kalman_odom.py @@ -0,0 +1,140 @@ +""" +test_kalman_odom.py — Unit tests for KalmanOdomFilter and StereoVO. + +Runs without a ROS2 environment (no rclpy imports). +""" + +import math +import sys +import os + +import numpy as np +import pytest + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) + +from saltybot_visual_odom.kalman_odom_filter import KalmanOdomFilter +from saltybot_visual_odom.stereo_vo import StereoVO, _rot_to_yaw + + +# ── KalmanOdomFilter tests ──────────────────────────────────────────────────── + +class TestKalmanOdomFilter: + + def test_initial_pose_zero(self): + kf = KalmanOdomFilter() + x, y, th = kf.pose + assert x == pytest.approx(0.0) + assert y == pytest.approx(0.0) + assert th == pytest.approx(0.0) + + def test_predict_straight(self): + kf = KalmanOdomFilter() + # Set velocity to 1 m/s forward, then predict 1s + kf.update_wheel(1.0, 0.0) + kf.predict(1.0) + x, y, th = kf.pose + assert x == pytest.approx(1.0, abs=0.2) + assert abs(y) < 0.1 + assert th == pytest.approx(0.0, abs=0.1) + + def test_predict_pure_rotation(self): + kf = KalmanOdomFilter() + kf.update_wheel(0.0, math.pi / 2) + kf.predict(1.0) + _, _, th = kf.pose + assert th == pytest.approx(math.pi / 2, abs=0.1) + + def test_angle_wrap(self): + kf = KalmanOdomFilter() + kf.update_wheel(0.0, math.pi) + kf.predict(1.1) # should wrap past π + _, _, th = kf.pose + assert -math.pi <= th <= math.pi + + def test_wheel_slip_inflates_covariance(self): + kf_normal = KalmanOdomFilter() + kf_slip = KalmanOdomFilter(slip_noise_multiplier=10.0) + for _ in range(5): + kf_normal.update_wheel(1.0, 0.0, slipping=False) + kf_slip.update_wheel(1.0, 0.0, slipping=True) + kf_normal.predict(0.1) + kf_slip.predict(0.1) + cov_normal = kf_normal.covariance_3x3[0, 0] + cov_slip = kf_slip.covariance_3x3[0, 0] + # Slip covariance should be larger + assert cov_slip > cov_normal + + def test_rtabmap_correction(self): + kf = KalmanOdomFilter() + # Drive straight for 5s to accumulate drift + for _ in range(50): + kf.update_wheel(1.0, 0.0) + kf.predict(0.1) + x_before, y_before, th_before = kf.pose + # RTAB-Map says we're actually 1m behind + kf.update_rtabmap(x_before - 1.0, 0.0, 0.0) + x_after, _, _ = kf.pose + # Pose should have moved toward correction + assert x_after < x_before + + def test_visual_odom_fusion(self): + kf = KalmanOdomFilter() + kf.update_wheel(1.0, 0.0, slipping=False) + kf.update_visual(0.95, 0.0, valid=True) # VO slightly different + kf.predict(0.1) + v, _ = kf.velocity + # Fused velocity should be between 0.95 and 1.0 + assert 0.9 <= v <= 1.05 + + +# ── StereoVO tests ──────────────────────────────────────────────────────────── + +class TestStereoVO: + + def _make_K(self) -> np.ndarray: + return np.array([ + [600.0, 0.0, 320.0], + [ 0.0, 600.0, 240.0], + [ 0.0, 0.0, 1.0], + ]) + + def test_no_K_returns_invalid(self): + vo = StereoVO() + prev = np.random.rand(10, 2).astype(np.float32) * 100 + curr = prev + np.random.rand(10, 2).astype(np.float32) * 2 + depth = np.ones((480, 640), dtype=np.float32) * 1.5 + est = vo.update(prev, curr, depth) + assert not est.valid + + def test_static_scene_near_zero_motion(self): + vo = StereoVO(K=self._make_K()) + # No real motion: prev_pts == curr_pts → zero translation + pts = np.random.rand(80, 2).astype(np.float32) + pts[:, 0] *= 640 + pts[:, 1] *= 480 + depth = np.ones((480, 640), dtype=np.float32) * 2.0 + est = vo.update(pts, pts.copy(), depth) + # With identical points, E is degenerate → may be invalid + # Just ensure it doesn't raise + assert isinstance(est.valid, bool) + + def test_loop_closure_correction(self): + vo = StereoVO(K=self._make_K()) + # Manually set accumulated pose + vo._pose[0, 3] = 5.0 + vo._pose[1, 3] = 0.0 + # RTAB-Map says we're at (4.5, 0) + vo.correct_from_loop_closure(4.5, 0.0, 0.0, alpha=1.0) + assert vo._pose[0, 3] == pytest.approx(4.5, abs=0.01) + + def test_reset_clears_state(self): + vo = StereoVO(K=self._make_K()) + vo._pose[0, 3] = 10.0 + vo.reset() + est = vo._invalid_est() + assert est.x == pytest.approx(0.0) + + +if __name__ == '__main__': + pytest.main([__file__, '-v'])