Trajectory visualizer

trajectory_display.py

#!/usr/bin/env python2

import rospy
import tf
import tf2_ros
import tf.transformations
from visualization_msgs.msg import Marker, MarkerArray
from geometry_msgs.msg import PoseStamped

import copy
import csv
import sys


def get_poses_from_file(fh):
    map_frame_id = rospy.get_param('~map_frame_id', 'map')
    base_frame_id = rospy.get_param('~base_frame_id', 'base_link')

    tf_l = tf.TransformListener()

    # We want the first point in the trajectory to be at the center of the vehicle
    trajectory_offset = [None, None]

    poses = list()

    csv_file = csv.reader(fh)
    # Skip header
    csv_file.next()
    for line in csv_file:
        _, x, y, theta, _ = line

        x = float(x)
        y = float(y)
        theta = float(theta)

        if trajectory_offset == [None, None]:
            trajectory_offset = [x, y]

        x -= trajectory_offset[0]
        y -= trajectory_offset[1]

        p = PoseStamped()
        p.header.frame_id = base_frame_id
        p.header.stamp = rospy.Time.now()
        p.pose.position.x = x
        p.pose.position.y = y
        p.pose.position.z = 0
        q = tf.transformations.quaternion_from_euler(0, 0, theta)
        p.pose.orientation.x = q[0]
        p.pose.orientation.y = q[1]
        p.pose.orientation.z = q[2]
        p.pose.orientation.w = q[3]

        # Transform into /map frame
        tf_l.waitForTransform(map_frame_id,
                              base_frame_id,
                              rospy.Time(0),
                              rospy.Duration(0.3))
        pose_in_map = tf_l.transformPose(map_frame_id, p)
        poses.append(pose_in_map)

    return poses


def plot_trajectory(poses):
    map_frame_id = rospy.get_param('~map_frame_id', 'map')
    stick_height = rospy.get_param('~stick_height', 0.3)
    trajectory_pub = rospy.Publisher('target_trajectory',
                          MarkerArray,
                          queue_size=1,
                          latch=True)

    ma = MarkerArray()

    # Line for the trajectory
    trajectory = Marker()
    trajectory.header.frame_id = map_frame_id
    trajectory.id = 0
    trajectory.type = Marker.LINE_STRIP
    trajectory.action = Marker.ADD
    trajectory.scale.x = 0.1
    trajectory.points = [pose.pose.position for pose in poses]
    trajectory.color.r = 0.0
    trajectory.color.g = 0.0
    trajectory.color.b = 1.0
    trajectory.color.a = 1.0
    ma.markers.append(trajectory)

    # Add sticks to the trajectory
    sticks = Marker()
    sticks.header.frame_id = map_frame_id
    sticks.id = 1
    sticks.type = Marker.LINE_LIST
    sticks.action = Marker.ADD
    sticks.scale.x = 0.03
    sticks.color.r = 0.0
    sticks.color.g = 0.6
    sticks.color.b = 0.8
    sticks.color.a = 1.0
    for pose in poses:
        floor = copy.deepcopy(pose.pose.position)
        sticks.points.append(floor)
        ceiling = copy.deepcopy(floor)
        ceiling.z += stick_height
        sticks.points.append(ceiling)
    ma.markers.append(sticks)

    trajectory_pub.publish(ma)


if __name__ == '__main__':
    rospy.init_node('trajectory_display')

    if len(sys.argv) < 2:
        sys.stderr.write('Not enough arguments:\n')
        sys.stderr.write('  {0} <csv file>\n'.format(sys.argv[0]))
        sys.exit(1)

    fn = sys.argv[1]
    print('Opening file: {0}'.format(fn))
    with open(fn) as fh:
        poses = get_poses_from_file(fh)
        plot_trajectory(poses)
        rospy.spin()