mGPT/data/humanml/common/visualize.py

import numpy as np
import matplotlib.pyplot as plt
import matplotlib
from matplotlib.animation import FuncAnimation, FFMpegWriter
from mpl_toolkits.mplot3d import Axes3D  # Required for 3D plot
import os
matplotlib.use('Agg')

xsens_t2m_kinematic_chain = [                         
    [0, 1, 15, 16, 17, 18],      # left leg
    [0, 1, 19, 20, 21, 22],      # right leg
    [0, 2, 3, 4, 5, 6],    # spine -> head
    [4, 7, 8, 9, 10],         # left arm
    [4, 11, 12, 13, 14]       # right arm
]
humanml3d_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]]

R = np.array([
    [0, 0, 1],
    [1, 0, 0],
    [0, 1, 0]
], dtype=np.float32)

## VISUALIZE ##
def visualize(out_path, joints, prefix_idx=0, length_segment=100):
    t2m_kinematic_chain = xsens_t2m_kinematic_chain if joints.shape[1] == 23 else humanml3d_t2m_kinematic_chain
    def update(frame_idx):
        ax.cla()  # clear current frame

        frame = joints[frame_idx].copy()  # shape: (23, 3)
        xs, ys, zs = frame[:, 0], frame[:, 1], frame[:, 2]

        ax.scatter(xs, ys, zs, c='b', s=2)
        color = 'lightgray' if frame_idx < prefix_idx else 'black'
        for chain in t2m_kinematic_chain:
            for i in range(len(chain) - 1):
                j1, j2 = chain[i], chain[i + 1]
                ax.plot(
                    [frame[j1, 0], frame[j2, 0]],
                    [frame[j1, 1], frame[j2, 1]],
                    [frame[j1, 2], frame[j2, 2]],
                    color=color, linewidth=1
                )
        # for i, (x, y, z) in enumerate(frame):
        #     ax.text(x, y, z, str(i), color='red', fontsize=5)

        ax.set_title(f'Frame {frame_idx}', fontsize=12)
        ax.set_xlim(xmin, xmax)
        ax.set_ylim(ymin, ymax)
        ax.set_zlim(zmin, zmax)
        ax.set_box_aspect((dx, dy, dz))
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')
        
    joints = joints @ R.T
    
    num_seg = joints.shape[0] // length_segment
    if num_seg == 0 :
        num_seg += 1
    for i in range(num_seg) :
        if i >= 1: continue
        
        start_idx = i * length_segment
        end_idx = min((i+1)*length_segment, joints.shape[0])
        joints_segment = joints[start_idx:end_idx]

        fig = plt.figure(figsize=(6, 6), dpi=150)
        ax = fig.add_subplot(111, projection='3d')

        all_joints = joints_segment.reshape(-1, 3)
        xmin, xmax = all_joints[:, 0].min()-0.3, all_joints[:, 0].max()+0.3
        ymin, ymax = all_joints[:, 1].min()-0.3, all_joints[:, 1].max()+0.3
        zmin, zmax = all_joints[:, 2].min()-0.3, all_joints[:, 2].max()+0.3 
        dx = max(xmax - xmin, 1e-6)
        dy = max(ymax - ymin, 1e-6)
        dz = max(zmax - zmin, 1e-6)
        
        # anim = FuncAnimation(fig, update_only_joints, frames=len(joints), interval=500)
        anim = FuncAnimation(fig, update, frames=[i for i in range(start_idx, end_idx)], interval=500)
        # anim.save(os.path.join(out_dir, f'{base_name}_ori_{cnt}_{i}.gif'), writer="pillow", fps=10)
        anim.save(out_path, writer="pillow", fps=10)
        plt.close(fig)
        
def visualize_3pts(out_path, joints, triads_T34=None, prefix_idx=0, length_segment=100):
    t2m_kinematic_chain = xsens_t2m_kinematic_chain if joints.shape[1] == 23 else humanml3d_t2m_kinematic_chain
    
    def undo_rotate_pose_matrix(T34_rotated, R):
        out = T34_rotated.copy()
        out[..., :3, :3] = (R[None, ...] @ out[..., :3, :3])   # inverse of (R.T @ ·) is (R @ ·)
        out[..., :3,  3] = out[..., :3,  3] @ R.T              # inverse of (· @ R) is (· @ R.T)
        return out

    def draw_triad(ax, R, t, axis_len=0.05, lw=1.5, alpha=0.95):
        """R: (3,3), t: (3,)"""
        o = t.reshape(3)
        x, y, z = R[:, 0], R[:, 1], R[:, 2]
        # X(red), Y(green), Z(blue)
        ax.quiver(*o, *x, length=axis_len, linewidth=lw, color='r', alpha=alpha)
        ax.quiver(*o, *y, length=axis_len, linewidth=lw, color='g', alpha=alpha)
        ax.quiver(*o, *z, length=axis_len, linewidth=lw, color='b', alpha=alpha)
    
    def update(frame_idx):
        ax.cla()  # clear current frame

        frame = joints[frame_idx].copy()  # shape: (23, 3)
        xs, ys, zs = frame[:, 0], frame[:, 1], frame[:, 2]

        ax.scatter(xs, ys, zs, c='b', s=2)
        color = 'lightgray' if frame_idx < prefix_idx else 'black'
        for chain in t2m_kinematic_chain:
            for i in range(len(chain) - 1):
                j1, j2 = chain[i], chain[i + 1]
                ax.plot(
                    [frame[j1, 0], frame[j2, 0]],
                    [frame[j1, 1], frame[j2, 1]],
                    [frame[j1, 2], frame[j2, 2]],
                    color=color, linewidth=1
                )
                
        if triads_T34 is not None and frame_idx < triads_T34.shape[0]:
            for j in range(triads_T34.shape[1]):
                M = triads_T34[frame_idx, j]
                Rloc = M[:3, :3] #orthonormalize(M[:3, :3])
                tloc = M[:3, 3]
                axis_len = max(dx, dy, dz) * 0.04
                draw_triad(ax, Rloc, tloc, axis_len=axis_len, lw=1.4, alpha=0.95)
                
                
        ax.set_title(f'Frame {frame_idx}', fontsize=12)
        ax.set_xlim(xmin, xmax)
        ax.set_ylim(ymin, ymax)
        ax.set_zlim(zmin, zmax)
        ax.set_box_aspect((dx, dy, dz))
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')
    
    joints = joints @ R.T
    if triads_T34 is not None :
        triads_T34=undo_rotate_pose_matrix(triads_T34, R)
    
    num_seg = joints.shape[0] // length_segment
    if num_seg == 0 :
        num_seg += 1
    for i in range(num_seg) :
        if i >= 1: continue
        
        start_idx = i * length_segment
        end_idx = min((i+1)*length_segment, joints.shape[0])
        joints_segment = joints[start_idx:end_idx]

        fig = plt.figure(figsize=(6, 6), dpi=150)
        ax = fig.add_subplot(111, projection='3d')

        all_joints = joints_segment.reshape(-1, 3)
        xmin, xmax = all_joints[:, 0].min()-0.3, all_joints[:, 0].max()+0.3
        ymin, ymax = all_joints[:, 1].min()-0.3, all_joints[:, 1].max()+0.3
        zmin, zmax = all_joints[:, 2].min()-0.3, all_joints[:, 2].max()+0.3 
        dx = max(xmax - xmin, 1e-6)
        dy = max(ymax - ymin, 1e-6)
        dz = max(zmax - zmin, 1e-6)
        
        # anim = FuncAnimation(fig, update_only_joints, frames=len(joints), interval=500)
        anim = FuncAnimation(fig, update, frames=[i for i in range(start_idx, end_idx)], interval=500)
        # anim.save(os.path.join(out_dir, f'{base_name}_ori_{cnt}_{i}.gif'), writer="pillow", fps=10)
        anim.save(out_path, writer="pillow", fps=10)
        plt.close(fig)
        

def visualize_two(out_path, joints_ref, joints_rst, prefix_idx=0, length_segment=50):
    """
    joints1, joints2: (T, J, 3) numpy arrays (e.g. J=23 or 22/24, etc.)
    joints1 in black, joints2 in light gray; if the two sequences differ in length, animate to the shorter one.
    """
    joints1 = joints_ref
    joints2 = joints_rst
    assert joints1.ndim == 3 and joints1.shape[-1] == 3, "joints1 shape must be (T, J, 3)"
    assert joints2.ndim == 3 and joints2.shape[-1] == 3, "joints2 shape must be (T, J, 3)"
    T = min(joints1.shape[0], joints2.shape[0])
    J1, J2 = joints1.shape[1], joints2.shape[1]

    if J1 == 23 and J2 == 23:
        t2m_kinematic_chain = xsens_t2m_kinematic_chain
    else:
        NotImplementedError()

    joints1 = joints1 @ R.T
    joints2 = joints2 @ R.T

    num_seg = T // length_segment
    if num_seg == 0:
        num_seg = 1

    for seg_idx in range(num_seg):
        if seg_idx >= 1: continue

        start_idx = seg_idx * length_segment
        end_idx = min((seg_idx + 1) * length_segment, T)

        seg1 = joints1[start_idx:end_idx]  # (t, J1, 3)
        seg2 = joints2[start_idx:end_idx]  # (t, J2, 3)

        fig = plt.figure(figsize=(6, 6), dpi=150)
        ax = fig.add_subplot(111, projection='3d')

        all_joints = np.concatenate([seg1.reshape(-1, 3), seg2.reshape(-1, 3)], axis=0)
        xmin, xmax = all_joints[:, 0].min() - 0.3, all_joints[:, 0].max() + 0.3
        ymin, ymax = all_joints[:, 1].min() - 0.3, all_joints[:, 1].max() + 0.3
        zmin, zmax = all_joints[:, 2].min() - 0.3, all_joints[:, 2].max() + 0.3
        dx = max(xmax - xmin, 1e-6)
        dy = max(ymax - ymin, 1e-6)
        dz = max(zmax - zmin, 1e-6)

        def plot_skeleton(frame, color_lines, size_points=2, lw=1):
            xs, ys, zs = frame[:, 0], frame[:, 1], frame[:, 2]
            ax.scatter(xs, ys, zs, c='b', s=size_points)
            for chain in t2m_kinematic_chain:
                for i in range(len(chain) - 1):
                    j1, j2 = chain[i], chain[i + 1]
                    if j1 < frame.shape[0] and j2 < frame.shape[0]:
                        ax.plot(
                            [frame[j1, 0], frame[j2, 0]],
                            [frame[j1, 1], frame[j2, 1]],
                            [frame[j1, 2], frame[j2, 2]],
                            color=color_lines, linewidth=lw
                        )

        def update(frame_idx):
            ax.cla()
            f1 = joints1[frame_idx]
            f2 = joints2[frame_idx]

            plot_skeleton(f1, color_lines='black', lw=1)
            plot_skeleton(f2, color_lines='lightgray', lw=1)

            ax.set_title(f'Frame {frame_idx}', fontsize=12)
            ax.set_xlim(xmin, xmax)
            ax.set_ylim(ymin, ymax)
            ax.set_zlim(zmin, zmax)
            ax.set_box_aspect((dx, dy, dz))
            ax.set_xlabel('X')
            ax.set_ylabel('Y')
            ax.set_zlabel('Z')


        frames = list(range(start_idx, end_idx))
        anim = FuncAnimation(fig, update, frames=frames, interval=500)
        save_path = out_path if num_seg == 1 else out_path.replace(".gif", f"_seg{seg_idx}.gif")
        anim.save(save_path, writer="pillow", fps=10)
        plt.close(fig)