SHARE
TWEET

Untitled

a guest Oct 18th, 2019 72 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. ## If you run into an "[NSApplication _setup] unrecognized selector" problem on macOS,
  2. ## try uncommenting the following snippet
  3.  
  4. # try:
  5. #     import matplotlib
  6. #     matplotlib.use('TkAgg')
  7. # except ImportError:
  8. #     pass
  9.  
  10. from skimage import color
  11. import cozmo
  12. import cozmo.song
  13. import numpy as np
  14. from numpy.linalg import inv
  15. import threading
  16. import time
  17. import sys
  18. import asyncio
  19. from PIL import Image
  20.  
  21. from markers import detect, annotator
  22.  
  23. from grid import CozGrid
  24. from gui import GUIWindow
  25. from particle import Particle, Robot
  26. from setting import *
  27. from particle_filter import *
  28. from utils import *
  29.  
  30. #particle filter functionality
  31. class ParticleFilter:
  32.  
  33.     def __init__(self, grid):
  34.         self.particles = Particle.create_random(PARTICLE_COUNT, grid)
  35.         self.grid = grid
  36.  
  37.     def update(self, odom, r_marker_list):
  38.  
  39.         # ---------- Motion model update ----------
  40.         self.particles = motion_update(self.particles, odom)
  41.  
  42.         # ---------- Sensor (markers) model update ----------
  43.         self.particles = measurement_update(self.particles, r_marker_list, self.grid)
  44.  
  45.         # ---------- Show current state ----------
  46.         # Try to find current best estimate for display
  47.         m_x, m_y, m_h, m_confident = compute_mean_pose(self.particles)
  48.         return (m_x, m_y, m_h, m_confident)
  49.  
  50. # tmp cache
  51. last_pose = cozmo.util.Pose(0,0,0,angle_z=cozmo.util.Angle(degrees=0))
  52. flag_odom_init = False
  53.  
  54. # goal location for the robot to drive to, (x, y, theta)
  55. goal = (6,10,0)
  56.  
  57. # map
  58. Map_filename = "map_arena.json"
  59. grid = CozGrid(Map_filename)
  60. gui = GUIWindow(grid, show_camera=True)
  61. pf = ParticleFilter(grid)
  62.  
  63. def compute_odometry(curr_pose, cvt_inch=True):
  64.     '''
  65.     Compute the odometry given the current pose of the robot (use robot.pose)
  66.  
  67.     Input:
  68.         - curr_pose: a cozmo.robot.Pose representing the robot's current location
  69.         - cvt_inch: converts the odometry into grid units
  70.     Returns:
  71.         - 3-tuple (dx, dy, dh) representing the odometry
  72.     '''
  73.  
  74.     global last_pose, flag_odom_init
  75.     last_x, last_y, last_h = last_pose.position.x, last_pose.position.y, \
  76.         last_pose.rotation.angle_z.degrees
  77.     curr_x, curr_y, curr_h = curr_pose.position.x, curr_pose.position.y, \
  78.         curr_pose.rotation.angle_z.degrees
  79.    
  80.     dx, dy = rotate_point(curr_x-last_x, curr_y-last_y, -last_h)
  81.     if cvt_inch:
  82.         dx, dy = dx / grid.scale, dy / grid.scale
  83.  
  84.     return (dx, dy, diff_heading_deg(curr_h, last_h))
  85.  
  86.  
  87. async def marker_processing(robot, camera_settings, show_diagnostic_image=False):
  88.     '''
  89.     Obtain the visible markers from the current frame from Cozmo's camera.
  90.     Since this is an async function, it must be called using await, for example:
  91.  
  92.         markers, camera_image = await marker_processing(robot, camera_settings, show_diagnostic_image=False)
  93.  
  94.     Input:
  95.         - robot: cozmo.robot.Robot object
  96.         - camera_settings: 3x3 matrix representing the camera calibration settings
  97.         - show_diagnostic_image: if True, shows what the marker detector sees after processing
  98.     Returns:
  99.         - a list of detected markers, each being a 3-tuple (rx, ry, rh)
  100.           (as expected by the particle filter's measurement update)
  101.         - a PIL Image of what Cozmo's camera sees with marker annotations
  102.     '''
  103.  
  104.     global grid
  105.  
  106.     # Wait for the latest image from Cozmo
  107.     image_event = await robot.world.wait_for(cozmo.camera.EvtNewRawCameraImage, timeout=30)
  108.  
  109.     # Convert the image to grayscale
  110.     image = np.array(image_event.image)
  111.     image = color.rgb2gray(image)
  112.    
  113.     # Detect the markers
  114.     markers, diag = detect.detect_markers(image, camera_settings, include_diagnostics=True)
  115.  
  116.     # Measured marker list for the particle filter, scaled by the grid scale
  117.     marker_list = [marker['xyh'] for marker in markers]
  118.     marker_list = [(x/grid.scale, y/grid.scale, h) for x,y,h in marker_list]
  119.  
  120.     # Annotate the camera image with the markers
  121.     if not show_diagnostic_image:
  122.         annotated_image = image_event.image.resize((image.shape[1] * 2, image.shape[0] * 2))
  123.         annotator.annotate_markers(annotated_image, markers, scale=2)
  124.     else:
  125.         diag_image = color.gray2rgb(diag['filtered_image'])
  126.         diag_image = Image.fromarray(np.uint8(diag_image * 255)).resize((image.shape[1] * 2, image.shape[0] * 2))
  127.         annotator.annotate_markers(diag_image, markers, scale=2)
  128.         annotated_image = diag_image
  129.  
  130.     return marker_list, annotated_image
  131.  
  132.  
  133. async def run(robot: cozmo.robot.Robot):
  134.  
  135.     global flag_odom_init, last_pose
  136.     global grid, gui, pf
  137.  
  138.     # start streaming
  139.     robot.camera.image_stream_enabled = True
  140.     robot.camera.color_image_enabled = False
  141.     robot.camera.enable_auto_exposure()
  142.     await robot.set_head_angle(cozmo.util.degrees(0)).wait_for_completed()
  143.  
  144.     # Obtain the camera intrinsics
  145.     fx, fy = robot.camera.config.focal_length.x_y
  146.     cx, cy = robot.camera.config.center.x_y
  147.     camera_settings = np.array([
  148.         [fx,  0, cx],
  149.         [ 0, fy, cy],
  150.         [ 0,  0,  1]
  151.     ], dtype=np.float)
  152.  
  153.     ###################
  154.     # gui.start()
  155.     # YOUR CODE HERE
  156.     await robot.set_head_angle(cozmo.util.degrees(5)).wait_for_completed()
  157.     while(True):
  158.         if robot.is_picked_up:
  159.             print("kidnapped")
  160.             await robot.play_song([cozmo.song.SongNote(noteType=cozmo.song.NoteTypes.D2_Sharp), cozmo.song.SongNote()]).wait_for_completed()
  161.             pf = ParticleFilter(grid)
  162.             robot.stop_all_motors()
  163.             await robot.set_head_angle(cozmo.util.degrees(0)).wait_for_completed()
  164.             continue
  165.  
  166.         old_pose = robot.pose
  167.         odom = compute_odometry(old_pose)
  168.         markers, camera_image = await marker_processing(robot, camera_settings)
  169.         m_x, m_y, m_h, m_confident = pf.update(odom, markers)
  170.  
  171.  
  172.         d_x = goal[0] - m_x
  173.         d_y = goal[1] - m_y
  174.         d_h = goal[2] - m_h
  175.  
  176.         gui.show_mean(m_x, m_y, m_h, m_confident)
  177.         gui.show_particles(pf.particles)
  178.         # gui.show_robot(robot)
  179.         gui.show_camera_image(camera_image)
  180.         # gui.update()
  181.         gui.updated.set()
  182.         if m_confident:
  183.             print((d_x, d_y, d_h))
  184.             robot.stop_all_motors()
  185.             await robot.turn_in_place(cozmo.util.degrees(d_h)).wait_for_completed()
  186.             await robot.turn_in_place(cozmo.util.degrees(math.degrees(math.atan2(d_y, d_x)))).wait_for_completed()
  187.             dist = math.sqrt(d_x*d_x+d_y*d_y)
  188.             moved = 0
  189.             while moved < dist:
  190.                 if robot.is_picked_up:
  191.                     print("kidnapped")
  192.                     await robot.play_song([cozmo.song.SongNote(noteType=cozmo.song.NoteTypes.D2_Sharp),
  193.                                            cozmo.song.SongNote()]).wait_for_completed()
  194.                     pf = ParticleFilter(grid)
  195.                     robot.stop_all_motors()
  196.                     await robot.set_head_angle(cozmo.util.degrees(0)).wait_for_completed()
  197.                     break
  198.                 current_move = min(dist - moved, 1.0)
  199.                 await robot.drive_straight(cozmo.util.distance_inches(current_move), cozmo.util.speed_mmps(40)).wait_for_completed()
  200.                 moved += current_move
  201.  
  202.             if robot.is_picked_up:
  203.                 print("kidnapped")
  204.                 await robot.play_song([cozmo.song.SongNote(noteType=cozmo.song.NoteTypes.D2_Sharp),
  205.                                        cozmo.song.SongNote()]).wait_for_completed()
  206.                 pf = ParticleFilter(grid)
  207.                 robot.stop_all_motors()
  208.                 await robot.set_head_angle(cozmo.util.degrees(0)).wait_for_completed()
  209.             else:
  210.                 await robot.turn_in_place(cozmo.util.radians(-math.atan2(d_y, d_x))).wait_for_completed()
  211.  
  212.                 robot.play_audio(cozmo.audio.AudioEvents.MusicHotPotatoLevel1Loop)
  213.                 break
  214.  
  215.         else:
  216.             # Try w/o ifelses
  217.             await robot.turn_in_place(cozmo.util.degrees(10)).wait_for_completed()
  218.  
  219.         last_pose = old_pose
  220.  
  221.  
  222.  
  223.     ###################
  224.  
  225. class CozmoThread(threading.Thread):
  226.    
  227.     def __init__(self):
  228.         threading.Thread.__init__(self, daemon=False)
  229.  
  230.     def run(self):
  231.         cozmo.robot.Robot.drive_off_charger_on_connect = False  # Cozmo can stay on his charger
  232.         cozmo.run_program(run, use_viewer=False)
  233.  
  234.  
  235. if __name__ == '__main__':
  236.  
  237.     # cozmo thread
  238.     cozmo_thread = CozmoThread()
  239.     cozmo_thread.start()
  240.  
  241.     # init
  242.     gui.show_particles(pf.particles)
  243.     gui.show_mean(0, 0, 0)
  244.     gui.start()
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
Not a member of Pastebin yet?
Sign Up, it unlocks many cool features!
 
Top