viewer4d_export.py

"""Export a Lift4D custom-run reconstruction into the interactive 4D viewer's
param format (the same ``meta.json`` + ``*.bin`` bundle under ``static/params/``
that ``static/js/gs4d_player.js`` consumes for the released example scenes).

The released example scenes are canonical deformable-Gaussian reconstructions
(a fixed canonical splat set skinned by sparse control nodes across frames).  A
custom run instead produces *independent* per-frame Gaussian splats (one posed
splat cloud per reconstructed frame, with a possibly different Gaussian count
each frame).  The viewer's format is structurally "canonical splats + per-frame
deformation", so we express the per-frame clouds through the viewer's existing
``base`` node-skinning path *without* changing the player:

  * pad/truncate every frame to a common count ``N`` (pad slots get opacity 0),
  * give every canonical Gaussian its own control node (``M == N``, ``Kb == 1``,
    ``base_idx[i] == i``, ``base_w[i] == 1``),
  * store each frame's actual per-Gaussian translation / rotation / scale in the
    per-frame base-node arrays, so ``_skin(f)`` reproduces exactly that frame's
    splat cloud (canonical arrays are neutral placeholders).

The player's skinning for the ``has_base``-only, ``opt_deform_rot`` case is
(``gs4d_player.js`` ``_skin``):

    d_xyz = Σ bW·bNt ; d_rot = Σ bW·bNr ; d_scale = Σ bW·bNs
    q = normalize(rot_raw + d_rot) ; s = max(scale + d_scale, eps)
    p_world = (pos + d_xyz) @ glob_M + glob_g
    cov_world = (glob_L·R(q)) diag(s²) (glob_L·R(q))^T

so with canonical pos=0, rot_raw=0, scale=0 and glob_M/glob_L = I, glob_g = 0
the world Gaussian at frame f is exactly (bNt[f], normalize(bNr[f]), bNs[f]).
"""
from __future__ import annotations

import json
from pathlib import Path

import numpy as np

# viewer space (y up, z toward viewer) matches the per-frame gr.Model3D splat
# export in orbit_render.write_viewer_ply: camera space (y down, z forward) is
# rotated by diag(1,-1,-1) and quats premultiplied by Rx(180deg) (wxyz 0,1,0,0).
_VIEW_FLIP = np.array([[1.0, 0.0, 0.0], [0.0, -1.0, 0.0], [0.0, 0.0, -1.0]], dtype=np.float32)
_VIEW_QUAT = np.array([0.0, 1.0, 0.0, 0.0], dtype=np.float32)  # wxyz, Rx(180 deg)


def _quat_mul_np(a, b):
    aw, ax, ay, az = a
    bw, bx, by, bz = b[:, 0], b[:, 1], b[:, 2], b[:, 3]
    return np.stack([
        aw * bw - ax * bx - ay * by - az * bz,
        aw * bx + ax * bw + ay * bz - az * by,
        aw * by - ax * bz + ay * bw + az * bx,
        aw * bz + ax * by - ay * bx + az * bw,
    ], axis=-1).astype(np.float32)


def _frame_arrays(packed, center):
    """(means, quats(wxyz), scales, colors[0..1], opac) -> viewer-space numpy.

    ``packed`` is one entry of the reconstruction's ``cam_frames`` list, i.e.
    (means, quats, scales, colors, opacities) camera-space tensors.  Recentre on
    the shared ``center`` and rotate into the viewer's y-up frame, matching the
    gr.Model3D per-frame splat export so both viewers agree.
    """
    means, quats, scales, colors, opac = packed
    m = means.detach().cpu().numpy().astype(np.float32)
    q = quats.detach().cpu().numpy().astype(np.float32)
    s = scales.detach().cpu().numpy().astype(np.float32)
    c = colors.detach().cpu().numpy().astype(np.float32)
    o = opac.detach().cpu().numpy().astype(np.float32).reshape(-1)

    m = (m - center) @ _VIEW_FLIP.T
    q = _quat_mul_np(_VIEW_QUAT, q)
    return m, q, s, c, o


def export_run(out_dir, cam_frames, fps: int = 12) -> str:
    """Write a viewer4d param bundle for a custom run into ``out_dir``.

    Args:
        out_dir: directory to write ``meta.json`` + ``*.bin`` into.
        cam_frames: list of per-frame (means, quats, scales, colors, opac)
            camera-space tensors, as built in app.reconstruct.
        fps: playback fps hint stored in meta.

    Returns:
        The bundle directory path (str).
    """
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)

    F = len(cam_frames)
    if F == 0:
        raise ValueError("no frames to export")

    # Shared recentring so the sequence stays framed the same way the orbit /
    # per-frame splat viewers recentre (each on its own median). Use a common
    # center (median of frame-0 means) so the object doesn't jump frame-to-frame.
    center0 = np.median(cam_frames[0][0].detach().cpu().numpy().astype(np.float32), axis=0)

    frames = [_frame_arrays(p, center0) for p in cam_frames]
    counts = [fr[0].shape[0] for fr in frames]
    N = int(max(counts))

    # Per-frame node arrays (M == N, one node per canonical Gaussian).
    bNt = np.zeros((F, N, 3), dtype=np.float32)   # base_node_trans  -> world position
    bNr = np.zeros((F, N, 4), dtype=np.float32)   # base_node_rot    -> rotation residual (== rotation, canon=0)
    bNs = np.zeros((F, N, 3), dtype=np.float32)   # base_node_scale  -> scale residual (== scale, canon=0)
    # Time-invariant per-Gaussian color/opacity: take the frame with the most
    # Gaussians as the color source; padded slots stay opacity 0 (invisible).
    src = int(np.argmax(counts))
    col = np.zeros((N, 3), dtype=np.uint8)
    opac = np.zeros((N,), dtype=np.float32)

    for f, (m, q, s, c, o) in enumerate(frames):
        n = m.shape[0]
        bNt[f, :n] = m
        bNr[f, :n] = q
        bNs[f, :n] = s
        # padded slots: identity rotation + tiny scale so normalize()/cov stay finite
        if n < N:
            bNr[f, n:, 0] = 1.0
            bNs[f, n:] = 1e-6

    m_s, q_s, s_s, c_s, o_s = frames[src]
    ns = m_s.shape[0]
    col[:ns] = np.clip(c_s * 255.0, 0, 255).astype(np.uint8)
    opac[:ns] = np.clip(o_s, 0.0, 1.0)
    # padded slots invisible (opacity already 0)

    # Neutral canonical arrays (all deformation flows through base nodes).
    positions = np.zeros((N, 3), dtype=np.float32)
    rot_raw = np.zeros((N, 4), dtype=np.float32)   # canon rot 0; q = normalize(0 + bNr)
    scale = np.zeros((N, 3), dtype=np.float32)     # canon scale 0; s = bNs

    # Per-frame identity globals (glob_M / glob_L = I, glob_g = 0).
    eyeF = np.tile(np.eye(3, dtype=np.float32).reshape(1, 9), (F, 1))
    glob_M = eyeF.copy()
    glob_L = eyeF.copy()
    glob_g = np.zeros((F, 3), dtype=np.float32)

    # KNN base skinning: node i, weight 1.
    base_idx = np.arange(N, dtype=np.uint32).reshape(N, 1)
    base_w = np.ones((N, 1), dtype=np.float32)

    def _w(name, arr):
        (out / name).write_bytes(np.ascontiguousarray(arr).tobytes())

    _w("positions.bin", positions)
    _w("rot_raw.bin", rot_raw)
    _w("scale.bin", scale)
    _w("color.bin", col)
    _w("opacity.bin", opac)
    _w("base_idx.bin", base_idx)
    _w("base_w.bin", base_w)
    _w("base_node_trans.bin", bNt)
    _w("base_node_rot.bin", bNr)
    _w("base_node_scale.bin", bNs)
    _w("glob_M.bin", glob_M)
    _w("glob_g.bin", glob_g)
    _w("glob_L.bin", glob_L)

    meta = {
        "scene": "custom",
        "N": N,
        "M": N,
        "F": F,
        "Kd": 0,
        "Kb": 1,
        "opt_deform_rot": True,
        "has_delta": False,
        "has_base": True,
        "has_delta_rot": False,
        "fps": int(fps),
        "cam_up": [0.0, 1.0, 0.0],
    }
    (out / "meta.json").write_text(json.dumps(meta))
    return str(out)