visualize.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Visualize motions generated by scripts/generate.py (.npz) in the browser.

Accepts a single .npz file or a folder of them (as written with --num_samples > 1);
the skeleton (core / g1 / soma) is detected from the joint count.

Examples:
    python scripts/visualize.py outputs/output.npz
    python scripts/visualize.py outputs/waves --port 2334
"""

import argparse
import time
from pathlib import Path

import numpy as np
import torch
import viser

from ardy.skeleton import (
    CoreSkeleton27,
    G1Skeleton34,
    SOMASkeleton30,
    SOMASkeleton77,
)
from ardy.viz.viser_utils import Character

SKELETON_BY_NBJOINTS = {
    27: (CoreSkeleton27, "core_skin"),
    30: (SOMASkeleton30, "soma_skin"),
    34: (G1Skeleton34, "g1_stl"),
    77: (SOMASkeleton77, "soma_skin"),
}


def parse_args():
    parser = argparse.ArgumentParser(description="Visualize generated motion .npz files in viser")
    parser.add_argument(
        "path",
        type=str,
        help="A .npz file from scripts/generate.py, or a folder of them (--num_samples > 1 output).",
    )
    parser.add_argument("--port", type=int, default=2334, help="Viser server port (default: 2334)")
    parser.add_argument("--fps", type=float, default=None, help="Playback FPS override (default: the file's fps)")
    return parser.parse_args()


def list_motion_files(path: str) -> list[Path]:
    p = Path(path)
    if p.is_dir():
        files = sorted(p.glob("*.npz"))
        if not files:
            raise FileNotFoundError(f"No .npz files found in folder {p}")
        return files
    if not p.exists():
        raise FileNotFoundError(f"{p} does not exist")
    return [p]


def load_motion(path: Path) -> dict:
    """Load one generated sample: joint positions/rotations (FK fallback), contacts, fps, prompt."""
    data = np.load(path)

    def get(key):
        return torch.from_numpy(np.asarray(data[key])).float() if key in data else None

    joints_pos = get("posed_joints")  # [T, J, 3]
    joints_rot = get("global_rot_mats")  # [T, J, 3, 3]

    if joints_pos is None or joints_rot is None:
        local_rot_mats = get("local_rot_mats")
        root_positions = get("root_positions")
        if local_rot_mats is None or root_positions is None:
            raise ValueError(f"{path} has neither posed joints nor local rotations to reconstruct them")
        skeleton_cls, _ = SKELETON_BY_NBJOINTS[local_rot_mats.shape[1]]
        fk_rot, fk_pos, _ = skeleton_cls().fk(local_rot_mats, root_positions)
        joints_pos = joints_pos if joints_pos is not None else fk_pos
        joints_rot = joints_rot if joints_rot is not None else fk_rot

    foot_contacts = get("foot_contacts")
    return {
        "joints_pos": joints_pos,
        "joints_rot": joints_rot,
        "foot_contacts": foot_contacts,
        "num_frames": joints_pos.shape[0],
        "nbjoints": joints_pos.shape[1],
        "fps": float(data["fps"]) if "fps" in data else 20.0,
        "text": str(data["text"]) if "text" in data else "",
    }


def create_character(server: viser.ViserServer, nbjoints: int) -> Character:
    if nbjoints not in SKELETON_BY_NBJOINTS:
        raise ValueError(f"Unsupported joint count {nbjoints}; expected one of {sorted(SKELETON_BY_NBJOINTS)}")
    skeleton_cls, mesh_mode = SKELETON_BY_NBJOINTS[nbjoints]
    skeleton = skeleton_cls()
    try:
        return Character(
            "motion",
            server,
            skeleton,
            create_skeleton_mesh=True,
            create_skinned_mesh=True,
            visible_skeleton=False,
            visible_skinned_mesh=True,
            mesh_mode=mesh_mode,
        )
    except Exception as e:
        # Skin assets unavailable — fall back to the skeleton-only view.
        print(f"Could not create skinned mesh ({e}); showing skeleton only")
        return Character(
            "motion",
            server,
            skeleton,
            create_skeleton_mesh=True,
            create_skinned_mesh=False,
            visible_skeleton=True,
        )


def main():
    args = parse_args()
    files = list_motion_files(args.path)

    server = viser.ViserServer(host="0.0.0.0", port=args.port, label="ARDY Motion Viewer")
    server.scene.set_up_direction("+y")
    server.scene.add_grid(
        "/grid",
        width=20.0,
        height=20.0,
        wxyz=viser.transforms.SO3.from_x_radians(-np.pi / 2.0).wxyz,
        position=(0.0, 0.0001, 0.0),
        fade_distance=20.0,
        infinite_grid=True,
    )

    @server.on_client_connect
    def _(client: viser.ClientHandle) -> None:
        client.camera.position = (3.0, 2.0, 3.0)
        client.camera.look_at = (0.0, 1.0, 0.0)

    state = {"motion": load_motion(files[0]), "character": None}
    state["character"] = create_character(server, state["motion"]["nbjoints"])

    gui_info = server.gui.add_markdown("")
    if len(files) > 1:
        gui_sample = server.gui.add_dropdown("Sample", options=[f.name for f in files], initial_value=files[0].name)
    gui_playing = server.gui.add_checkbox("Playing", initial_value=True)
    gui_frame = server.gui.add_slider("Frame", min=0, max=state["motion"]["num_frames"] - 1, step=1, initial_value=0)
    gui_show_mesh = server.gui.add_checkbox("Show Mesh", initial_value=True)
    gui_show_skeleton = server.gui.add_checkbox("Show Skeleton", initial_value=state["character"].skinned_mesh is None)

    def show_frame(frame_idx: int) -> None:
        motion = state["motion"]
        frame_idx = min(frame_idx, motion["num_frames"] - 1)
        contacts = motion["foot_contacts"][frame_idx] if motion["foot_contacts"] is not None else None
        state["character"].set_pose(motion["joints_pos"][frame_idx], motion["joints_rot"][frame_idx], contacts)

    def show_file(path: Path) -> None:
        motion = load_motion(path)
        if motion["nbjoints"] != state["motion"]["nbjoints"]:
            state["character"].clear()
            state["character"] = create_character(server, motion["nbjoints"])
        state["motion"] = motion
        gui_frame.max = motion["num_frames"] - 1
        gui_frame.value = 0
        gui_info.content = (
            f"**Prompt:** {motion['text']}  \n"
            f"**File:** {path.name} — {motion['num_frames']} frames @ {motion['fps']:g} fps"
        )
        show_frame(0)

    @gui_frame.on_update
    def _(_event) -> None:
        show_frame(int(gui_frame.value))

    @gui_show_mesh.on_update
    def _(_event) -> None:
        state["character"].set_skinned_mesh_visibility(gui_show_mesh.value)

    @gui_show_skeleton.on_update
    def _(_event) -> None:
        if state["character"].skeleton_mesh is not None:
            state["character"].skeleton_mesh.set_visibility(gui_show_skeleton.value)

    if len(files) > 1:

        @gui_sample.on_update
        def _(_event) -> None:
            show_file(next(f for f in files if f.name == gui_sample.value))

    show_file(files[0])
    print(f"Viewer running at http://localhost:{args.port} — Ctrl+C to stop")

    # Schedule frames against an absolute clock so per-frame work doesn't slow
    # or unevenly pace the playback.
    next_tick = time.perf_counter()
    while True:
        fps = args.fps if args.fps is not None else state["motion"]["fps"]
        if gui_playing.value and state["motion"]["num_frames"] > 1:
            gui_frame.value = (int(gui_frame.value) + 1) % state["motion"]["num_frames"]
        next_tick += 1.0 / fps
        delay = next_tick - time.perf_counter()
        if delay > 0:
            time.sleep(delay)
        else:
            next_tick = time.perf_counter()  # fell behind; don't try to catch up


if __name__ == "__main__":
    main()