Camera Sim Update
This commit is contained in:
@@ -1,128 +1,85 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Optical Flow - Velocity estimation from downward camera."""
|
||||
"""
|
||||
Optical Flow - Velocity estimation from downward camera.
|
||||
Lucas-Kanade sparse optical flow for GPS-denied velocity estimation.
|
||||
"""
|
||||
|
||||
import rclpy
|
||||
from rclpy.node import Node
|
||||
from sensor_msgs.msg import Image, Range
|
||||
from geometry_msgs.msg import TwistStamped
|
||||
from cv_bridge import CvBridge
|
||||
import time
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
class OpticalFlowNode(Node):
|
||||
|
||||
def __init__(self):
|
||||
super().__init__('optical_flow_node')
|
||||
|
||||
self.bridge = CvBridge()
|
||||
|
||||
class OpticalFlowEstimator:
|
||||
def __init__(self, focal_length=500.0, min_altitude=0.3, max_altitude=10.0):
|
||||
self.focal_length = focal_length
|
||||
self.min_altitude = min_altitude
|
||||
self.max_altitude = max_altitude
|
||||
|
||||
self.prev_frame = None
|
||||
self.prev_points = None
|
||||
self.current_altitude = 1.0
|
||||
|
||||
self.declare_parameter('window_size', 15)
|
||||
self.declare_parameter('max_level', 3)
|
||||
self.declare_parameter('min_altitude', 0.3)
|
||||
self.declare_parameter('max_altitude', 10.0)
|
||||
self.declare_parameter('focal_length', 500.0)
|
||||
|
||||
self.window_size = self.get_parameter('window_size').value
|
||||
self.max_level = self.get_parameter('max_level').value
|
||||
self.min_altitude = self.get_parameter('min_altitude').value
|
||||
self.max_altitude = self.get_parameter('max_altitude').value
|
||||
self.focal_length = self.get_parameter('focal_length').value
|
||||
|
||||
self.lk_params = dict(
|
||||
winSize=(self.window_size, self.window_size),
|
||||
maxLevel=self.max_level,
|
||||
criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03)
|
||||
)
|
||||
|
||||
self.feature_params = dict(
|
||||
maxCorners=100,
|
||||
qualityLevel=0.3,
|
||||
minDistance=7,
|
||||
blockSize=7
|
||||
)
|
||||
|
||||
self.velocity = np.zeros(2)
|
||||
self.prev_time = None
|
||||
|
||||
self.image_sub = self.create_subscription(
|
||||
Image, '/uav/camera/downward/image_raw', self.image_callback, 10)
|
||||
|
||||
self.altitude_sub = self.create_subscription(
|
||||
Range, '/uav/rangefinder/range', self.altitude_callback, 10)
|
||||
|
||||
self.velocity_pub = self.create_publisher(
|
||||
TwistStamped, '/uav/optical_flow/velocity', 10)
|
||||
|
||||
self.get_logger().info('Optical Flow Node Started')
|
||||
|
||||
def altitude_callback(self, msg):
|
||||
if msg.range > self.min_altitude and msg.range < self.max_altitude:
|
||||
self.current_altitude = msg.range
|
||||
|
||||
def image_callback(self, msg):
|
||||
try:
|
||||
frame = self.bridge.imgmsg_to_cv2(msg, 'bgr8')
|
||||
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
||||
|
||||
current_time = self.get_clock().now()
|
||||
|
||||
if self.prev_frame is not None and self.prev_points is not None:
|
||||
new_points, status, error = cv2.calcOpticalFlowPyrLK(
|
||||
self.prev_frame, gray, self.prev_points, None, **self.lk_params)
|
||||
|
||||
if new_points is not None:
|
||||
good_new = new_points[status == 1]
|
||||
good_old = self.prev_points[status == 1]
|
||||
|
||||
if len(good_new) > 10:
|
||||
flow = good_new - good_old
|
||||
avg_flow = np.mean(flow, axis=0)
|
||||
|
||||
if self.prev_time is not None:
|
||||
dt = (current_time - self.prev_time).nanoseconds / 1e9
|
||||
if dt > 0:
|
||||
pixel_velocity = avg_flow / dt
|
||||
|
||||
self.velocity[0] = pixel_velocity[0] * self.current_altitude / self.focal_length
|
||||
self.velocity[1] = pixel_velocity[1] * self.current_altitude / self.focal_length
|
||||
|
||||
self.publish_velocity()
|
||||
|
||||
self.prev_points = cv2.goodFeaturesToTrack(gray, **self.feature_params)
|
||||
self.prev_frame = gray
|
||||
self.prev_time = current_time
|
||||
|
||||
except Exception as e:
|
||||
self.get_logger().error(f'Optical flow error: {e}')
|
||||
|
||||
def publish_velocity(self):
|
||||
vel_msg = TwistStamped()
|
||||
vel_msg.header.stamp = self.get_clock().now().to_msg()
|
||||
vel_msg.header.frame_id = 'base_link'
|
||||
|
||||
vel_msg.twist.linear.x = float(self.velocity[0])
|
||||
vel_msg.twist.linear.y = float(self.velocity[1])
|
||||
vel_msg.twist.linear.z = 0.0
|
||||
|
||||
self.velocity_pub.publish(vel_msg)
|
||||
self.current_altitude = 1.0
|
||||
|
||||
self.velocity = np.zeros(2)
|
||||
|
||||
def main(args=None):
|
||||
rclpy.init(args=args)
|
||||
node = OpticalFlowNode()
|
||||
try:
|
||||
rclpy.spin(node)
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
finally:
|
||||
node.destroy_node()
|
||||
rclpy.shutdown()
|
||||
self.lk_params = dict(
|
||||
winSize=(15, 15),
|
||||
maxLevel=3,
|
||||
criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03),
|
||||
)
|
||||
|
||||
self.feature_params = dict(
|
||||
maxCorners=100, qualityLevel=0.3, minDistance=7, blockSize=7
|
||||
)
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
self.on_velocity = None
|
||||
|
||||
print("[OF] Optical Flow Estimator initialized")
|
||||
|
||||
def set_altitude(self, altitude):
|
||||
if self.min_altitude < altitude < self.max_altitude:
|
||||
self.current_altitude = altitude
|
||||
|
||||
def process_frame(self, camera_name, frame):
|
||||
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
||||
current_time = time.time()
|
||||
|
||||
if self.prev_frame is not None and self.prev_points is not None:
|
||||
new_points, status, error = cv2.calcOpticalFlowPyrLK(
|
||||
self.prev_frame, gray, self.prev_points, None, **self.lk_params
|
||||
)
|
||||
|
||||
if new_points is not None:
|
||||
good_new = new_points[status.flatten() == 1]
|
||||
good_old = self.prev_points[status.flatten() == 1]
|
||||
|
||||
if len(good_new) > 10:
|
||||
flow = good_new - good_old
|
||||
avg_flow = np.mean(flow, axis=0)
|
||||
|
||||
if self.prev_time is not None:
|
||||
dt = current_time - self.prev_time
|
||||
if dt > 0:
|
||||
pixel_velocity = avg_flow / dt
|
||||
self.velocity[0] = (
|
||||
pixel_velocity[0]
|
||||
* self.current_altitude
|
||||
/ self.focal_length
|
||||
)
|
||||
self.velocity[1] = (
|
||||
pixel_velocity[1]
|
||||
* self.current_altitude
|
||||
/ self.focal_length
|
||||
)
|
||||
|
||||
if self.on_velocity:
|
||||
self.on_velocity(self.velocity)
|
||||
|
||||
self.prev_points = cv2.goodFeaturesToTrack(gray, **self.feature_params)
|
||||
self.prev_frame = gray
|
||||
self.prev_time = current_time
|
||||
|
||||
return self.velocity
|
||||
|
||||
def get_velocity(self):
|
||||
return self.velocity.copy()
|
||||
|
||||
Reference in New Issue
Block a user