252 lines
8.8 KiB
Python
252 lines
8.8 KiB
Python
"""
|
|
This code is borrowed from https://github.com/EricGuo5513/text-to-motion
|
|
Modified to be compatible with matplotlib 3.8.0+
|
|
"""
|
|
|
|
import torch
|
|
import numpy as np
|
|
|
|
import math
|
|
import matplotlib
|
|
import matplotlib.pyplot as plt
|
|
from mpl_toolkits.mplot3d import Axes3D
|
|
from matplotlib.animation import FuncAnimation, PillowWriter
|
|
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
|
|
import mpl_toolkits.mplot3d.axes3d as p3
|
|
from tempfile import TemporaryDirectory
|
|
from pathlib import Path
|
|
from uuid import uuid4
|
|
|
|
# Define a kinematic tree for the skeletal struture
|
|
kit_kinematic_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]]
|
|
|
|
kit_raw_offsets = np.array(
|
|
[
|
|
[0, 0, 0],
|
|
[0, 1, 0],
|
|
[0, 1, 0],
|
|
[0, 1, 0],
|
|
[0, 1, 0],
|
|
[1, 0, 0],
|
|
[0, -1, 0],
|
|
[0, -1, 0],
|
|
[-1, 0, 0],
|
|
[0, -1, 0],
|
|
[0, -1, 0],
|
|
[1, 0, 0],
|
|
[0, -1, 0],
|
|
[0, -1, 0],
|
|
[0, 0, 1],
|
|
[0, 0, 1],
|
|
[-1, 0, 0],
|
|
[0, -1, 0],
|
|
[0, -1, 0],
|
|
[0, 0, 1],
|
|
[0, 0, 1]
|
|
]
|
|
)
|
|
|
|
t2m_raw_offsets = np.array([[0,0,0],
|
|
[1,0,0],
|
|
[-1,0,0],
|
|
[0,1,0],
|
|
[0,-1,0],
|
|
[0,-1,0],
|
|
[0,1,0],
|
|
[0,-1,0],
|
|
[0,-1,0],
|
|
[0,1,0],
|
|
[0,0,1],
|
|
[0,0,1],
|
|
[0,1,0],
|
|
[1,0,0],
|
|
[-1,0,0],
|
|
[0,0,1],
|
|
[0,-1,0],
|
|
[0,-1,0],
|
|
[0,-1,0],
|
|
[0,-1,0],
|
|
[0,-1,0],
|
|
[0,-1,0]])
|
|
|
|
t2m_kinematic_chain = [[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]]
|
|
t2m_left_hand_chain = [[20, 22, 23, 24], [20, 34, 35, 36], [20, 25, 26, 27], [20, 31, 32, 33], [20, 28, 29, 30]]
|
|
t2m_right_hand_chain = [[21, 43, 44, 45], [21, 46, 47, 48], [21, 40, 41, 42], [21, 37, 38, 39], [21, 49, 50, 51]]
|
|
|
|
|
|
def qinv(q):
|
|
assert q.shape[-1] == 4, 'q must be a tensor of shape (*, 4)'
|
|
mask = torch.ones_like(q)
|
|
mask[..., 1:] = -mask[..., 1:]
|
|
return q * mask
|
|
|
|
|
|
def qrot(q, v):
|
|
"""
|
|
Rotate vector(s) v about the rotation described by quaternion(s) q.
|
|
Expects a tensor of shape (*, 4) for q and a tensor of shape (*, 3) for v,
|
|
where * denotes any number of dimensions.
|
|
Returns a tensor of shape (*, 3).
|
|
"""
|
|
assert q.shape[-1] == 4
|
|
assert v.shape[-1] == 3
|
|
assert q.shape[:-1] == v.shape[:-1]
|
|
|
|
original_shape = list(v.shape)
|
|
# print(q.shape)
|
|
q = q.contiguous().view(-1, 4)
|
|
v = v.contiguous().view(-1, 3)
|
|
|
|
qvec = q[:, 1:]
|
|
uv = torch.cross(qvec, v, dim=1)
|
|
uuv = torch.cross(qvec, uv, dim=1)
|
|
return (v + 2 * (q[:, :1] * uv + uuv)).view(original_shape)
|
|
|
|
|
|
def recover_root_rot_pos(data):
|
|
rot_vel = data[..., 0]
|
|
r_rot_ang = torch.zeros_like(rot_vel).to(data.device)
|
|
'''Get Y-axis rotation from rotation velocity'''
|
|
r_rot_ang[..., 1:] = rot_vel[..., :-1]
|
|
r_rot_ang = torch.cumsum(r_rot_ang, dim=-1)
|
|
|
|
r_rot_quat = torch.zeros(data.shape[:-1] + (4,)).to(data.device)
|
|
r_rot_quat[..., 0] = torch.cos(r_rot_ang)
|
|
r_rot_quat[..., 2] = torch.sin(r_rot_ang)
|
|
|
|
r_pos = torch.zeros(data.shape[:-1] + (3,)).to(data.device)
|
|
r_pos[..., 1:, [0, 2]] = data[..., :-1, 1:3]
|
|
'''Add Y-axis rotation to root position'''
|
|
r_pos = qrot(qinv(r_rot_quat), r_pos)
|
|
|
|
r_pos = torch.cumsum(r_pos, dim=-2)
|
|
|
|
r_pos[..., 1] = data[..., 3]
|
|
return r_rot_quat, r_pos
|
|
|
|
|
|
def recover_from_ric(data, joints_num):
|
|
r_rot_quat, r_pos = recover_root_rot_pos(data)
|
|
positions = data[..., 4:(joints_num - 1) * 3 + 4]
|
|
positions = positions.view(positions.shape[:-1] + (-1, 3))
|
|
|
|
'''Add Y-axis rotation to local joints'''
|
|
positions = qrot(qinv(r_rot_quat[..., None, :]).expand(positions.shape[:-1] + (4,)), positions)
|
|
|
|
'''Add root XZ to joints'''
|
|
positions[..., 0] += r_pos[..., 0:1]
|
|
positions[..., 2] += r_pos[..., 2:3]
|
|
|
|
'''Concate root and joints'''
|
|
positions = torch.cat([r_pos.unsqueeze(-2), positions], dim=-2)
|
|
|
|
return positions
|
|
|
|
def plot_3d_motion(save_path, kinematic_tree, joints, distance, elevation, rotation, poselinewidth, title, figsize=(10, 10), fps=120, radius=4, visualization="original", save_as_pil_lists=False):
|
|
matplotlib.use('Agg')
|
|
|
|
title_sp = title.split(' ')
|
|
if len(title_sp) > 20:
|
|
title = '\n'.join([' '.join(title_sp[:10]), ' '.join(title_sp[10:20]), ' '.join(title_sp[20:])])
|
|
elif len(title_sp) > 10:
|
|
title = '\n'.join([' '.join(title_sp[:10]), ' '.join(title_sp[10:])])
|
|
|
|
def init():
|
|
ax.set_xlim3d([-radius / 4, radius / 4])
|
|
ax.set_ylim3d([0, radius / 2])
|
|
ax.set_zlim3d([0, radius / 2])
|
|
if visualization == "original":
|
|
fig.suptitle(title, fontsize=20)
|
|
ax.grid(b=False)
|
|
fig.add_axes(ax) #As a user, you do not instantiate Axes directly, but use Axes creation methods instead; e.g. from pyplot or Figure: subplots, subplot_mosaic or Figure.add_axes.
|
|
if visualization == "pseudo-openpose":
|
|
ax.set_facecolor("black")
|
|
|
|
def plot_xzPlane(minx, maxx, miny, minz, maxz):
|
|
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)
|
|
|
|
# (seq_len, joints_num, 3)
|
|
data = joints.copy().reshape(len(joints), -1, 3)
|
|
fig = plt.figure(figsize=figsize)
|
|
ax = p3.Axes3D(fig)
|
|
init()
|
|
MINS = data.min(axis=0).min(axis=0)
|
|
MAXS = data.max(axis=0).max(axis=0)
|
|
if visualization == "pseudo-openpose":
|
|
colors = [('#009900', '#009900', '#009933', '#009966'), #left leg
|
|
('#009999', '#009999', '#006699', '#003399'), #right leg
|
|
('#000099', ), #body
|
|
('#990000', '#990000', '#996600', '#999900'), #left arm
|
|
('#993300', '#993300', '#669900', '#339900'), #right arm
|
|
'darkblue', 'darkblue', 'darkblue', 'darkblue', 'darkblue',
|
|
'darkred', 'darkred', 'darkred', 'darkred', 'darkred']
|
|
else:
|
|
colors = ['red', 'blue', 'black', 'red', 'blue',
|
|
'darkblue', 'darkblue', 'darkblue', 'darkblue', 'darkblue',
|
|
'darkred', 'darkred', 'darkred', 'darkred', 'darkred']
|
|
|
|
|
|
frame_number = data.shape[0]
|
|
|
|
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):
|
|
for line in ax.lines:
|
|
line.remove()
|
|
for collection in ax.collections:
|
|
collection.remove()
|
|
ax.view_init(elev=elevation, azim=rotation)
|
|
ax._dist = distance #ax.dist = 7.5, https://matplotlib.org/stable/api/prev_api_changes/api_changes_3.8.0.html#axes3d
|
|
if visualization == "original":
|
|
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(kinematic_tree, colors)):
|
|
if i < 5:
|
|
linewidth = 4.0
|
|
else:
|
|
linewidth = 2.0
|
|
if visualization == "pseudo-openpose":
|
|
for j in range(len(data[index, chain, 0]) - 1):
|
|
ax.plot3D(data[index, chain, 0][j:j+2], data[index, chain, 1][j:j+2], data[index, chain, 2][j:j+2], linewidth=poselinewidth,
|
|
color=color[j % len(color)])
|
|
|
|
else:
|
|
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([])
|
|
|
|
ani = FuncAnimation(fig, update, frames=frame_number, interval=1000 / fps, repeat=False)
|
|
|
|
if save_as_pil_lists:
|
|
pil_writer = PillowWriter()
|
|
with TemporaryDirectory() as tmpdir:
|
|
ani.save(Path(tmpdir, f"{uuid4()}.gif"), writer=pil_writer)
|
|
plt.close()
|
|
return pil_writer._frames
|
|
else:
|
|
ani.save(save_path, fps=fps)
|
|
plt.close()
|