Update runner_script

scripts/runner_script.py

@@ -2,6 +2,9 @@
 from lab2.msg import XYHV, XYHVPath
 from lab2.srv import FollowPath
 from std_msgs.msg import Header
+from geometry_msgs.msg import PoseStamped
+import tf.transformations
+import tf
 
 import numpy as np
 import pickle
@@ -10,22 +13,20 @@ import rospy
 
 
 def saw():
-    """
-    Generates a sawtooth path
-    """
-    # Change alpha to get a different frequency
     t = np.linspace(0, 20, 100)
-    alpha = 0.5
-    saw = signal.sawtooth(alpha * np.pi * t)
+    saw = signal.sawtooth(0.5 * np.pi * t)
     configs = [[x, y, 0] for (x, y) in zip(t, saw)]
     return configs
 
+def wave():
+    t = np.linspace(0, 20, 100)
+    y = np.sin(t)
+    theta = np.cos(t)
+    configs = [[x, y, _theta] for (x, y, _theta) in zip(t, y, theta)]
+    return configs
+
 
 def circle():
-    """
-    Generates a circular path.
-    """
-    # Change radius to modify the size of the circle
     waypoint_sep = 0.1
     radius = 2.5
     center = [0, radius]
@@ -36,10 +37,6 @@ def circle():
 
 
 def left_turn():
-    """
-    Generates a path that goes straight and turns left 90 degrees.
-    """
-    # Change turn_radius and straight_len to modify the path
     waypoint_sep = 0.1
     turn_radius = 1.5
     straight_len = 10.0
@@ -54,10 +51,6 @@ def left_turn():
 
 
 def right_turn():
-    """
-    Generates a path that goes straight and turns right 90 degrees.
-    """
-    # Change turn_radius and straight_len to modify the path
     waypoint_sep = 0.1
     turn_radius = 1.5
     straight_len = 10.0
@@ -72,14 +65,45 @@ def right_turn():
 
 
 def cse022_path():
-    with open('cse022_path.pickle', 'r') as f:
+    import os
+    dir_path = os.path.dirname(os.path.realpath(__file__))
+    with open(os.path.join(dir_path, 'cse022_path.pickle'), 'r') as f:
         p = pickle.load(f)
     return p
 
 
-plans = {'circle': circle, 'left turn': left_turn, 'right turn': right_turn, 'saw': saw, 'cse022 real path': cse022_path}
-plan_names = ['circle', 'left turn', 'right turn', 'saw', 'cse022 real path']
-
+plans = {'circle': circle, 'left turn': left_turn, 'right turn': right_turn, 'wave': wave, 'cse022 real path': cse022_path}
+plan_names = ['circle', 'left turn', 'right turn', 'wave', 'cse022 real path']
+
+def get_current_pose():
+    car_pose_topic = \
+            ('/car_pose' if int(rospy.get_param('/controller/use_sim_pose')) == 1
+            else '/pf/inferred_pose')
+    print("Listening to {} for initial pose".format(car_pose_topic))
+    car_pose_msg = rospy.wait_for_message(car_pose_topic, PoseStamped)
+    tf_listener = tf.TransformListener()
+    msg = tf_listener.transformPose("map", car_pose_msg)
+    car_pose = [msg.pose.position.x,
+            msg.pose.position.y,
+            rosquaternion_to_angle(msg.pose.orientation)]
+    return car_pose
+
+def rosquaternion_to_angle(q):
+    """Convert a quaternion _message_ into an angle in radians.
+    The angle represents the yaw.
+    This is not just the z component of the quaternion."""
+    x, y, z, w = q.x, q.y, q.z, q.w
+    _, _, yaw = tf.transformations.euler_from_quaternion((x, y, z, w))
+    return yaw
+
+def shift_zero_pose(configs, shift_by):
+    # TODO (optional) you may edit how you shift all poses
+    shift_by[2] = 0
+    shifted_configs = []
+    _configs = np.array(configs) + np.array(shift_by)
+    for c in _configs:
+        shifted_configs.append(list(c))
+    return shifted_configs
 
 def generate_plan():
     print "Which plan would you like to generate? "
@@ -89,7 +113,9 @@ def generate_plan():
     if index >= len(plan_names):
         print "Wrong number. Exiting."
         exit()
-    return plans[plan_names[index]]()
+    if plan_names[index] == 'cse022 real path':
+        return plans[plan_names[index]]()
+    return shift_zero_pose(plans[plan_names[index]](), get_current_pose())
 
 
 if __name__ == '__main__':