Initial Commit

src/localization/position_estimator.py

@@ -0,0 +1,188 @@
+#!/usr/bin/env python3
+"""Position Estimator - Fuses sensors for local position estimate."""
+
+import rclpy
+from rclpy.node import Node
+from geometry_msgs.msg import PoseStamped, TwistStamped, PoseWithCovarianceStamped
+from sensor_msgs.msg import Imu
+from nav_msgs.msg import Odometry
+import numpy as np
+from scipy.spatial.transform import Rotation
+
+
+class PositionEstimator(Node):
+    
+    def __init__(self):
+        super().__init__('position_estimator')
+        
+        self.declare_parameter('fusion_method', 'weighted_average')
+        self.declare_parameter('vo_weight', 0.6)
+        self.declare_parameter('optical_flow_weight', 0.3)
+        self.declare_parameter('imu_weight', 0.1)
+        self.declare_parameter('update_rate', 50.0)
+        
+        self.fusion_method = self.get_parameter('fusion_method').value
+        self.vo_weight = self.get_parameter('vo_weight').value
+        self.of_weight = self.get_parameter('optical_flow_weight').value
+        self.imu_weight = self.get_parameter('imu_weight').value
+        
+        self.position = np.zeros(3)
+        self.velocity = np.zeros(3)
+        self.orientation = np.array([0.0, 0.0, 0.0, 1.0])
+        
+        self.vo_pose = None
+        self.of_velocity = None
+        self.imu_data = None
+        
+        self.prev_time = None
+        
+        self.vo_sub = self.create_subscription(
+            PoseStamped, '/uav/visual_odometry/pose', self.vo_callback, 10)
+        
+        self.of_sub = self.create_subscription(
+            TwistStamped, '/uav/optical_flow/velocity', self.of_callback, 10)
+        
+        self.imu_sub = self.create_subscription(
+            Imu, '/uav/imu/data', self.imu_callback, 10)
+        
+        self.pose_pub = self.create_publisher(
+            PoseWithCovarianceStamped, '/uav/local_position/pose', 10)
+        
+        self.odom_pub = self.create_publisher(
+            Odometry, '/uav/local_position/odom', 10)
+        
+        update_rate = self.get_parameter('update_rate').value
+        self.timer = self.create_timer(1.0 / update_rate, self.update_estimate)
+        
+        self.get_logger().info(f'Position Estimator Started - {self.fusion_method}')
+    
+    def vo_callback(self, msg):
+        self.vo_pose = msg
+    
+    def of_callback(self, msg):
+        self.of_velocity = msg
+    
+    def imu_callback(self, msg):
+        self.imu_data = msg
+    
+    def update_estimate(self):
+        current_time = self.get_clock().now()
+        
+        if self.prev_time is not None:
+            dt = (current_time.nanoseconds - self.prev_time.nanoseconds) / 1e9
+            if dt <= 0:
+                dt = 0.02
+        else:
+            dt = 0.02
+        
+        self.prev_time = current_time
+        
+        if self.fusion_method == 'weighted_average':
+            self.weighted_average_fusion(dt)
+        else:
+            self.simple_fusion(dt)
+        
+        self.publish_estimate()
+    
+    def weighted_average_fusion(self, dt):
+        total_weight = 0.0
+        weighted_pos = np.zeros(3)
+        weighted_vel = np.zeros(3)
+        
+        if self.vo_pose is not None:
+            vo_pos = np.array([
+                self.vo_pose.pose.position.x,
+                self.vo_pose.pose.position.y,
+                self.vo_pose.pose.position.z
+            ])
+            weighted_pos += self.vo_weight * vo_pos
+            total_weight += self.vo_weight
+            
+            self.orientation = np.array([
+                self.vo_pose.pose.orientation.x,
+                self.vo_pose.pose.orientation.y,
+                self.vo_pose.pose.orientation.z,
+                self.vo_pose.pose.orientation.w
+            ])
+        
+        if self.of_velocity is not None:
+            of_vel = np.array([
+                self.of_velocity.twist.linear.x,
+                self.of_velocity.twist.linear.y,
+                0.0
+            ])
+            weighted_vel += self.of_weight * of_vel
+        
+        if self.imu_data is not None:
+            imu_accel = np.array([
+                self.imu_data.linear_acceleration.x,
+                self.imu_data.linear_acceleration.y,
+                self.imu_data.linear_acceleration.z - 9.81
+            ])
+            self.velocity += imu_accel * dt * self.imu_weight
+        
+        if total_weight > 0:
+            self.position = weighted_pos / total_weight
+        
+        if self.of_velocity is not None:
+            self.velocity = weighted_vel
+    
+    def simple_fusion(self, dt):
+        if self.vo_pose is not None:
+            self.position = np.array([
+                self.vo_pose.pose.position.x,
+                self.vo_pose.pose.position.y,
+                self.vo_pose.pose.position.z
+            ])
+    
+    def publish_estimate(self):
+        pose_msg = PoseWithCovarianceStamped()
+        pose_msg.header.stamp = self.get_clock().now().to_msg()
+        pose_msg.header.frame_id = 'odom'
+        
+        pose_msg.pose.pose.position.x = float(self.position[0])
+        pose_msg.pose.pose.position.y = float(self.position[1])
+        pose_msg.pose.pose.position.z = float(self.position[2])
+        
+        pose_msg.pose.pose.orientation.x = float(self.orientation[0])
+        pose_msg.pose.pose.orientation.y = float(self.orientation[1])
+        pose_msg.pose.pose.orientation.z = float(self.orientation[2])
+        pose_msg.pose.pose.orientation.w = float(self.orientation[3])
+        
+        pose_msg.pose.covariance = [
+            0.1, 0, 0, 0, 0, 0,
+            0, 0.1, 0, 0, 0, 0,
+            0, 0, 0.1, 0, 0, 0,
+            0, 0, 0, 0.05, 0, 0,
+            0, 0, 0, 0, 0.05, 0,
+            0, 0, 0, 0, 0, 0.05
+        ]
+        
+        self.pose_pub.publish(pose_msg)
+        
+        odom_msg = Odometry()
+        odom_msg.header = pose_msg.header
+        odom_msg.child_frame_id = 'base_link'
+        odom_msg.pose = pose_msg.pose
+        
+        odom_msg.twist.twist.linear.x = float(self.velocity[0])
+        odom_msg.twist.twist.linear.y = float(self.velocity[1])
+        odom_msg.twist.twist.linear.z = float(self.velocity[2])
+        
+        self.odom_pub.publish(odom_msg)
+
+
+def main(args=None):
+    rclpy.init(args=args)
+    node = PositionEstimator()
+    try:
+        rclpy.spin(node)
+    except KeyboardInterrupt:
+        pass
+    finally:
+        node.destroy_node()
+        rclpy.shutdown()
+
+
+if __name__ == '__main__':
+    main()