# 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()