visualization/MEI138/render.py

"""MEI 138D rendering: SMPL-H mesh with skeleton overlay.

Pipeline:
    1. Recover joint positions (for skeleton + camera)
    2. Recover SMPL-H mesh vertices from MEI features
    3. Render mesh with pyrender (shared camera from joints)
    4. Overlay skeleton on mesh images
"""

import numpy as np

from .recovery import recover_joint_positions
from .representation import mei_to_smplh
from ..tools.render_skeleton import (
    get_smpl22_chains,
    render_skeleton_frames,
    compute_camera_params,
)
from ..tools.smplh import (
    load_smplh_model,
    get_J0,
    smplh_forward,
    LEFT_HAND_MEAN_AA,
    RIGHT_HAND_MEAN_AA,
)
from ..tools.render_mesh import render_mesh_frames


def _recover_mesh(motion):
    """Recover SMPL-H mesh vertices from MEI 138D features.

    Uses neutral gender and betas=0 (uniform body shape), consistent
    with the joint recovery pipeline.

    Args:
        motion: (T, 138) denormalized MEI features.

    Returns:
        verts: (T, 6890, 3) mesh vertices.
        faces: (F, 3) triangle indices.
    """
    T = motion.shape[0]

    # Decode MEI -> SMPL-H parameters
    out = mei_to_smplh(motion.astype(np.float64), fps=30.0)
    root_orient = out["root_orient"].astype(np.float32)  # (T, 3)
    body_pose = out["body_pose"].astype(np.float32)      # (T, 63)
    transl_abs = out["transl"].astype(np.float32)        # (T, 3) pelvis absolute

    # Assemble full 52-joint pose: root(3) + body(63) + left_hand(45) + right_hand(45)
    lh = np.tile(LEFT_HAND_MEAN_AA, (T, 1))   # (T, 45)
    rh = np.tile(RIGHT_HAND_MEAN_AA, (T, 1))  # (T, 45)
    poses_aa = np.concatenate([root_orient, body_pose, lh, rh], axis=-1)  # (T, 156)
    poses_aa = poses_aa.reshape(T, 52, 3)

    # Load neutral SMPL-H model
    gender = "neutral"
    betas = np.zeros(16, dtype=np.float32)
    model = load_smplh_model(gender)

    # Convert pelvis absolute position -> SMPLX convention (trans = pelvis - J0)
    J0 = get_J0(model, betas)
    transl = transl_abs - J0[None, :]

    # GPU forward pass
    verts = smplh_forward(model, gender, betas, poses_aa, transl)
    faces = model["f"]

    return verts, faces


def render_frames(motion: np.ndarray) -> list:
    """Render 138D MEI motion features to image frames.

    Renders SMPL-H mesh with skeleton overlay on top.

    Args:
        motion: (T, 138) denormalized MEI features.

    Returns:
        list of np.ndarray images (H, W, 3), uint8.
    """
    # 1. Recover joint positions (for skeleton + camera)
    joints = recover_joint_positions(motion)

    # 2. Recover SMPL-H mesh vertices
    verts, faces = _recover_mesh(motion)

    # 3. Compute camera from joints (shared between mesh and skeleton)
    cam_params = compute_camera_params(joints)

    # 4. Render mesh
    mesh_images = render_mesh_frames(verts, faces, cam_params)

    # 5. Overlay skeleton on mesh
    return render_skeleton_frames(joints, get_smpl22_chains(), canvas_images=mesh_images)