Crazyflie_backyard_flyer.py

"""
Created on Tue Oct 24 16:17:28 2017

@author: Adrien
"""

import argparse
import time
from enum import Enum

import numpy as np

from udacidrone import Drone
from udacidrone.connection import CrazyflieConnection  # 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.target_position = np.array([0.0, 0.0, 0.0])
        self.all_waypoints = []
        self.in_mission = True
        self.check_state = {}

        # 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):
        ### added code ###
        if self.flight_state == States.MANUAL:
            self.takeoff(0.5)
            self.target_position[2] = 0.5
            self.flight_state = States.TAKEOFF
        ### added code ###

        if self.flight_state == States.TAKEOFF:
            if -1.0 * self.local_position[2] > 0.95 * self.target_position[2]:
                self.all_waypoints = self.calculate_box()
                self.waypoint_transition()

        elif self.flight_state == States.WAYPOINT:
            # DEBUG
            print("curr pos: ({}, {}, {}), desired pos: ({}, {}, {})".format(
                self.local_position[0], self.local_position[1], self.local_position[2],
                self.target_position[0], self.target_position[1], self.target_position[2]))

            if np.linalg.norm(self.target_position[0:2] - self.local_position[0:2]) < 0.2:
                print("reached waypoint!!!!")
                if len(self.all_waypoints) > 0:
                    self.waypoint_transition()
                else:
                    if np.linalg.norm(self.local_velocity[0:2]) < 0.5:
                        self.landing_transition()

    def velocity_callback(self):
        ### added code ###
        if self.flight_state == States.LANDING:
            if abs(self.local_position[2] < 0.01):
                print("vehicle is down!")
                self.manual_transition()
        ### added code ###

        # if self.flight_state == States.LANDING:
        #     # TODO: need a better stop criteria, this is incredibly specific to
        #     if self.global_position[2] - self.global_home[2] < 0.1:
        #         if abs(self.local_position[2]) < 0.01:
        #             self.disarming_transition()

    def state_callback(self):
        if self.in_mission:
            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:
                if ~self.armed & ~self.guided:
                    self.manual_transition()

    def calculate_box(self):
        print("Setting Home")
        # local_waypoints = [[10.0, 0.0, -3.0], [10.0, 10.0, -3.0], [0.0, 10.0, -3.0], [0.0, 0.0, -3.0]]
        cp = self.local_position
        cp[2] = 0
        local_waypoints = [cp + [1.0, 0.0, 0.5], cp + [1.0, 1.0, 0.5], cp + [.0, 1.0, 0.5], cp + [0.0, 0.0, 0.5]]

        return local_waypoints

    def arming_transition(self):
        print("arming transition")
        self.take_control()
        self.arm()
        self.set_home_position(self.global_position[0], self.global_position[1],
                               self.global_position[2])  # set the current location to be the home position

        self.flight_state = States.ARMING

    def takeoff_transition(self):
        print("takeoff transition")
        # self.global_home = np.copy(self.global_position)  # can't write to this variable!
        target_altitude = 3.0
        self.target_position[2] = target_altitude
        self.takeoff(target_altitude)
        self.flight_state = States.TAKEOFF

    def waypoint_transition(self):
        print("waypoint transition")
        self.target_position = self.all_waypoints.pop(0)
        print('target position', self.target_position)
        self.cmd_position(self.target_position[0], self.target_position[1], self.target_position[2], 0.0)
        self.flight_state = States.WAYPOINT

    def landing_transition(self):
        print("landing transition")
        self.land()
        self.flight_state = States.LANDING

    def disarming_transition(self):
        print("disarm transition")
        self.disarm()
        self.release_control()
        self.flight_state = States.DISARMING

    def manual_transition(self):
        print("manual transition")
        self.stop()
        self.in_mission = False
        self.flight_state = States.MANUAL

    def start(self):
        self.start_log("Logs", "NavLog.txt")
        # self.connect()

        print("starting connection")
        # self.connection.start()

        super().start()

        # Only required if they do threaded
        # while self.in_mission:
        #    pass

        self.stop_log()


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

    # conn = MavlinkConnection('tcp:{0}:{1}'.format(args.host, args.port))
    # conn = WebSocketConnection('ws://{0}:{1}'.format(args.host, args.port))
    uri = 'radio://0/80/2M'
    conn = CrazyflieConnection(uri)
    drone = BackyardFlyer(conn)
    time.sleep(2)
    drone.start()