예제 #1
0
class Drone:
    def __init__(self,
                 test_flying,
                 mambo_addr,
                 use_wifi=True,
                 use_vision=True):
        """
        Constructor for the Drone Superclass.
        When writing your code, you may redefine __init__() to have additional
        attributes for use in processing elsewhere. If you do this, be sure to
        call super().__init__() with relevant args in order properly
        initialize all mambo-related things.

        test_flying:    bool : run flight code if True
        mambo_addr:     str  : BLE address of drone
        use_wifi:       bool : connect with WiFi instead of BLE if True
        use_vision:     bool : turn on DroneVisionGUI module if True
        """
        self.test_flying = test_flying
        self.use_wifi = use_wifi
        self.use_vision = use_vision
        self.mamboAddr = mambo_addr
        self.mambo = Mambo(self.mamboAddr, use_wifi=self.use_wifi)

        self.mambo.set_user_sensor_callback(self.sensor_cb, args=None)
        if self.use_vision:
            self.mamboVision = DroneVisionGUI(
                self.mambo,
                is_bebop=False,
                buffer_size=200,
                user_code_to_run=self.flight_func,
                user_args=None)
            self.mamboVision.set_user_callback_function(
                self.vision_cb, user_callback_args=None)

    def execute(self):
        """
        Connects to drone and executes flight_func as well as any vision
        handling when needed.
        Run this after initializing your subclass in your main method to
        start the test/script.
        """
        print("trying to connect to self.mambo now")
        success = self.mambo.connect(num_retries=3)
        print("connected: %s" % success)

        if (success):
            self.mambo.smart_sleep(1)
            self.mambo.ask_for_state_update()
            self.mambo.smart_sleep(1)

            if self.use_vision:
                print("starting vision")
                self.mamboVision.open_video()
            else:
                print("starting flight without vision")
                self.flight_func(None, None)

            if self.use_vision:
                print("ending vision")
                self.mamboVision.close_video()
            self.mambo.smart_sleep(3)
            self.mambo.disconnect()

    def flight_func(self, mamboVision, args):
        """
        Method to me run for flight. This is defined by the user outside of this
        class.
        When writing your code, define a new class for each test that inherits
        this class. Redefine your flight_func in your subclass to match
        your flight plan.
        your redefinition of flight_func must include the mamboVision and args
        arguments to properly work.
        NOTE: after takeoff it is smart to do the following:
            print("sensor calib")
            while self.mambo.sensors.speed_ts == 0:
                continue
        This gives the sensors time to fully initialize and send back proper
        readings for use in your sensor callback method.
        Cannot be done before takeoff; sensors send nothing at this time.
        """
        pass

    def vision_cb(self, args):
        """
        Callback function for every vision-handle update Again, defined by the
        user outside of the class in a specific subclass for every test script.
        Your vision_cb must include the self and args arguments to work.
        """
        pass

    def sensor_cb(self, args):
        """
        Callback function for every sensor update. Again, defined by the
        user outside of the class in a specific subclass for every test script.
        Your sensor_cb must include the self and args arguments to work.
        """
        pass
class DroneColorSegTest:
    def __init__(self, testFlying, mamboAddr, use_wifi):
        self.bb = [0, 0, 0, 0]
        self.bb_threshold = 4000
        self.bb_trigger = False

        self.testFlying = testFlying
        self.mamboAddr = mamboAddr
        self.use_wifi = use_wifi
        self.mambo = Mambo(self.mamboAddr, self.use_wifi)
        self.mamboVision = DroneVisionGUI(
            self.mambo,
            is_bebop=False,
            buffer_size=200,
            user_code_to_run=self.mambo_fly_function,
            user_args=None)

    def color_segmentation(self, args):
        img = self.mamboVision.get_latest_valid_picture()

        if img is not None:
            [((x1, y1), (x2, y2)), ln_color] = cd_color_segmentation(img)
            self.bb = [x1, y1, x2, y2]

            bb_size = self.get_bb_size()
            print('bb_size:', bb_size)
            # cv2.imwrite('test_file.png', img) # uncomment to save latest pic
            if bb_size >= self.bb_threshold:
                print('orange detected')
                self.bb_trigger = True
            # else:
            # self.bb_trigger = False
        else:
            print('image is None')

    def get_bb_size(self):
        ''' returns area of self.bb (bounding box) '''
        return (self.bb[2] - self.bb[0]) * (self.bb[3] - self.bb[1])

    def mambo_fly_function(self, mamboVision, args):
        """
        self.mambo takes off and yaws slowly in a circle until the vision processing
        detects a large patch of orange. It then performs a flip and lands.
        """

        if (self.testFlying):
            print("taking off!")
            self.mambo.safe_takeoff(5)

            if (self.mambo.sensors.flying_state != "emergency"):
                print("flying state is %s" % self.mambo.sensors.flying_state)
                print("Flying direct: going up")
                self.mambo.fly_direct(roll=0,
                                      pitch=0,
                                      yaw=0,
                                      vertical_movement=15,
                                      duration=2)

                print("flying state is %s" % self.mambo.sensors.flying_state)
                print("yawing slowly")
                for deg in range(36):
                    self.mambo.turn_degrees(10)
                    if self.bb_trigger:
                        break
                    self.mambo.smart_sleep(1)

                print("flying state is %s" % self.mambo.sensors.flying_state)
                print("flip left")
                success = self.mambo.flip(direction="left")
                print("self.mambo flip result %s" % success)
                self.mambo.smart_sleep(2)

            print("landing")
            print("flying state is %s" % self.mambo.sensors.flying_state)
            self.mambo.safe_land(5)
        else:
            print("Sleeeping for 15 seconds - move the self.mambo around")
            self.mambo.smart_sleep(15)

        # done doing vision demo
        print("Ending the sleep and vision")
        self.mamboVision.close_video()

        self.mambo.smart_sleep(5)

        print("disconnecting")
        self.mambo.disconnect()

    def run_test(self):
        print("trying to connect to self.mambo now")
        success = self.mambo.connect(num_retries=3)
        print("connected: %s" % success)

        if (success):
            # get the state information
            print("sleeping")
            self.mambo.smart_sleep(1)
            self.mambo.ask_for_state_update()
            self.mambo.smart_sleep(1)

            print("Preparing to open vision")
            self.mamboVision.set_user_callback_function(
                self.color_segmentation, user_callback_args=None)
            self.mamboVision.open_video()