Spaces:
Sleeping
Sleeping
File size: 4,739 Bytes
f87d582 |
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 |
import math
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import FuncAnimation, FFMpegFileWriter
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import mpl_toolkits.mplot3d.axes3d as p3
# import cv2
from textwrap import wrap
from moviepy.editor import VideoClip
from moviepy.video.io.bindings import mplfig_to_npimage
def list_cut_average(ll, intervals):
if intervals == 1:
return ll
bins = math.ceil(len(ll) * 1.0 / intervals)
ll_new = []
for i in range(bins):
l_low = intervals * i
l_high = l_low + intervals
l_high = l_high if l_high < len(ll) else len(ll)
ll_new.append(np.mean(ll[l_low:l_high]))
return ll_new
def plot_3d_motion(save_path, kinematic_tree, joints, title, dataset, figsize=(3, 3), fps=120, radius=3,
vis_mode='default', gt_frames=[]):
matplotlib.use('Agg')
title_per_frame = type(title) == list
if title_per_frame:
assert len(title) == len(joints), 'Title length should match the number of frames'
title = ['\n'.join(wrap(s, 20)) for s in title]
else:
title = '\n'.join(wrap(title, 20))
def init():
ax.set_xlim3d([-radius / 2, radius / 2])
ax.set_ylim3d([0, radius])
ax.set_zlim3d([-radius / 3., radius * 2 / 3.])
# print(title)
# fig.suptitle(title, fontsize=10) # Using dynamic title instead
ax.grid(b=False)
def plot_xzPlane(minx, maxx, miny, minz, maxz):
## Plot a plane XZ
verts = [
[minx, miny, minz],
[minx, miny, maxz],
[maxx, miny, maxz],
[maxx, miny, minz]
]
xz_plane = Poly3DCollection([verts])
xz_plane.set_facecolor((0.5, 0.5, 0.5, 0.5))
ax.add_collection3d(xz_plane)
# return ax
# (seq_len, joints_num, 3)
data = joints.copy().reshape(len(joints), -1, 3)
# preparation related to specific datasets
if dataset == 'kit':
data *= 0.003 # scale for visualization
elif dataset == 'humanml':
data *= 1.3 # scale for visualization
elif dataset in ['humanact12', 'uestc']:
data *= -1.5 # reverse axes, scale for visualization
fig = plt.figure(figsize=figsize)
plt.tight_layout()
ax = p3.Axes3D(fig)
init()
MINS = data.min(axis=0).min(axis=0)
MAXS = data.max(axis=0).max(axis=0)
colors_blue = ["#4D84AA", "#5B9965", "#61CEB9", "#34C1E2", "#80B79A"] # GT color
colors_orange = ["#DD5A37", "#D69E00", "#B75A39", "#FF6D00", "#DDB50E"] # Generation color
colors = colors_orange
if vis_mode == 'upper_body': # lower body taken fixed to input motion
colors[0] = colors_blue[0]
colors[1] = colors_blue[1]
elif vis_mode == 'gt':
colors = colors_blue
n_frames = data.shape[0]
# print(dataset.shape)
height_offset = MINS[1]
data[:, :, 1] -= height_offset
trajec = data[:, 0, [0, 2]] # memorize original x,z pelvis values
# locate x,z pelvis values of ** each frame ** at zero
data[..., 0] -= data[:, 0:1, 0]
data[..., 2] -= data[:, 0:1, 2]
# print(trajec.shape)
def update(index):
# sometimes index is equal to n_frames/fps due to floating point issues. in such case, we duplicate the last frame
index = min(n_frames-1, int(index*fps))
ax.clear()
ax.view_init(elev=120, azim=-90)
ax.dist = 7.5
# Dynamic title
if title_per_frame:
_title = title[index]
else:
_title = title
_title += f' [{index}]'
fig.suptitle(_title, fontsize=10)
plot_xzPlane(MINS[0] - trajec[index, 0], MAXS[0] - trajec[index, 0], 0, MINS[2] - trajec[index, 1],
MAXS[2] - trajec[index, 1])
used_colors = colors_blue if index in gt_frames else colors
for i, (chain, color) in enumerate(zip(kinematic_tree, used_colors)):
if i < 5:
linewidth = 4.0
else:
linewidth = 2.0
ax.plot3D(data[index, chain, 0], data[index, chain, 1], data[index, chain, 2], linewidth=linewidth,
color=color)
# print(trajec[:index, 0].shape)
plt.axis('off')
ax.set_axis_off()
ax.set_xticklabels([])
ax.set_yticklabels([])
ax.set_zticklabels([])
# Hide grid lines
ax.grid(False)
# Hide axes ticks
ax.set_xticks([])
ax.set_yticks([])
ax.set_zticks([])
return mplfig_to_npimage(fig)
ani = VideoClip(update)
plt.close()
return ani |