You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

CozmoController.py 6.5KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187
  1. import os
  2. import asyncio
  3. import cozmo
  4. import PIL.Image
  5. import PIL.ImageFont
  6. import PIL.ImageTk
  7. import cv2
  8. import numpy
  9. import tkinter
  10. import tensorflow as tf
  11. import time
  12. from CozmoDriver import CozmoDriver
  13. from tensorflow.python.platform import gfile
  14. USE_LOGGING = False
  15. class CozmoController:
  16. #Only for testing
  17. def printTensors(self, pb_file):
  18. # read pb into graph_def
  19. with tf.io.gfile.GFile(pb_file, "rb") as f:
  20. graph_def = tf.compat.v1.GraphDef()
  21. graph_def.ParseFromString(f.read())
  22. # import graph_def
  23. with tf.Graph().as_default() as graph:
  24. tf.import_graph_def(graph_def)
  25. # print operations
  26. for op in graph.get_operations():
  27. print(op.name)
  28. def __init__(self, sess, logits):
  29. self._robot = None
  30. self._tk_root = 0
  31. self._tk_label_input = 0
  32. self._tk_label_output = 0
  33. self.session = sess
  34. self.logits = logits
  35. self.cozmoDriver = 0
  36. self._counter = 0
  37. self._frameDecisonModifier = 1
  38. self._startTime = time.time()
  39. # Testing
  40. #self.printTensors(GRAPH_PB_PATH)
  41. if USE_LOGGING:
  42. cozmo.setup_basic_logging()
  43. cozmo.connect(self.run)
  44. def on_new_camera_image(self, event, *, image:cozmo.world.CameraImage, **kw):
  45. '''
  46. Gets called on every new image that arrives from cozmos camera
  47. :param event:
  48. :param image:
  49. :param kw:
  50. :return:
  51. '''
  52. elapsedTime = time.time() - self._startTime
  53. self._startTime = time.time()
  54. #print("Elapsed Time: {}".format(elapsedTime))
  55. self._counter = self._counter + 1
  56. raw_image = image.raw_image
  57. # Convert PIL Image to OpenCV Image
  58. cv2_image = cv2.cvtColor(numpy.array(raw_image), cv2.COLOR_RGB2BGR)
  59. # Apply edge filter
  60. cv2_edges = self.auto_canny(cv2_image)
  61. #print("Shape of the image {}".format(cv2_edges.shape))
  62. # Crop image
  63. cv2_edges = cv2_edges[60:180, :]
  64. # resize image
  65. cv2_edges = cv2.resize(cv2_edges, (80, 30))
  66. # Convert output image back to PIL image
  67. pil_edges = PIL.Image.fromarray(cv2.cvtColor(cv2_edges, cv2.COLOR_GRAY2RGB))
  68. # Display input and output feed
  69. display_image_input = PIL.ImageTk.PhotoImage(image=image.annotate_image())
  70. display_image_output = PIL.ImageTk.PhotoImage(image=pil_edges)
  71. self._tk_label_input.imgtk = display_image_input
  72. self._tk_label_input.configure(image=display_image_input)
  73. self._tk_label_output.imgtk = display_image_output
  74. self._tk_label_output.configure(image=display_image_output)
  75. self._tk_root.update()
  76. cv2_edges = numpy.array(cv2_edges).reshape(-1, 30, 80, 1)
  77. #get prediction from model
  78. prediciton = self.session.run("action:0", {logits: cv2_edges, "action_masks:0": [[1,1,1,1]]},)
  79. #Control on which frame a decision should produce a new action
  80. if(self._counter % self._frameDecisonModifier == 0):
  81. if(numpy.argmax(prediciton) == 0):
  82. #behaves to strict no movement anymores
  83. #self.cozmoDriver.stop()
  84. pass
  85. elif(numpy.argmax(prediciton) == 1):
  86. self.cozmoDriver.drive()
  87. elif(numpy.argmax(prediciton) == 2):
  88. self.cozmoDriver.driveRight()
  89. elif(numpy.argmax(prediciton) == 3):
  90. self.cozmoDriver.driveLeft()
  91. # Auto-paramter Canny edge detection adapted from:
  92. # http://www.pyimagesearch.com/2015/04/06/zero-parameter-automatic-canny-edge-detection-with-python-and-opencv/
  93. def auto_canny(self, img, sigma=0.33):
  94. '''
  95. Canny algorithem with automatic parameters in opencv
  96. :param img:
  97. :param sigma:
  98. :return:
  99. '''
  100. blurred = cv2.GaussianBlur(img, (3, 3), 0)
  101. v = numpy.median(blurred)
  102. lower = int(max(0, (1.0 - sigma) * v))
  103. upper = int(min(255, (1.0 + sigma) * v))
  104. edged = cv2.Canny(blurred, lower-33, upper)
  105. return edged
  106. async def set_up_cozmo(self, coz_conn):
  107. '''
  108. initialize cozmo to get images from camera and let him drive
  109. :param coz_conn:
  110. :return:
  111. '''
  112. asyncio.set_event_loop(coz_conn._loop)
  113. self._robot = await coz_conn.wait_for_robot()
  114. #enable camera stream
  115. self._robot.camera.image_stream_enabled = True
  116. #event handler for new images
  117. self._robot.add_event_handler(cozmo.world.EvtNewCameraImage, self.on_new_camera_image)
  118. #set cozmos head to lowest position
  119. await self._robot.set_head_angle(cozmo.robot.MIN_HEAD_ANGLE).wait_for_completed()
  120. #set cozmos lift to highest position
  121. await self._robot.set_lift_height(1).wait_for_completed()
  122. self.cozmoDriver = CozmoDriver(self._robot)
  123. async def run(self, coz_conn):
  124. # Set up Cozmo
  125. await self.set_up_cozmo(coz_conn)
  126. self.cozmo = CozmoDriver(self._robot)
  127. self._tk_root = tkinter.Tk()
  128. # TODO: ESC to exit
  129. self._tk_label_input = tkinter.Label(self._tk_root)
  130. self._tk_label_output = tkinter.Label(self._tk_root)
  131. self._tk_label_input.pack()
  132. self._tk_label_output.pack()
  133. while True:
  134. await asyncio.sleep(0)
  135. if __name__ == '__main__':
  136. # original trained brain
  137. GRAPH_PB_PATH = '..\\Models\\1_original_trained_brain\\CozmoLearningBrain\\frozen_graph_def.pb'
  138. # brain with stop movement allowed and slower turnspeed in training
  139. #GRAPH_PB_PATH = '..\\Models\\2_modified_turnspeed_actionmask_brain\\CozmoLearningBrain\\frozen_graph_def.pb'
  140. # brain with a new actionmask trained that remembers the 3 latest actions it took
  141. #GRAPH_PB_PATH = '..\\Models\\3_modified_actionmask_memory_3_brain\\CozmoLearningBrain\\frozen_graph_def.pb'
  142. # brain with a new actionmask trained that remembers the 5 latest actions it took
  143. #GRAPH_PB_PATH = '..\\Models\\4_modified_actionmask_memory_5_brain\\CozmoLearningBrain\\frozen_graph_def.pb'
  144. with tf.compat.v1.Session() as sess:
  145. print("load model graph")
  146. with tf.io.gfile.GFile(GRAPH_PB_PATH, 'rb') as f:
  147. graph_def = tf.compat.v1.GraphDef()
  148. graph_def.ParseFromString(f.read())
  149. sess.graph.as_default()
  150. tf.import_graph_def(graph_def, name='')
  151. graph_nodes = [n for n in graph_def.node]
  152. logits = sess.graph.get_tensor_by_name("visual_observation_0:0")
  153. CozmoController(sess, logits)