File size: 4,575 Bytes
45950ff | 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 | import cv2, io
import numpy as np
import torch
import matplotlib.pyplot as plt
import matplotlib
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import mpl_toolkits.mplot3d.axes3d as p3
from textwrap import wrap
def _vis_3d_motion(args, figsize=(10, 10), fps=120, radius=4):
matplotlib.use('Agg')
joints, out_name, title = args
data = joints.copy().reshape(len(joints), -1, 3)
nb_joints = joints.shape[1]
smpl_kinetic_chain = [
[0, 11, 12, 13, 14, 15],
[0, 16, 17, 18, 19, 20],
[0, 1, 2, 3, 4], [3, 5, 6, 7],
[3, 8, 9, 10]
] if nb_joints == 21 else [
[0, 2, 5, 8, 11], [0, 1, 4, 7, 10],
[0, 3, 6, 9, 12, 15], [9, 14, 17, 19, 21],
[9, 13, 16, 18, 20]
]
limits = 1000 if nb_joints == 21 else 2
MINS = data.min(axis=0).min(axis=0)
MAXS = data.max(axis=0).max(axis=0)
colors = ['red', 'blue', 'black', 'red', 'blue',
'darkblue', 'darkblue', 'darkblue', 'darkblue', 'darkblue',
'darkred', 'darkred', 'darkred', 'darkred', 'darkred']
frame_number = data.shape[0]
# print(data.shape)
height_offset = MINS[1]
data[:, :, 1] -= height_offset
trajec = data[:, 0, [0, 2]]
data[..., 0] -= data[:, 0:1, 0]
data[..., 2] -= data[:, 0:1, 2]
def update(index):
def init():
ax.set_xlim(-limits, limits)
ax.set_ylim(-limits, limits)
ax.set_zlim(0, limits)
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)
fig = plt.figure(figsize=(480/96., 320/96.), dpi=96) if nb_joints == 21 else plt.figure(figsize=(10, 10), dpi=96)
if title is not None :
wraped_title = '\n'.join(wrap(title, 40))
fig.suptitle(wraped_title, fontsize=16)
ax = p3.Axes3D(fig, auto_add_to_figure=False)
fig.add_axes(ax)
init()
for coll in ax.lines:
coll.remove()
for coll in ax.collections:
coll.remove()
ax.view_init(elev=110, azim=-90)
ax.dist = 7.5 # type: ignore
plot_xzPlane(MINS[0] - trajec[index, 0], MAXS[0] - trajec[index, 0], 0,
MINS[2] - trajec[index, 1], MAXS[2] - trajec[index, 1])
if index > 1:
ax.plot3D(trajec[:index, 0] - trajec[index, 0],
np.zeros_like(trajec[:index, 0]),
trajec[:index, 1] - trajec[index, 1],
linewidth=1.0, color='blue')
for i, (chain, color) in enumerate(zip(smpl_kinetic_chain, colors)):
linewidth = 4.0 if i < 5 else 2.0
ax.plot3D(data[index, chain, 0], data[index, chain, 1], data[index, chain, 2],
linewidth=linewidth, color=color)
plt.axis('off')
ax.set_xticklabels([])
ax.set_yticklabels([])
ax.set_zticklabels([]) # type: ignore
if out_name is not None :
plt.savefig(out_name, dpi=96)
plt.close()
else :
io_buf = io.BytesIO()
fig.savefig(io_buf, format='raw', dpi=96)
io_buf.seek(0)
arr = np.reshape(np.frombuffer(io_buf.getvalue(), dtype=np.uint8),
newshape=(int(fig.bbox.bounds[3]), int(fig.bbox.bounds[2]), -1))
io_buf.close()
plt.close()
return arr
out = []
for i in range(frame_number) :
out.append(update(i))
out = np.stack(out, axis=0)
return torch.from_numpy(out)
def vis_3d_motion(smpl_joints_batch, title_batch=None, outname=None, fps=30) :
out = []
for i in range(len(smpl_joints_batch)) :
out.append(_vis_3d_motion([smpl_joints_batch[i], None, title_batch[i] if title_batch is not None else None]))
if outname is not None:
import imageio.v3 as iio
images = np.array(out[-1]).astype(np.uint8)[..., :3]
with iio.imopen(outname[i], "w", plugin="pyav") as writer:
writer.init_video_stream("libx264", fps=int(fps))
writer._video_stream.options = {"crf": str(17)}
writer.write(images)
|