visualize.py
1#!/usr/bin/env python
2
3
36
37# Author: Ryan Luna
38
39import sys
40import yaml
41from math import sin, cos
42import numpy as np
43import matplotlib.pyplot as plt
44import matplotlib.animation as animation
45import matplotlib.patches as patches
46
47# Reads a YAML world specification from the worldFile.
48# Returns the set of obstacles and the bounds of the world.
49def readWorldFile(worldFile):
50 handle = open(worldFile, "r")
51 # use safe_load instead load
52 dataMap = yaml.safe_load(handle)
53 handle.close()
54
55 # Read in bounds
56 xlower = float(dataMap["bounds"]["x"]["lower"])
57 xupper = float(dataMap["bounds"]["x"]["upper"])
58 ylower = float(dataMap["bounds"]["y"]["lower"])
59 yupper = float(dataMap["bounds"]["y"]["upper"])
60
61 bounds = [[xlower, xupper], [ylower, yupper]]
62
63 obs = {}
64
65 if "obstacles" in dataMap:
66 for r in dataMap["obstacles"]:
67 for k in r.keys():
68 if k not in obs:
69 obs[k] = []
70 obs[k].append(r[k])
71 return obs, bounds
72
73
74def readPathFile(pathFile):
75 # Read in file line by line
76 lines = [line.rstrip() for line in open(pathFile) if len(line.rstrip()) > 0]
77
78 # first line is meta data
79 metadata = lines[0]
80 metadata = metadata.split(" ")
81 if len(metadata) != 5:
82 raise RuntimeError(
83 "Malformed path file. Expected first line with # links, link length, originX, originY, xySlices"
84 )
85
86 numLinks = int(metadata[0])
87 linkLength = float(metadata[1])
88 origin = (float(metadata[2]), float(metadata[3]))
89 slices = int(metadata[4])
90
91 path = []
92
93 for l in lines[1:]:
94 entries = l.split(" ")
95 if len(entries) != numLinks:
96 raise RuntimeError(
97 "Malformed path file. Path entries must have length = # links"
98 )
99 config = [float(e) for e in entries]
100 path.append(config)
101
102 return numLinks, linkLength, origin, path, slices
103
104
105def plotPolygon(axes, data, color, alpha=1.0, edgecolor=None):
106 points = []
107 for d in data:
108 for k in d.keys():
109 if k == "vertex":
110 coords = d[k].strip()
111 if coords.startswith("("):
112 coords = coords[1:-1]
113 coords = coords.strip().split(",")
114 pt = [float(coords[0]), float(coords[1])]
115 points.append(pt)
116
117 else:
118 raise RuntimeError('Expected "vertex", but got ', k)
119
120 if len(points) > 0:
121 arr = np.array(points)
122 axes.add_patch(
123 patches.Polygon(
124 arr, facecolor=color, alpha=alpha, fill=True, edgecolor=edgecolor
125 )
126 )
127
128
129def plotObstacles(axes, obstacles, bounds):
130 for k, v in obstacles.items():
131 for o in v:
132 if k == "polygon":
133 plotPolygon(axes, o, "0.4", edgecolor="0.4")
134 else:
135 raise RuntimeError("Unknown geometry type: ", k)
136
137 plt.axis([bounds[0][0], bounds[0][1], bounds[1][0], bounds[1][1]])
138
139
140def plotChain(axes, angles, origin, color="blue"):
141 x = origin[0]
142 y = origin[1]
143 angle = 0.0
144
145 linkLength = 1.0 / len(angles)
146 for theta in angles:
147 angle += theta
148
149 xN = x + (cos(angle) * linkLength)
150 yN = y + (sin(angle) * linkLength)
151
152 axes.plot([x, xN], [y, yN], "-", color=color)
153
154 x = xN
155 y = yN
156
157
158def plotChainFrame(frame_num, ax, path, obstacles, bounds, origin, rate, slices):
159 ax.clear()
160 plotObstacles(ax, obstacles, bounds)
161
162 # Hold for 1 second at the beginning and end
163 if frame_num < rate:
164 configuration = path[0]
165 elif frame_num >= len(path) + rate:
166 configuration = path[-1]
167 else:
168 configuration = path[frame_num - rate]
169
170 plotChain(ax, configuration, origin)
171
172
173def AnimatePath(obstacles, bounds, numLinks, linkLength, origin, path, slices):
174 fig = plt.figure()
175 axes = plt.axes()
176 framerate = 30
177
178 numFrames = len(path)
179 anim = animation.FuncAnimation(
180 fig,
181 plotChainFrame,
182 fargs=[axes, path, obstacles, bounds, origin, slices, framerate],
183 frames=numFrames + (framerate * 2),
184 interval=(1.0 / framerate) * 1000,
185 blit=False,
186 )
187
188 filename = "animation.mp4"
189 anim.save(filename, fps=framerate)
190
191
192if __name__ == "__main__":
193 worldFile = "./world.yaml"
194 pathFile = "./manipulator_path.txt"
195
196 obs, bounds = readWorldFile(worldFile)
197 numLinks, linkLength, origin, path, slices = readPathFile(pathFile)
198
199 print("Creating animation for planar kinematic chain... ")
200 AnimatePath(obs, bounds, numLinks, linkLength, origin, path, slices)
201 print("Done")