189 lines
6.2 KiB
Python
189 lines
6.2 KiB
Python
#!/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()
|