OpenMotionLab/MotionGPT

1# Inspired by2# - https://github.com/anindita127/Complextext2animation/blob/main/src/utils/visualization.py3# - https://github.com/facebookresearch/QuaterNet/blob/main/common/visualization.py4 5from typing import List, Tuple6import numpy as np7from mGPT.utils.joints import mmm_kinematic_tree, mmm_to_smplh_scaling_factor8 9mmm_colors = ['black', 'magenta', 'red', 'green', 'blue']10 11 12def init_axis(fig, title, radius=1.5, dist=10):13    ax = fig.add_subplot(1, 1, 1, projection='3d')14    ax.view_init(elev=20., azim=-60)15 16    fact = 217    ax.set_xlim3d([-radius / fact, radius / fact])18    ax.set_ylim3d([-radius / fact, radius / fact])19    ax.set_zlim3d([0, radius])20 21    ax.set_aspect('auto')22    ax.set_xticklabels([])23    ax.set_yticklabels([])24    ax.set_zticklabels([])25 26    ax.set_axis_off()27 28    ax.dist = dist29    ax.grid(b=False)30 31    ax.set_title(title, loc='center', wrap=True)32    return ax33 34 35def plot_floor(ax, minx, maxx, miny, maxy, minz):36    from mpl_toolkits.mplot3d.art3d import Poly3DCollection37    # Plot a plane XZ38    verts = [39        [minx, miny, minz],40        [minx, maxy, minz],41        [maxx, maxy, minz],42        [maxx, miny, minz]43    ]44    xz_plane = Poly3DCollection([verts], zorder=1)45    xz_plane.set_facecolor((0.5, 0.5, 0.5, 1))46    ax.add_collection3d(xz_plane)47 48    # Plot a bigger square plane XZ49    radius = max((maxx - minx), (maxy - miny))50 51    # center +- radius52    minx_all = (maxx + minx) / 2 - radius53    maxx_all = (maxx + minx) / 2 + radius54 55    miny_all = (maxy + miny) / 2 - radius56    maxy_all = (maxy + miny) / 2 + radius57 58    verts = [59        [minx_all, miny_all, minz],60        [minx_all, maxy_all, minz],61        [maxx_all, maxy_all, minz],62        [maxx_all, miny_all, minz]63    ]64    xz_plane = Poly3DCollection([verts], zorder=1)65    xz_plane.set_facecolor((0.5, 0.5, 0.5, 0.5))66    ax.add_collection3d(xz_plane)67    return ax68 69 70def update_camera(ax, root, radius=1.5):71    fact = 272    ax.set_xlim3d([-radius / fact + root[0], radius / fact + root[0]])73    ax.set_ylim3d([-radius / fact + root[1], radius / fact + root[1]])74 75 76def render_animation(joints: np.ndarray, output: str = "notebook", title: str = "",77                     fps: float = 12.5,78                     kinematic_tree: List[List[int]] = mmm_kinematic_tree,79                     colors: List[str] = mmm_colors,80                     figsize: Tuple[int] = (4, 4),81                     fontsize: int = 15):82    import matplotlib.pyplot as plt83    from matplotlib.animation import FuncAnimation84    import matplotlib.patheffects as pe85    plt.rcParams.update({'font.size': fontsize})86 87    # Z is gravity here88    x, y, z = 0, 1, 289 90    # Convert mmm joints for visualization91    # into smpl-h "scale" and axis92    joints = joints.copy()[..., [2, 0, 1]] * mmm_to_smplh_scaling_factor93 94    # Create a figure and initialize 3d plot95    fig = plt.figure(figsize=figsize)96    ax = init_axis(fig, title)97 98    # Create spline line99    trajectory = joints[:, 0, [x, y]]100    avg_segment_length = np.mean(np.linalg.norm(np.diff(trajectory, axis=0), axis=1)) + 1e-3101    draw_offset = int(25 / avg_segment_length)102    spline_line, = ax.plot(*trajectory.T, zorder=10, color="white")103 104    # Create a floor105    minx, miny, _ = joints.min(axis=(0, 1))106    maxx, maxy, _ = joints.max(axis=(0, 1))107    plot_floor(ax, minx, maxx, miny, maxy, 0)108 109    # Put the character on the floor110    height_offset = np.min(joints[:, :, z])  # Min height111    joints = joints.copy()112    joints[:, :, z] -= height_offset113 114    # Initialization for redrawing115    lines = []116    initialized = False117 118    def update(frame):119        nonlocal initialized120        skeleton = joints[frame]121 122        root = skeleton[0]123        update_camera(ax, root)124 125        for index, (chain, color) in enumerate(zip(reversed(kinematic_tree), reversed(colors))):126            if not initialized:127                lines.append(ax.plot(skeleton[chain, x],128                                     skeleton[chain, y],129                                     skeleton[chain, z], linewidth=8.0, color=color, zorder=20,130                                     path_effects=[pe.SimpleLineShadow(), pe.Normal()]))131 132            else:133                lines[index][0].set_xdata(skeleton[chain, x])134                lines[index][0].set_ydata(skeleton[chain, y])135                lines[index][0].set_3d_properties(skeleton[chain, z])136 137        left = max(frame - draw_offset, 0)138        right = min(frame + draw_offset, trajectory.shape[0])139 140        spline_line.set_xdata(trajectory[left:right, 0])141        spline_line.set_ydata(trajectory[left:right, 1])142        spline_line.set_3d_properties(np.zeros_like(trajectory[left:right, 0]))143        initialized = True144 145    fig.tight_layout()146    frames = joints.shape[0]147    anim = FuncAnimation(fig, update, frames=frames, interval=1000 / fps, repeat=False)148 149    if output == "notebook":150        from IPython.display import HTML151        HTML(anim.to_jshtml())152    else:153        anim.save(output, writer='ffmpeg', fps=fps)154 155    plt.close()156