Initial Commit

src/control/uav_controller.py

@@ -0,0 +1,249 @@
+#!/usr/bin/env python3
+"""UAV Controller - Flight control using local position only."""
+
+import rclpy
+from rclpy.node import Node
+from geometry_msgs.msg import PoseStamped, TwistStamped, Twist
+from mavros_msgs.srv import CommandBool, SetMode, CommandTOL
+from mavros_msgs.msg import State
+from std_msgs.msg import String, Bool
+import numpy as np
+from enum import Enum
+
+
+class FlightState(Enum):
+    DISARMED = 0
+    ARMED = 1
+    TAKING_OFF = 2
+    HOVERING = 3
+    NAVIGATING = 4
+    LANDING = 5
+    EMERGENCY = 6
+
+
+class UAVController(Node):
+    
+    def __init__(self):
+        super().__init__('uav_controller')
+        
+        self.declare_parameter('takeoff_altitude', 5.0)
+        self.declare_parameter('position_tolerance', 0.3)
+        self.declare_parameter('namespace', '/uav')
+        
+        self.takeoff_altitude = self.get_parameter('takeoff_altitude').value
+        self.position_tolerance = self.get_parameter('position_tolerance').value
+        ns = self.get_parameter('namespace').value
+        
+        self.state = FlightState.DISARMED
+        self.mavros_state = None
+        self.current_pose = None
+        self.target_pose = None
+        self.home_position = None
+        
+        self.state_sub = self.create_subscription(
+            State, f'{ns}/mavros/state', self.state_callback, 10)
+        
+        self.local_pose_sub = self.create_subscription(
+            PoseStamped, f'{ns}/mavros/local_position/pose', self.pose_callback, 10)
+        
+        self.cmd_sub = self.create_subscription(
+            String, f'{ns}/controller/command', self.command_callback, 10)
+        
+        self.setpoint_sub = self.create_subscription(
+            PoseStamped, f'{ns}/setpoint_position', self.setpoint_callback, 10)
+        
+        self.vel_sub = self.create_subscription(
+            Twist, f'{ns}/cmd_vel_safe', self.velocity_callback, 10)
+        
+        self.setpoint_pub = self.create_publisher(
+            PoseStamped, f'{ns}/mavros/setpoint_position/local', 10)
+        
+        self.vel_pub = self.create_publisher(
+            TwistStamped, f'{ns}/mavros/setpoint_velocity/cmd_vel', 10)
+        
+        self.status_pub = self.create_publisher(
+            String, f'{ns}/controller/status', 10)
+        
+        self.arming_client = self.create_client(CommandBool, f'{ns}/mavros/cmd/arming')
+        self.set_mode_client = self.create_client(SetMode, f'{ns}/mavros/set_mode')
+        self.takeoff_client = self.create_client(CommandTOL, f'{ns}/mavros/cmd/takeoff')
+        self.land_client = self.create_client(CommandTOL, f'{ns}/mavros/cmd/land')
+        
+        self.setpoint_timer = self.create_timer(0.05, self.publish_setpoint)
+        self.status_timer = self.create_timer(0.5, self.publish_status)
+        
+        self.get_logger().info('UAV Controller Started - GPS-denied mode')
+    
+    def state_callback(self, msg):
+        self.mavros_state = msg
+        if msg.armed and self.state == FlightState.DISARMED:
+            self.state = FlightState.ARMED
+    
+    def pose_callback(self, msg):
+        self.current_pose = msg
+        if self.home_position is None and self.mavros_state and self.mavros_state.armed:
+            self.home_position = np.array([
+                msg.pose.position.x,
+                msg.pose.position.y,
+                msg.pose.position.z
+            ])
+            self.get_logger().info(f'Home position set: {self.home_position}')
+    
+    def setpoint_callback(self, msg):
+        self.target_pose = msg
+    
+    def velocity_callback(self, msg):
+        if self.state not in [FlightState.NAVIGATING, FlightState.HOVERING]:
+            return
+        
+        vel_msg = TwistStamped()
+        vel_msg.header.stamp = self.get_clock().now().to_msg()
+        vel_msg.header.frame_id = 'base_link'
+        vel_msg.twist = msg
+        self.vel_pub.publish(vel_msg)
+    
+    def command_callback(self, msg):
+        cmd = msg.data.lower()
+        
+        if cmd == 'arm':
+            self.arm()
+        elif cmd == 'disarm':
+            self.disarm()
+        elif cmd == 'takeoff':
+            self.takeoff()
+        elif cmd == 'land':
+            self.land()
+        elif cmd == 'rtl':
+            self.return_to_launch()
+        elif cmd == 'hold':
+            self.hold_position()
+        elif cmd == 'guided':
+            self.set_mode('GUIDED')
+    
+    def arm(self):
+        if not self.arming_client.wait_for_service(timeout_sec=5.0):
+            self.get_logger().error('Arming service not available')
+            return False
+        
+        req = CommandBool.Request()
+        req.value = True
+        future = self.arming_client.call_async(req)
+        future.add_done_callback(self.arm_response)
+        return True
+    
+    def arm_response(self, future):
+        try:
+            result = future.result()
+            if result.success:
+                self.state = FlightState.ARMED
+                self.get_logger().info('Vehicle armed')
+            else:
+                self.get_logger().error('Arming failed')
+        except Exception as e:
+            self.get_logger().error(f'Arming error: {e}')
+    
+    def disarm(self):
+        req = CommandBool.Request()
+        req.value = False
+        future = self.arming_client.call_async(req)
+        self.state = FlightState.DISARMED
+    
+    def set_mode(self, mode):
+        if not self.set_mode_client.wait_for_service(timeout_sec=5.0):
+            return False
+        
+        req = SetMode.Request()
+        req.custom_mode = mode
+        future = self.set_mode_client.call_async(req)
+        return True
+    
+    def takeoff(self):
+        self.set_mode('GUIDED')
+        self.arm()
+        
+        if self.current_pose is not None:
+            self.target_pose = PoseStamped()
+            self.target_pose.header.frame_id = 'odom'
+            self.target_pose.pose.position.x = self.current_pose.pose.position.x
+            self.target_pose.pose.position.y = self.current_pose.pose.position.y
+            self.target_pose.pose.position.z = self.takeoff_altitude
+            self.target_pose.pose.orientation.w = 1.0
+            
+            self.state = FlightState.TAKING_OFF
+            self.get_logger().info(f'Taking off to {self.takeoff_altitude}m')
+    
+    def land(self):
+        if self.current_pose is not None:
+            self.target_pose = PoseStamped()
+            self.target_pose.header.frame_id = 'odom'
+            self.target_pose.pose.position.x = self.current_pose.pose.position.x
+            self.target_pose.pose.position.y = self.current_pose.pose.position.y
+            self.target_pose.pose.position.z = 0.0
+            self.target_pose.pose.orientation.w = 1.0
+            
+            self.state = FlightState.LANDING
+            self.get_logger().info('Landing')
+    
+    def return_to_launch(self):
+        if self.home_position is not None:
+            self.target_pose = PoseStamped()
+            self.target_pose.header.frame_id = 'odom'
+            self.target_pose.pose.position.x = float(self.home_position[0])
+            self.target_pose.pose.position.y = float(self.home_position[1])
+            self.target_pose.pose.position.z = self.takeoff_altitude
+            self.target_pose.pose.orientation.w = 1.0
+            
+            self.state = FlightState.NAVIGATING
+            self.get_logger().info('Returning to local home position')
+    
+    def hold_position(self):
+        if self.current_pose is not None:
+            self.target_pose = PoseStamped()
+            self.target_pose.header = self.current_pose.header
+            self.target_pose.pose = self.current_pose.pose
+            self.state = FlightState.HOVERING
+    
+    def publish_setpoint(self):
+        if self.target_pose is None:
+            return
+        
+        self.target_pose.header.stamp = self.get_clock().now().to_msg()
+        self.setpoint_pub.publish(self.target_pose)
+        
+        if self.current_pose is not None:
+            error = np.array([
+                self.target_pose.pose.position.x - self.current_pose.pose.position.x,
+                self.target_pose.pose.position.y - self.current_pose.pose.position.y,
+                self.target_pose.pose.position.z - self.current_pose.pose.position.z
+            ])
+            distance = np.linalg.norm(error)
+            
+            if self.state == FlightState.TAKING_OFF and distance < self.position_tolerance:
+                self.state = FlightState.HOVERING
+                self.get_logger().info('Takeoff complete')
+            elif self.state == FlightState.LANDING and self.current_pose.pose.position.z < 0.2:
+                self.state = FlightState.ARMED
+                self.get_logger().info('Landing complete')
+    
+    def publish_status(self):
+        status_msg = String()
+        armed = 'ARMED' if (self.mavros_state and self.mavros_state.armed) else 'DISARMED'
+        mode = self.mavros_state.mode if self.mavros_state else 'UNKNOWN'
+        status_msg.data = f'{self.state.name}|{armed}|{mode}'
+        self.status_pub.publish(status_msg)
+
+
+def main(args=None):
+    rclpy.init(args=args)
+    node = UAVController()
+    try:
+        rclpy.spin(node)
+    except KeyboardInterrupt:
+        pass
+    finally:
+        node.destroy_node()
+        rclpy.shutdown()
+
+
+if __name__ == '__main__':
+    main()