def draw_coordinate_axes(ax: plt.Figure,
traj: trajectory.PosePath3D,
plot_mode: PlotMode,
marker_scale: float = 0.1,
x_color: str = "r",
y_color: str = "g",
z_color: str = "b") -> None:
"""
Draws a coordinate frame axis for each pose of a trajectory.
:param ax: plot axis
:param traj: trajectory.PosePath3D or trajectory.PoseTrajectory3D object
:param plot_mode: PlotMode value
:param marker_scale: affects the size of the marker (1. * marker_scale)
:param x_color: color of the x-axis
:param y_color: color of the y-axis
:param z_color: color of the z-axis
"""
if marker_scale <= 0:
return
unit_x = np.array([1 * marker_scale, 0, 0, 1])
unit_y = np.array([0, 1 * marker_scale, 0, 1])
unit_z = np.array([0, 0, 1 * marker_scale, 1])
# Transform start/end vertices of each axis to global frame.
x_vertices = np.array([[p[:3, 3], p.dot(unit_x)[:3]]
for p in traj.poses_se3])
y_vertices = np.array([[p[:3, 3], p.dot(unit_y)[:3]]
for p in traj.poses_se3])
z_vertices = np.array([[p[:3, 3], p.dot(unit_z)[:3]]
for p in traj.poses_se3])
n = traj.num_poses
# Concatenate all line segment vertices in order x, y, z.
vertices = np.concatenate((x_vertices, y_vertices, z_vertices)).reshape(
(n * 2 * 3, 3))
# Concatenate all colors per line segment in order x, y, z.
colors = np.array(n * [x_color] + n * [y_color] + n * [z_color])
markers = colored_line_collection(vertices, colors, plot_mode, step=2)
ax.add_collection(markers)
