backyard_flyer.py

import argparse
import time
from enum import Enum
from typing import List

import numpy as np

from udacidrone import Drone
from udacidrone.connection import MavlinkConnection, WebSocketConnection  # noqa: F401
from udacidrone.messaging import MsgID


class States(Enum):
    MANUAL = 0
    ARMING = 1
    TAKEOFF = 2
    WAYPOINT = 3
    LANDING = 4
    DISARMING = 5


class BackyardFlyer(Drone):

    def __init__(self, connection):
        super().__init__(connection)
        self.in_mission = True
        self.check_state = {}
        self.target_position = [0, 0, 0]

        # Prepare the plan.
        self.target_altitude = 3
        self.all_waypoints = self.calculate_box(self.target_altitude)

        # Initial state.
        self.flight_state = States.MANUAL

        # Register all your callbacks here.
        self.register_callback(MsgID.LOCAL_POSITION, self.local_position_callback)
        self.register_callback(MsgID.LOCAL_VELOCITY, self.velocity_callback)
        self.register_callback(MsgID.STATE, self.state_callback)

    def local_position_callback(self):
        """
        This triggers when `MsgID.LOCAL_POSITION` is received and self.local_position contains new data
        """
        if self.flight_state == States.TAKEOFF:
            # Coordinate conversion.
            altitude = -1.0 * self.local_position[2]

            # Check if altitude is within 95% of target.
            if altitude > 0.95 * self.target_position[2]:
                self.waypoint_transition()

        elif self.flight_state == States.WAYPOINT:
            # Determine distance to waypoint.
            target_xy = self.target_position[0:2]
            local_xy = self.local_position[0:2]
            distance = np.linalg.norm(target_xy - local_xy)

            # When we're close to the waypoint, move on to the next one.
            # This is a bit sloppy, but then again - so is our control.
            if distance < 1.0:
                if len(self.all_waypoints) > 0:
                    self.waypoint_transition()
                else:
                    # Wait for the movement to settle, then initiate landing.
                    velocity_xy = self.local_velocity[0:2]
                    if np.linalg.norm(velocity_xy) < 0.5:
                        self.landing_transition()

    def velocity_callback(self):
        """
        This triggers when `MsgID.LOCAL_VELOCITY` is received and self.local_velocity contains new data
        """
        if self.flight_state == States.LANDING:
            if ((self.global_position[2] - self.global_home[2] < 0.1) and abs(self.local_position[2]) < 0.01):
                self.disarming_transition()

    def state_callback(self):
        """
        This triggers when `MsgID.STATE` is received and self.armed and self.guided contain new data
        """
        if not self.in_mission:
            return
        if self.flight_state == States.MANUAL:
            self.arming_transition()
        elif self.flight_state == States.ARMING:
            if self.armed:
                self.takeoff_transition()
        elif self.flight_state == States.DISARMING:
            # Transitioning to manual control while being armed
            # is dangerous, since we can't know the state of the manual controls.
            if not self.armed:
                self.manual_transition()

    def calculate_box(self, target_altitude: int) -> List[List[float]]:
        """
        1. Return waypoints to fly a box
        """
        local_waypoints = [[10.0,  0.0, target_altitude],
                           [10.0, 10.0, target_altitude],
                           [ 0.0, 10.0, target_altitude],
                           [ 0.0,  0.0, target_altitude]]
        return local_waypoints

    def arming_transition(self):
        """
        1. Take control of the drone
        2. Pass an arming command
        3. Set the home location to current position
        4. Transition to the ARMING state
        """
        print("arming transition")
        self.take_control()
        self.arm()

        # Set the current location to be the home position.
        self.set_home_position(self.global_position[0],
                               self.global_position[1],
                               self.global_position[2])

        self.flight_state = States.ARMING

    def takeoff_transition(self):
        """
        1. Set target_position altitude to 3.0m
        2. Command a takeoff to 3.0m
        3. Transition to the TAKEOFF state
        """
        print("takeoff transition")

        self.target_position[2] = self.target_altitude
        self.takeoff(self.target_altitude)
        self.flight_state = States.TAKEOFF

    def waypoint_transition(self):
        """
        1. Command the next waypoint position
        2. Transition to WAYPOINT state
        """
        print("waypoint transition")

        # We're treating the waypoints as a queue.
        # This way, new waypoints could always be added to the end,
        # and we'd always try to fly to the next waypoint without
        # having to keep track of a separate index.
        # A trajectory planner could issue an entirely new set of
        # waypoints every once in a while (according to current conditions)
        # and this logic would still apply (sans concurrency issues, but
        # then again, in Python we have the GIL).
        self.target_position = self.all_waypoints.pop(0)

        north, east = self.target_position[0], self.target_position[1]
        altitude, heading = self.target_position[2], 0

        self.cmd_position(north, east, altitude, heading)
        self.flight_state = States.WAYPOINT

    def landing_transition(self):
        """
        1. Command the drone to land
        2. Transition to the LANDING state
        """
        print("landing transition")
        self.land()
        self.flight_state = States.LANDING

    def disarming_transition(self):
        """
        1. Command the drone to disarm
        2. Transition to the DISARMING state
        """
        print("disarm transition")
        self.disarm()
        self.flight_state = States.DISARMING

    def manual_transition(self):
        """This method is provided

        1. Release control of the drone
        2. Stop the connection (and telemetry log)
        3. End the mission
        4. Transition to the MANUAL state
        """
        print("manual transition")

        self.release_control()
        self.stop()
        self.in_mission = False
        self.flight_state = States.MANUAL

    def start(self):
        """This method is provided

        1. Open a log file
        2. Start the drone connection
        3. Close the log file
        """
        print("Creating log file")
        self.start_log("Logs", "NavLog.txt")
        print("starting connection")
        self.connection.start()
        print("Closing log file")
        self.stop_log()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--port', type=int, default=5760, help='Port number')
    parser.add_argument('--host', type=str, default='127.0.0.1', help="host address, i.e. '127.0.0.1'")
    args = parser.parse_args()

    conn = MavlinkConnection('tcp:{0}:{1}'.format(args.host, args.port), threaded=False, PX4=False)
    #conn = WebSocketConnection('ws://{0}:{1}'.format(args.host, args.port))
    drone = BackyardFlyer(conn)
    time.sleep(2)
    drone.start()