visualize.py

"""
MuJoCo 离屏渲染：DOF + root pose → 带 mesh 的机器人动画视频 (MP4)。
"""
import os

# 设置 MuJoCo 离屏渲染后端 (必须在 import mujoco 之前)
# EGL 用于有 GPU 的环境, osmesa 用于纯 CPU 环境 (如 HF Spaces)
if 'MUJOCO_GL' not in os.environ:
    try:
        import ctypes
        ctypes.cdll.LoadLibrary('libEGL.so.1')
        os.environ['MUJOCO_GL'] = 'egl'
    except OSError:
        os.environ['MUJOCO_GL'] = 'osmesa'

import tempfile
import numpy as np
import mujoco as mj
import imageio

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
XML_PATH = os.path.join(BASE_DIR, 'assets', 'g1_mocap_29dof.xml')

# inference 输出 DOF 顺序 → MuJoCo XML joint 顺序的映射
JOINT_MAPPING = [
    0, 6, 12,
    1, 7, 13,
    2, 8, 14,
    3, 9, 15, 22,
    4, 10, 16, 23,
    5, 11, 17, 24,
    18, 25,
    19, 26,
    20, 27,
    21, 28,
]


def create_skeleton_animation(dof, rot_quat, transl, fps=30,
                              width=1920, height=1088,
                              camera_distance=2.5, camera_elevation=-15):
    """渲染 G1 机器人 mesh 动画并输出 MP4 视频。

    Args:
        dof: (T, 29) 关节角度
        rot_quat: (T, 4) wxyz 四元数
        transl: (T, 3) 根位置
        fps: 输出视频帧率
        width, height: 渲染分辨率
        camera_distance: 摄像机距离
        camera_elevation: 摄像机俯仰角 (度)

    Returns:
        str: MP4 视频文件路径
    """
    model = mj.MjModel.from_xml_path(XML_PATH)
    # 确保离屏帧缓冲区足够大
    model.vis.global_.offwidth = max(model.vis.global_.offwidth, width)
    model.vis.global_.offheight = max(model.vis.global_.offheight, height)
    data = mj.MjData(model)
    renderer = mj.Renderer(model, height=height, width=width)

    # 摄像机设置
    camera = mj.MjvCamera()
    camera.distance = camera_distance
    camera.elevation = camera_elevation
    camera.azimuth = 90

    pelvis_id = model.body('pelvis').id

    T = dof.shape[0]

    # 映射 DOF 顺序
    mapped_dof = np.zeros((T, model.nq - 7), dtype=np.float32)
    for i, jm in enumerate(JOINT_MAPPING):
        mapped_dof[:, jm] = dof[:, i]

    # 渲染到视频
    video_path = os.path.join(tempfile.gettempdir(), 'g1_motion.mp4')
    writer = imageio.get_writer(video_path, fps=fps, codec='libx264')

    for t in range(T):
        data.qpos[:3] = transl[t]
        data.qpos[3:7] = rot_quat[t]  # wxyz, MuJoCo 原生格式
        data.qpos[7:] = mapped_dof[t]

        mj.mj_forward(model, data)

        # 摄像机跟踪 pelvis
        camera.lookat[:] = data.xpos[pelvis_id]

        renderer.update_scene(data, camera=camera)
        img = renderer.render()
        writer.append_data(img)

    writer.close()
    renderer.close()

    return video_path