def plot_skeleton_on_axes3d(skel,
skel_desc,
ax: Axes3D,
invert=True,
alpha=1.0):
ax.set_xlabel('x')
ax.set_ylabel('z')
ax.set_zlabel('y')
# NOTE: y and z axes are swapped
xs = skel.narrow(-1, 0, 1).numpy()
ys = skel.narrow(-1, 2, 1).numpy()
zs = skel.narrow(-1, 1, 1).numpy()
# Correct aspect ratio (https://stackoverflow.com/a/21765085)
max_range = np.array(
[xs.max() - xs.min(),
ys.max() - ys.min(),
zs.max() - zs.min()]).max() / 2.0
mid_x = (xs.max() + xs.min()) * 0.5
mid_y = (ys.max() + ys.min()) * 0.5
mid_z = (zs.max() + zs.min()) * 0.5
ax.set_xlim(mid_x - max_range, mid_x + max_range)
ax.set_ylim(mid_y - max_range, mid_y + max_range)
ax.set_zlim(mid_z - max_range, mid_z + max_range)
ax.set_aspect('equal')
if invert:
ax.invert_zaxis()
# Set starting view
ax.view_init(elev=20, azim=-100)
get_joint_metadata = _make_joint_metadata_fn(skel_desc)
for joint_id, joint in enumerate(skel):
meta = get_joint_metadata(joint_id)
color = 'magenta'
if meta['group'] == 'left':
color = 'blue'
if meta['group'] == 'right':
color = 'red'
parent = skel[meta['parent']]
offset = parent - joint
ax.quiver(
[joint[0]],
[joint[2]],
[joint[1]],
[offset[0]],
[offset[2]],
[offset[1]],
color=color,
alpha=alpha,
)
ax.scatter(xs, ys, zs, color='grey', alpha=alpha)
