tools/render_joint_frames.py - master - koda/FiVI

tools/render_joint_frames.py
123 строки4 KB
koda
Prepare production-ready grasp planner
18 июл 2026, 09:39
18 июл 2026, 09:39f540458
from __future__ import annotations

import argparse
from pathlib import Path
import sys

import mujoco
import numpy as np
from PIL import Image

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from src.viewer import _line, _sphere


def _joint_id(model: mujoco.MjModel, name: str) -> int:
    result = int(mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, name))
    if result < 0:
        raise ValueError(f"joint not found: {name}")
    return result


def _decorate(scene: mujoco.MjvScene, model: mujoco.MjModel, data: mujoco.MjData) -> None:
    colors = ((1.0, 0.1, 0.1, 1.0), (0.1, 1.0, 0.1, 1.0), (0.1, 0.4, 1.0, 1.0))
    for joint_id in range(int(model.njnt)):
        name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, joint_id)
        if name is None or not name.startswith("rh_"):
            continue
        origin = np.asarray(data.xanchor[joint_id], dtype=np.float64)
        body_id = int(model.jnt_bodyid[joint_id])
        rotation = np.asarray(data.xmat[body_id], dtype=np.float64).reshape(3, 3)
        _sphere(scene, origin, 0.0012, (1.0, 0.9, 0.1, 1.0))
        for axis_index, color in enumerate(colors):
            _line(scene, origin, origin + 0.012 * rotation[:, axis_index], 0.00025, color)
        axis = np.asarray(data.xaxis[joint_id], dtype=np.float64)
        _line(scene, origin - 0.011 * axis, origin + 0.011 * axis, 0.0004, (1.0, 0.1, 1.0, 1.0))


def _render(
    model: mujoco.MjModel,
    data: mujoco.MjData,
    output: Path,
    lookat: np.ndarray,
    *,
    collision: bool,
) -> None:
    renderer = mujoco.Renderer(model, height=480, width=640)
    try:
        camera = mujoco.MjvCamera()
        camera.type = mujoco.mjtCamera.mjCAMERA_FREE
        camera.lookat[:] = lookat
        camera.distance = 0.23
        camera.azimuth = 145
        camera.elevation = -12
        option = mujoco.MjvOption()
        option.geomgroup[3] = int(collision)
        renderer.update_scene(data, camera=camera, scene_option=option)
        _decorate(renderer.scene, model, data)
        Image.fromarray(renderer.render()).save(output)
    finally:
        renderer.close()


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--xml", type=Path, default=ROOT / "models" / "hand" / "hand.xml")
    parser.add_argument(
        "--output",
        type=Path,
        default=ROOT / "logs" / "hand_frame_analysis" / "articulation",
    )
    arguments = parser.parse_args()
    xml_path = arguments.xml.expanduser().resolve()
    output = arguments.output.expanduser().resolve()
    if not xml_path.is_file():
        parser.error(f"model does not exist: {xml_path}")
    output.mkdir(parents=True, exist_ok=True)
    model = mujoco.MjModel.from_xml_path(str(xml_path))
    joints = {
        "ff": ("rh_FFJ4", "rh_FFJ3", "rh_FFJ2", "rh_FFJ1"),
        "mf": ("rh_MFJ4", "rh_MFJ3", "rh_MFJ2", "rh_MFJ1"),
        "rf": ("rh_RFJ4", "rh_RFJ3", "rh_RFJ2", "rh_RFJ1"),
        "lf": ("rh_LFJ4", "rh_LFJ3", "rh_LFJ2", "rh_LFJ1"),
        "th": ("rh_THJ5", "rh_THJ4", "rh_THJ2", "rh_THJ1"),
    }
    for prefix, (abduction, flexion, pip, dip) in joints.items():
        poses = {
            "neutral": {},
            "pip_30": {pip: np.deg2rad(30.0)},
            "pip_60": {pip: np.deg2rad(60.0)},
            "dip_20": {dip: np.deg2rad(20.0)},
            "mcp_flexion": {flexion: np.deg2rad(45.0)},
            "mcp_abduction": {abduction: np.deg2rad(20.0)},
        }
        for name, values in poses.items():
            data = mujoco.MjData(model)
            for joint_name, value in values.items():
                joint_id = _joint_id(model, joint_name)
                data.qpos[int(model.jnt_qposadr[joint_id])] = value
            mujoco.mj_forward(model, data)
            pivot = np.asarray(data.xanchor[_joint_id(model, pip)], dtype=np.float64)
            lookat = pivot + np.array((0.0, 0.0, 0.025))
            _render(
                model,
                data,
                output / f"{prefix}_{name}.png",
                lookat,
                collision=False,
            )
            if name == "neutral":
                _render(
                    model,
                    data,
                    output / f"{prefix}_neutral_collision.png",
                    lookat,
                    collision=True,
                )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())