motionfix-demo / renderer /matplotlib.py
atnikos's picture
basic functionalities working
517683d
# From TEMOS: temos/render/anim.py
# Inspired by
# - https://github.com/anindita127/Complextext2animation/blob/main/src/utils/visualization.py
# - https://github.com/facebookresearch/QuaterNet/blob/main/common/visualization.py
import os
import logging
from dataclasses import dataclass
from typing import List, Tuple, Optional
import numpy as np
from src.tools.rifke import canonicalize_rotation
logger = logging.getLogger("matplotlib.animation")
logger.setLevel(logging.ERROR)
colors = ("black", "magenta", "red", "green", "blue")
KINEMATIC_TREES = {
"smpljoints": [
[0, 3, 6, 9, 12, 15],
[9, 13, 16, 18, 20],
[9, 14, 17, 19, 21],
[0, 1, 4, 7, 10],
[0, 2, 5, 8, 11],
],
"guoh3djoints": [ # no hands
[0, 3, 6, 9, 12, 15],
[9, 13, 16, 18, 20],
[9, 14, 17, 19, 21],
[0, 1, 4, 7, 10],
[0, 2, 5, 8, 11],
],
}
@dataclass
class MatplotlibRender:
jointstype: str = "smpljoints"
fps: float = 20.0
colors: List[str] = colors
figsize: int = 4
fontsize: int = 15
canonicalize: bool = False
def __call__(
self,
joints,
output,
fps=None,
highlights=None,
title: str = "",
canonicalize=None,
):
canonicalize = canonicalize if canonicalize is not None else self.canonicalize
fps = fps if fps is not None else self.fps
if joints.shape[1] == 24:
# remove the hands
joints = joints[:, :22]
render_animation(
joints,
title=title,
highlights=highlights,
output=output,
jointstype=self.jointstype,
fps=self.fps,
colors=self.colors,
figsize=(self.figsize, self.figsize),
fontsize=self.fontsize,
canonicalize=canonicalize,
)
def init_axis(fig, title, radius=1.5):
ax = fig.add_subplot(1, 1, 1, projection="3d")
ax.view_init(elev=20.0, azim=-60)
fact = 2
ax.set_xlim3d([-radius / fact, radius / fact])
ax.set_ylim3d([-radius / fact, radius / fact])
ax.set_zlim3d([0, radius])
ax.set_aspect("auto")
ax.set_xticklabels([])
ax.set_yticklabels([])
ax.set_zticklabels([])
ax.set_axis_off()
ax.grid(b=False)
ax.set_title(title, loc="center", wrap=True)
return ax
def plot_floor(ax, minx, maxx, miny, maxy, minz):
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
# Plot a plane XZ
verts = [
[minx, miny, minz],
[minx, maxy, minz],
[maxx, maxy, minz],
[maxx, miny, minz],
]
xz_plane = Poly3DCollection([verts], zorder=1)
xz_plane.set_facecolor((0.5, 0.5, 0.5, 1))
ax.add_collection3d(xz_plane)
# Plot a bigger square plane XZ
radius = max((maxx - minx), (maxy - miny))
# center +- radius
minx_all = (maxx + minx) / 2 - radius
maxx_all = (maxx + minx) / 2 + radius
miny_all = (maxy + miny) / 2 - radius
maxy_all = (maxy + miny) / 2 + radius
verts = [
[minx_all, miny_all, minz],
[minx_all, maxy_all, minz],
[maxx_all, maxy_all, minz],
[maxx_all, miny_all, minz],
]
xz_plane = Poly3DCollection([verts], zorder=1)
xz_plane.set_facecolor((0.5, 0.5, 0.5, 0.5))
ax.add_collection3d(xz_plane)
return ax
def update_camera(ax, root, radius=1.5):
fact = 2
ax.set_xlim3d([-radius / fact + root[0], radius / fact + root[0]])
ax.set_ylim3d([-radius / fact + root[1], radius / fact + root[1]])
def render_animation(
joints: np.ndarray,
output: str = "notebook",
highlights: Optional[np.ndarray] = None,
jointstype: str = "smpljoints",
title: str = "",
fps: float = 20.0,
colors: List[str] = colors,
figsize: Tuple[int] = (4, 4),
fontsize: int = 15,
canonicalize: bool = False,
agg=True,
):
if agg:
import matplotlib
matplotlib.use("Agg")
if highlights is not None:
assert len(highlights) == len(joints)
assert jointstype in KINEMATIC_TREES
kinematic_tree = KINEMATIC_TREES[jointstype]
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import matplotlib.patheffects as pe
mean_fontsize = fontsize
# heuristic to change fontsize
fontsize = mean_fontsize - (len(title) - 30) / 20
plt.rcParams.update({"font.size": fontsize})
# Z is gravity here
x, y, z = 0, 1, 2
joints = joints.copy()
if canonicalize:
joints = canonicalize_rotation(joints, jointstype=jointstype)
# Create a figure and initialize 3d plot
fig = plt.figure(figsize=figsize)
ax = init_axis(fig, title)
# Create spline line
trajectory = joints[:, 0, [x, y]]
avg_segment_length = (
np.mean(np.linalg.norm(np.diff(trajectory, axis=0), axis=1)) + 1e-3
)
draw_offset = int(25 / avg_segment_length)
(spline_line,) = ax.plot(*trajectory.T, zorder=10, color="white")
# Create a floor
minx, miny, _ = joints.min(axis=(0, 1))
maxx, maxy, _ = joints.max(axis=(0, 1))
plot_floor(ax, minx, maxx, miny, maxy, 0)
# Put the character on the floor
height_offset = np.min(joints[:, :, z]) # Min height
joints = joints.copy()
joints[:, :, z] -= height_offset
# Initialization for redrawing
lines = []
initialized = False
def update(frame):
nonlocal initialized
skeleton = joints[frame]
root = skeleton[0]
update_camera(ax, root)
hcolors = colors
if highlights is not None and highlights[frame]:
hcolors = ("red", "red", "red", "red", "red")
for index, (chain, color) in enumerate(
zip(reversed(kinematic_tree), reversed(hcolors))
):
if not initialized:
lines.append(
ax.plot(
skeleton[chain, x],
skeleton[chain, y],
skeleton[chain, z],
linewidth=6.0,
color=color,
zorder=20,
path_effects=[pe.SimpleLineShadow(), pe.Normal()],
)
)
else:
lines[index][0].set_xdata(skeleton[chain, x])
lines[index][0].set_ydata(skeleton[chain, y])
lines[index][0].set_3d_properties(skeleton[chain, z])
lines[index][0].set_color(color)
left = max(frame - draw_offset, 0)
right = min(frame + draw_offset, trajectory.shape[0])
spline_line.set_xdata(trajectory[left:right, 0])
spline_line.set_ydata(trajectory[left:right, 1])
spline_line.set_3d_properties(np.zeros_like(trajectory[left:right, 0]))
initialized = True
fig.tight_layout()
frames = joints.shape[0]
anim = FuncAnimation(fig, update, frames=frames, interval=1000 / fps, repeat=False)
if output == "notebook":
from IPython.display import HTML
HTML(anim.to_jshtml())
else:
# anim.save(output, writer='ffmpeg', fps=fps)
anim.save(output, fps=fps)
plt.close()