KinematicChainPathPlot.py
1 #!/usr/bin/env python
2 
3 
36 
37 # Author: Mark Moll
38 
39 # This is a helper program to visualize the output of the KinematicChainBenchmark.cpp.
40 # Sample usage (when running from the OMPL build directory):
41 #
42 # ./bin/demo_KinematicChainBenchmark 10 1
43 # /path/to/KinematicChainPathPlot.py 10 # produces kinematic_10.pdf
44 # /path/to/KinematicChainPathPlot.py 10 movie # produces kinematic_10.mp4
45 
46 from sys import argv, exit
47 import numpy as np
48 import matplotlib.pyplot as plt
49 from matplotlib.patches import PathPatch
50 from matplotlib.path import Path
51 from matplotlib import animation, cm, colors
52 from matplotlib.backends.backend_pdf import PdfPages
53 
54 fig = plt.figure()
55 plt.axis('equal')
56 ax = plt.axes(xlim=(-1, 1), ylim=(-1, 1))
57 ln, = plt.plot([], [], animated=True)
58 
59 def getPositions(pose, linkLength, color = 'red'):
60  angles = np.cumsum(pose)
61  return PathPatch(Path(
62  np.insert(linkLength * np.cumsum([np.cos(angles), np.sin(angles)],
63  axis = 1), 0, 0., axis = 1).T), facecolor='none', edgecolor = color)
64 
65 def drawEnvironment(env):
66  ax.clear()
67  plt.gca().add_patch(env)
68 
69 def drawPose(index, env, poses, linkLength):
70  drawEnvironment(env)
71  if index == -1:
72  cMap = cm.ScalarMappable(
73  norm = colors.Normalize(vmin=0, vmax=poses.shape[0] - 1),
74  cmap = plt.get_cmap('viridis') )
75  for (i,pose) in enumerate(poses):
76  plt.gca().add_patch(getPositions(pose, linkLength,
77  cMap.to_rgba(i)))
78  else:
79  plt.gca().add_patch(getPositions(poses[index,:], linkLength))
80  return ln,
81 
82 def makeMovie(fname, env, poses, linkLength):
83  ani = animation.FuncAnimation(plt.gcf(), drawPose,
84  fargs=(env, poses, linkLength), frames = poses.shape[0],
85  blit = True)
86  ani.save(fname, bitrate = 300, fps = 20)
87 
88 if __name__ == '__main__':
89  if len(argv) < 2:
90  print('Usage: %s num_links [movie]' % argv[0])
91  exit(1)
92 
93  dims = int(argv[1])
94  coords = np.loadtxt('environment_%d.dat' % dims)
95  env = PathPatch(Path(coords), facecolor='none', edgecolor='black')
96  poses = np.loadtxt('kinematic_path_%d.dat' % dims)
97  linkLength = 1. / dims
98 
99  if len(argv) > 2:
100  makeMovie('kinematic_%d.mp4' % dims, env, poses, linkLength)
101  else:
102  drawPose(-1, env, poses, linkLength)
103  fig = plt.gcf()
104  fig.savefig('kinematic_%d.pdf' % dims)