main loop for two arms rough draft (#5)

src/KeyboardController.py

@@ -595,10 +595,12 @@ class KeyboardController(Controller):
 
 if __name__=="__main__":
     #con = KeyboardController(disableExtruder=True)
-    con = KeyboardController(disableExtruder=True, robot=2)
+    con1 = KeyboardController(disableExtruder=True, robot=1)
+    con2 = KeyboardController(disableExtruder=True, robot=2)
     #con = KeyboardController()
 
     print('Collecting configrations...')
-    configuations = con.prompt_configurations()
+    configuations = con1.prompt_configurations()
+    configuations = con2.prompt_configurations()
     print('Your configurations:')
     print(configuations)

src/main2.py

+import numpy as np
+import pickle
+import time
+import copy
+import sys
+
+# ROS Libraries
+import rospy
+from std_msgs.msg import Bool
+from trajectory_msgs.msg import (
+    JointTrajectoryPoint,
+    JointTrajectory
+)
+
+## Our libraries
+from utils import load_plan, JTraj2ROS, speed_multiplier
+from Controller import Controller
+import gcode2contour as gc
+
+
+def execute_plan(plan,
+        r1_con,
+        r2_con,
+        contours=None,
+        contour_speed=0.2,
+        move_speed=1,
+        n_trajs=None):
+    """
+    Executes a plan on MoveIt
+
+    Single arm only
+    skips the first n trajectories in the plan
+    """
+    # plan to the start position of the plan
+    r1_con.group.go(plan.trajs[0][0].positions[0])
+    r2_con.group.go(plan.trajs[1][0].positions[0])
+    controllers = [r1_con, r2_con]
+
+    # execute trajectories in plan
+    plan_len = max([a[-1] for a in plan.cumul_time.values()])
+
+    s = time.time()
+    expected_end_time = s + plan_len
+    i_r = [0,0]
+    while time.time() < expected_end_time:
+
+        # Find which arm and which trajectory is next to be executed
+        next_arm = np.argmin([plan.cumul_time[0][i_r[0]], plan.cumul_time[1][i_r[1]]])
+        next_time = min(plan.cumul_time[0][i_r[0]], plan.cumul_time[1][i_r[1]])
+        i_r[next_arm] += 1
+        t = plan.trajs[next_arm][i_r[next_arm]]
+
+        # wait till time to execute next trajectory
+        time.sleep(next_time - (time.time() - s ) - 1) # 1s extra
+
+        print("Printing {}th trajectory in the plan, which is contour {}".format(i, t.contour))
+        if t.contour is not None and contours is not None:
+            controllers[next_arm].exec_ctraj(speed_multiplier(t, 1/contour_speed), contour=contours[t.contour])
+        else:
+            controllers[next_arm].exec_ctraj(speed_multiplier(t, 1/move_speed))
+
+    return
+
+
+if __name__ == '__main__':
+    plan2arms = load_plan('two_arm_test', contours=True)
+
+    r1_con = Controller(disable_extruder=True, robot=1)
+    r2_con = Controller(disable_extruder=True, robot=2)
+
+    execute_plan(plan2arms, r1_con, r2_con, trajs=200, move_speed=0.5)