Files
simulation/src/vision/visual_odometry.py
2026-02-12 14:29:32 -05:00

120 lines
3.8 KiB
Python

#!/usr/bin/env python3
"""
Visual Odometry - Camera-based position estimation without GPS.
ORB/SIFT feature matching with essential matrix decomposition.
"""
import time
import cv2
import numpy as np
from scipy.spatial.transform import Rotation
class VisualOdometry:
def __init__(self, camera_matrix=None, detector_type="ORB", min_features=100):
self.detector_type = detector_type
self.min_features = min_features
if camera_matrix is not None:
self.camera_matrix = np.array(camera_matrix).reshape(3, 3)
else:
self.camera_matrix = np.array(
[[500.0, 0.0, 320.0], [0.0, 500.0, 240.0], [0.0, 0.0, 1.0]]
)
self.dist_coeffs = np.zeros(5)
if detector_type == "ORB":
self.detector = cv2.ORB_create(nfeatures=500)
self.matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=False)
elif detector_type == "SIFT":
self.detector = cv2.SIFT_create(nfeatures=500)
self.matcher = cv2.BFMatcher(cv2.NORM_L2, crossCheck=False)
else:
self.detector = cv2.ORB_create(nfeatures=500)
self.matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=False)
self.prev_frame = None
self.prev_keypoints = None
self.prev_descriptors = None
self.prev_time = None
self.position = np.zeros(3)
self.orientation = np.eye(3)
self.velocity = np.zeros(3)
self.on_pose = None
print(f"[VO] Visual Odometry initialized ({detector_type})")
def process_frame(self, camera_name, frame):
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
keypoints, descriptors = self.detector.detectAndCompute(gray, None)
current_time = time.time()
if (
self.prev_frame is not None
and self.prev_descriptors is not None
and len(keypoints) >= self.min_features
):
matches = self._match_features(self.prev_descriptors, descriptors)
if len(matches) >= self.min_features:
self._estimate_motion(
self.prev_keypoints, keypoints, matches, current_time
)
self.prev_frame = gray
self.prev_keypoints = keypoints
self.prev_descriptors = descriptors
self.prev_time = current_time
def _match_features(self, desc1, desc2):
if desc1 is None or desc2 is None:
return []
matches = self.matcher.knnMatch(desc1, desc2, k=2)
good = []
for pair in matches:
if len(pair) == 2:
m, n = pair
if m.distance < 0.7 * n.distance:
good.append(m)
return good
def _estimate_motion(self, prev_kp, curr_kp, matches, current_time):
if len(matches) < 5:
return
pts1 = np.float32([prev_kp[m.queryIdx].pt for m in matches])
pts2 = np.float32([curr_kp[m.trainIdx].pt for m in matches])
E, mask = cv2.findEssentialMat(
pts1, pts2, self.camera_matrix, method=cv2.RANSAC, prob=0.999, threshold=1.0
)
if E is None:
return
_, R, t, mask = cv2.recoverPose(E, pts1, pts2, self.camera_matrix, mask=mask)
scale = 1.0
translation = scale * t.flatten()
if self.prev_time is not None:
dt = current_time - self.prev_time
if dt > 0:
self.velocity = translation / dt
self.position += self.orientation @ translation
self.orientation = R @ self.orientation
if self.on_pose:
r = Rotation.from_matrix(self.orientation)
self.on_pose(self.position.copy(), r.as_quat())
def get_pose(self):
r = Rotation.from_matrix(self.orientation)
return self.position.copy(), r.as_quat()
def get_velocity(self):
return self.velocity.copy()