-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathvisualization.py
177 lines (150 loc) · 5.58 KB
/
visualization.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
'''Functions for generating visualizations from NiTE skeletons'''
import logging
from math import inf
from matplotlib import pyplot as plt
from matplotlib import animation
from mpl_toolkits.mplot3d import Axes3D
from skeletons import (NITE_JOINT_HEAD, NITE_JOINT_LEFT_ELBOW,
NITE_JOINT_LEFT_FOOT, NITE_JOINT_LEFT_HAND,
NITE_JOINT_LEFT_HIP, NITE_JOINT_LEFT_KNEE,
NITE_JOINT_LEFT_SHOULDER, NITE_JOINT_NECK,
NITE_JOINT_RIGHT_ELBOW, NITE_JOINT_RIGHT_FOOT,
NITE_JOINT_RIGHT_HAND, NITE_JOINT_RIGHT_HIP,
NITE_JOINT_RIGHT_KNEE, NITE_JOINT_RIGHT_SHOULDER,
NITE_JOINT_TORSO)
# pylint: disable=C0103
log = logging.getLogger(__name__)
links = [
(NITE_JOINT_HEAD, NITE_JOINT_NECK),
(NITE_JOINT_NECK, NITE_JOINT_LEFT_SHOULDER),
(NITE_JOINT_NECK, NITE_JOINT_RIGHT_SHOULDER),
(NITE_JOINT_LEFT_SHOULDER, NITE_JOINT_LEFT_ELBOW),
(NITE_JOINT_LEFT_ELBOW, NITE_JOINT_LEFT_HAND),
(NITE_JOINT_RIGHT_SHOULDER, NITE_JOINT_RIGHT_ELBOW),
(NITE_JOINT_RIGHT_ELBOW, NITE_JOINT_RIGHT_HAND),
(NITE_JOINT_NECK, NITE_JOINT_TORSO),
(NITE_JOINT_TORSO, NITE_JOINT_LEFT_HIP),
(NITE_JOINT_LEFT_HIP, NITE_JOINT_LEFT_KNEE),
(NITE_JOINT_LEFT_KNEE, NITE_JOINT_LEFT_FOOT),
(NITE_JOINT_TORSO, NITE_JOINT_RIGHT_HIP),
(NITE_JOINT_RIGHT_HIP, NITE_JOINT_RIGHT_KNEE),
(NITE_JOINT_RIGHT_KNEE, NITE_JOINT_RIGHT_FOOT)
]
def find_bounds(skeleton_data):
'''Utility function to find the limits of each dimension
Args:
skeleton_data ((skeleton, offset) list): The data to animate
Returns:
((x min, x_max), (y min, y max), (z min, z max))
'''
x_min = inf
x_max = -inf
y_min = inf
y_max = -inf
z_min = inf
z_max = -inf
for frame in skeleton_data:
for joint in frame[0].joints:
x_min, x_max = min(x_min, joint.position.x), max(x_max, joint.position.x)
y_min, y_max = min(y_min, joint.position.y), max(y_max, joint.position.y)
z_min, z_max = min(z_min, joint.position.z), max(z_max, joint.position.z)
return ((x_min, x_max), (y_min, y_max), (z_min, z_max))
def initialize_plots(skeleton_data):
'''Utility function to find the limits of each dimension
Args:
skeleton_data ((skeleton, offset) list): The data to animate
Returns:
figure, axes, plots
'''
fig = plt.figure()
axes = fig.add_subplot(111, projection='3d')
axes.view_init(elev=270, azim=90)
(x_min, x_max), (y_min, y_max), (z_min, z_max) = find_bounds(skeleton_data)
# We do seemingly weird assignments of bounds here, but it's due to the conventions used by NiTE
axes.set_xlabel('X')
axes.set_ylabel('Y')
axes.set_zlabel('Z')
axes.set_xlim(left=x_max, right=x_min)
axes.set_ylim(bottom=y_min, top=y_max)
axes.set_zlim(bottom=z_max, top=z_min)
log.info('Making inital plots')
plots = {link: axes.plot([0], [0], [0], 'bo-')[0] for link in links}
return fig, axes, plots
def make_video(skeleton_data, out_filename, show):
'''Creates an animation from skeleton data
Args:
skeleton_data ((skeleton, offset) list): The data to animate
out_filename (string): The file in which to save the video
show (bool): Whether or not to preview the video before saving
Returns:
Nothing
'''
# Note that we make the assumption that the interval between frames is equal throughout. This
# seems to hold mostly true, empirically.
interval_microseconds = skeleton_data[1][1] - skeleton_data[0][1]
interval = int(interval_microseconds / 1e3)
fig, axes, plots = initialize_plots(skeleton_data)
def init():
'''Initializes the animation'''
for link in links:
plots[link].set_xdata([0])
plots[link].set_ydata([0])
plots[link].set_3d_properties([0])
return tuple(plots.values())
def animate(i):
'''Render each frame'''
frame = skeleton_data[i][0]
skel = {link: None for link in links}
for joint in frame.joints:
skel[joint.jointType] = joint.position
for link in links:
top, bottom = link
if skel[top] and skel[bottom]:
xs = (skel[top].x, skel[bottom].x)
ys = (skel[top].y, skel[bottom].y)
zs = (skel[top].z, skel[bottom].z)
plots[link].set_xdata(xs)
plots[link].set_ydata(ys)
plots[link].set_3d_properties(zs)
return tuple(plots.values())
log.info('Creating animation')
video = animation.FuncAnimation(
fig,
animate,
init_func=init,
frames=len(skeleton_data),
interval=interval,
blit=True
)
if show:
plt.show()
log.info(f'Saving video to {out_filename}')
video.save(out_filename, fps=30, extra_args=['-vcodec', 'libx264'])
def make_image(skeleton_data, out_filename, show):
'''Generate a static image of the skeleton
Args:
skeleton_data ((skeleton, offset) list): The data to visualize. The first provided frame is
chosen.
out_filename (string): The file in which to save the image.
show (bool): Whether or not to preview the video before saving
Returns:
Nothing
'''
fig, axes, plots = initialize_plots(skeleton_data)
frame = skeleton_data[0][0]
skel = {link: None for link in links}
for joint in frame.joints:
skel[joint.jointType] = joint.position
for link in links:
top, bottom = link
if skel[top] and skel[bottom]:
xs = (skel[top].x, skel[bottom].x)
ys = (skel[top].y, skel[bottom].y)
zs = (skel[top].z, skel[bottom].z)
plots[link].set_xdata(xs)
plots[link].set_ydata(ys)
plots[link].set_3d_properties(zs)
if show:
plt.show()
log.info(f'Saving image to {out_filename}')
plt.savefig(out_filename)