go_to_cube.py

#!/usr/bin/env python3
#!c:/Python35/python3.exe -u
import asyncio
import sys
import cv2
import numpy as np
import cozmo
import time
import os
from glob import glob

from find_cube import *

try:
    from PIL import ImageDraw, ImageFont
except ImportError:
    sys.exit('run `pip3 install --user Pillow numpy` to run this example')


def nothing(x):
    pass


PINK_LOWER = np.array(
    np.array([351 / 2, .69 * 255, .76 * 255]).round(), np.uint8)
PINK_UPPER = np.array(
    np.array([358 / 2, .81 * 255, .80 * 255]).round(), np.uint8)

IMAGE_WIDTH = 320

GREEN_LOWER = np.array([0, 0, 0])
GREEN_UPPER = np.array([179, 255, 60])


# Define a decorator as a subclass of Annotator; displays the keypoint
class BoxAnnotator(cozmo.annotate.Annotator):

    cube = None

    def apply(self, image, scale):
        d = ImageDraw.Draw(image)
        bounds = (0, 0, image.width, image.height)

        if BoxAnnotator.cube is not None:

            #double size of bounding box to match size of rendered image
            BoxAnnotator.cube = np.multiply(BoxAnnotator.cube, 2)

            #define and display bounding box with params:
            #msg.img_topLeft_x, msg.img_topLeft_y, msg.img_width, msg.img_height
            box = cozmo.util.ImageBox(
                BoxAnnotator.cube[0] - BoxAnnotator.cube[2] / 2,
                BoxAnnotator.cube[1] - BoxAnnotator.cube[2] / 2,
                BoxAnnotator.cube[2], BoxAnnotator.cube[2])
            cozmo.annotate.add_img_box_to_image(image, box, "green", text=None)

            BoxAnnotator.cube = None


async def run(robot: cozmo.robot.Robot):

    robot.world.image_annotator.annotation_enabled = True
    robot.world.image_annotator.add_annotator('box', BoxAnnotator)

    robot.camera.image_stream_enabled = True
    robot.camera.color_image_enabled = True
    robot.camera.enable_auto_exposure = True

    robot.set_head_angle(cozmo.util.degrees(0))

    gain, exposure, mode = 390, 3, 1

    try:

        while True:
            event = await robot.world.wait_for(
                cozmo.camera.EvtNewRawCameraImage,
                timeout=30)  #get camera image
            if event.image is not None:
                image = cv2.cvtColor(
                    np.asarray(event.image), cv2.COLOR_BGR2RGB)

                if mode == 1:
                    robot.camera.enable_auto_exposure = True
                else:
                    robot.camera.set_manual_exposure(exposure, fixed_gain)

                #find the cube
                cube = find_cube(image, PINK_LOWER, PINK_UPPER)
                print(cube)
                BoxAnnotator.cube = cube

                #motion
                # if no keypoint, start turning til there is one
                if (cube is None):
                    robot.turn_in_place(
                        cozmo.util.degrees(37),
                        in_parallel=True).wait_for_completed()
                else:  # turn until it is in the center
                    last_turn = 0
                    oscillations = 0
                    delta = (IMAGE_WIDTH / 2) - cube[0]
                    if last_turn == np.sign(delta) * -1:
                        oscillations += 1
                    else:
                        oscillations = 0
                    if abs(delta) > 30:
                        # assume <15 oscillations
                        robot.turn_in_place(
                            cozmo.util.degrees(
                                np.sign(delta) * (15 - oscillations)),
                            in_parallel=True).wait_for_completed()
                    else:
                        size = np.pi * (cube[2]**2)
                        if size < (320 - 40) * (240 - 30):
                            robot.drive_straight(
                                cozmo.util.distance_inches(6),
                                cozmo.util.speed_mmps(400),
                                in_parallel=True).wait_for_completed()
                            
    except KeyboardInterrupt:
        print("")
        print("Exit requested by user")
    except cozmo.RobotBusy as e:
        print(e)
    #cv2.destroyAllWindows()


if __name__ == '__main__':
    cozmo.run_program(run, use_viewer=True, force_viewer_on_top=True)