02_box_loading.py

"""Load all box assets under assets/box and arrange them in a gallery grid.

GLB assets in this folder use mixed units:
  - cube* / card* meshes are authored in centimeters (apply scale=0.01)
  - printers* meshes are already in meters (scale=1.0)
"""

from __future__ import annotations

import sys
from dataclasses import dataclass
from pathlib import Path

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

import genesis as gs
import gs_nyx.nyx_py_sdk as nps
import trimesh

from tools.rt_nyx import RtNyx, RtNyxConfig

BOX_DIR = ROOT / "assets" / "box"

# Genesis simulation uses meters; many exported GLBs still store vertex data in cm.
CM_TO_M = 0.01


@dataclass(frozen=True)
class BoxAsset:
    path: Path
    scale: float
    size_x: float
    size_y: float
    size_z: float
    z_min: float


def infer_scale(max_extent: float) -> float:
    """Return scale to convert mesh units to meters."""
    return CM_TO_M if max_extent > 1.0 else 1.0


def inspect_box(path: Path) -> BoxAsset:
    mesh = trimesh.load(path, force="mesh")
    if not isinstance(mesh, trimesh.Trimesh):
        mesh = mesh.dump(concatenate=True)

    extents = mesh.bounds[1] - mesh.bounds[0]
    scale = infer_scale(float(extents.max()))

    # GLB is Y-up; Genesis converts (X, Y, Z) -> (X, -Z, Y) when loading.
    size_x = float(extents[0]) * scale
    size_y = float(extents[2]) * scale
    size_z = float(extents[1]) * scale
    z_min = float(mesh.bounds[0][1]) * scale

    return BoxAsset(path=path, scale=scale, size_x=size_x, size_y=size_y, size_z=size_z, z_min=z_min)


def layout_boxes(
    boxes: list[BoxAsset],
    *,
    cols: int = 10,
    col_gap: float = 0.25,
    row_gap: float = 0.5,
    spawn_height: float = 0.5,
) -> list[tuple[BoxAsset, tuple[float, float, float]]]:
    placed: list[tuple[BoxAsset, tuple[float, float, float]]] = []
    x = 0.0
    y = 0.0
    row_depth = 0.0

    for i, box in enumerate(boxes):
        if i > 0 and i % cols == 0:
            y += row_depth + row_gap
            x = 0.0
            row_depth = 0.0

        # align=False keeps the GLB origin; z_min is the bottom in link frame.
        pos = (x + box.size_x / 2, y + box.size_y / 2, spawn_height - box.z_min)
        placed.append((box, pos))

        x += box.size_x + col_gap
        row_depth = max(row_depth, box.size_y)

    return placed


def discover_boxes() -> list[BoxAsset]:
    paths = sorted(BOX_DIR.glob("*.glb"))
    if not paths:
        raise FileNotFoundError(f"No GLB files found in {BOX_DIR}")
    return [inspect_box(path) for path in paths]


def main() -> None:
    gs.init(backend=gs.cuda, theme="light")  # type: ignore[attr-defined]
    assert gs.logger is not None

    boxes = discover_boxes()
    placed = layout_boxes(boxes)

    scene_center_x = max(pos[0] for _, pos in placed) / 2
    scene_center_y = max(pos[1] for _, pos in placed) / 2
    scene_span = max(scene_center_x, scene_center_y, 1.0)
    camera_pos = (scene_center_x, scene_center_y - scene_span * 1.4, scene_span * 0.9)
    camera_lookat = (scene_center_x, scene_center_y, 0.3)

    scene = gs.Scene(
        sim_options=gs.options.SimOptions(
            dt=0.01,
            gravity=(0, 0, -10.0),
        ),
        viewer_options=gs.options.ViewerOptions(
            res=(1280, 720),
            camera_pos=camera_pos,
            camera_lookat=camera_lookat,
            camera_fov=50,
            max_FPS=60,
        ),
        vis_options=gs.options.VisOptions(
            show_world_frame=False,
            show_link_frame=False,
            show_cameras=False,
            plane_reflection=True,
            ambient_light=(0.15, 0.15, 0.15),
        ),
        renderer=gs.renderers.Rasterizer(),
        show_viewer=True,
    )

    rt_nyx = RtNyx(
        scene,
        config=RtNyxConfig(
            res=(1920, 1080),
            preview_size=(960, 540),
            pos=camera_pos,
            lookat=camera_lookat,
            fov=50,
            spp=32,
            record_fps=24,
            output_dir="out",
        ),
    )

    scene.add_entity(gs.morphs.Plane())

    for box, pos in placed:
        rel_path = box.path.relative_to(ROOT).as_posix()
        gs.logger.info(
            f"Loading box ~<{rel_path}>~: scale=~<{box.scale:g}>~, "
            f"size=(~<{box.size_x:.2f}>~, ~<{box.size_y:.2f}>~, ~<{box.size_z:.2f}>~) m."
        )
        scene.add_entity(
            gs.morphs.Mesh(
                file=rel_path,
                scale=box.scale,
                pos=pos,
                align=False,
                fixed=False,
                collision=True,
            ),
        )

    env_map = nps.EnvironmentMapAsset()
    env_map.texture = "assets/kloppenheim_07_puresky_4k.hdr" # sky background texture
    env_map.layout = nps.EEnvMapLayout.LongLat
    env_map.multiplier = 8

    rt_nyx.add_camera(
        env_maps=(env_map,),
        lights=[
            {"type": "directional", "dir": (-0.3, 0.6, -1.0), "color": (1.0, 1.0, 1.0), "intensity": 5.0},
            {"type": "directional", "dir": (0.5, -0.2, -0.8), "color": (1.0, 1.0, 1.0), "intensity": 2.0},
        ],
    )

    scene.build(n_envs=1)
    rt_nyx.setup()

    try:
        while True:
            scene.step()
            if not rt_nyx.update():
                break
    finally:
        rt_nyx.close()


if __name__ == "__main__":
    main()