tests/test_pose_mirror.py

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

"""Mesh-level validation of PoseMirror_SOMA and PoseMirror_MHR.

For each mirror class, compares the parameter-mirrored mesh against a geometric
mesh mirror (flip X + mirror_vert_indices) in Nova topology, excluding facial
inner geometry.

Pose data paths are resolved exclusively from environment variables.
Tests skip gracefully when the variables are not set or the data is unavailable,
so CI and external users are unaffected.

Environment variables:
    SOMA_POSE_NPZ  Path to a single Nova pose .npz file (key: pose_local or transforms).
    SOMA_POSE_DIR  Directory of Nova pose .npz files (first file is used).
    MHR_POSE_NPZ   Path to a single MHR pose .npz file (key: pose_params).
    MHR_POSE_DIR    Directory of MHR pose .npz files (first file is used).

Usage:
    SOMA_POSE_NPZ="path/to/soma_pose.npz" MHR_POSE_NPZ="path/to/mhr_pose.npz" \
        pytest tests/test_pose_mirror.py -v
"""

from __future__ import annotations

import os
from pathlib import Path

import numpy as np
import pytest
import torch
from scipy.sparse import csc_matrix

REPO_ROOT = Path(__file__).resolve().parents[1]
ASSETS_DIR = REPO_ROOT / "assets"
CORE_ASSET = ASSETS_DIR / "SOMA_neutral.npz"

MAX_FRAMES = 100
BATCH_SIZE = 64


# ---------------------------------------------------------------------------
# Data-path resolution helpers
# ---------------------------------------------------------------------------


def _resolve_npz_path(env_single: str, env_dir: str) -> Path | None:
    """Return the first .npz found via the given environment variables.

    Checks *env_single* first (should point to a single file), then
    *env_dir* (should point to a directory; the first .npz is used).
    Returns ``None`` when neither variable is set or no file is found.
    """
    if val := os.environ.get(env_single):
        p = Path(val)
        return p if p.is_file() else None

    if val := os.environ.get(env_dir):
        search_dir = Path(val)
        if search_dir.is_dir():
            npzs = sorted(search_dir.glob("*.npz"))
            if npzs:
                return npzs[0]
    return None


def _load_mhr_param_names(npz_path: Path):
    """Extract the first 204 parameter names from parameter_transform.npz."""
    with np.load(npz_path, allow_pickle=False) as data:
        return list(data["parameter_names"][:204])


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture(scope="module")
def device():
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


@pytest.fixture(scope="module")
def soma_rig(device):
    """Load Nova rig data needed for mesh-level mirror comparison."""
    if not CORE_ASSET.is_file():
        pytest.skip(
            f"Core asset not found: {CORE_ASSET}. Run `git lfs pull` to fetch LFS-tracked files."
        )

    rig_data = np.load(CORE_ASSET, allow_pickle=False)

    if "mirror_vert_indices" not in rig_data:
        pytest.skip(
            "SOMA_neutral.npz does not contain 'mirror_vert_indices'. "
            "Either update the NPZ or point the test at an NPZ that includes it."
        )

    sw_sp = csc_matrix(
        (
            rig_data["skinning_weights_data"],
            rig_data["skinning_weights_indices"],
            rig_data["skinning_weights_indptr"],
        ),
        shape=rig_data["skinning_weights_shape"],
    ).todense()

    facial_inner = np.concatenate(
        [
            rig_data["segment_eye_bags"],
            rig_data["segment_mouth_bag"],
        ]
    )
    body_mask = np.ones(len(rig_data["bind_shape"]), dtype=bool)
    body_mask[facial_inner] = False

    return dict(
        joint_names=rig_data["joint_names"].tolist(),
        joint_parent_ids=torch.from_numpy(rig_data["joint_parent_ids"].astype(np.int32).copy()).to(
            device
        ),
        bind_pose_world=torch.from_numpy(rig_data["bind_pose_world"]).to(device),
        bind_shape=torch.from_numpy(rig_data["bind_shape"]).to(device),
        skinning_weights=torch.from_numpy(np.array(sw_sp)).to(device),
        mirror_vert_indices=rig_data["mirror_vert_indices"],
        body_mask=body_mask,
    )


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


@torch.no_grad()
def test_soma_pose_mirror(device, soma_rig):
    """PoseMirror_SOMA: param-mirrored mesh should match geometric mesh mirror."""
    from soma.geometry.batched_skinning import BatchedSkinning
    from soma.geometry.rig_utils import (
        PoseMirror_SOMA,
        joint_local_to_world,
        joint_world_to_local,
    )

    npz_path = _resolve_npz_path("SOMA_POSE_NPZ", "SOMA_POSE_DIR")
    if npz_path is None:
        pytest.skip("Nova pose data not available (set SOMA_POSE_NPZ or SOMA_POSE_DIR)")

    soma_npz = np.load(npz_path)
    key = "pose_local" if "pose_local" in soma_npz else "transforms"
    pose_local_np = soma_npz[key].astype(np.float32)
    n_frames = min(pose_local_np.shape[0], MAX_FRAMES)
    pose_local = torch.from_numpy(pose_local_np[:n_frames]).to(device)

    joint_parent_ids = soma_rig["joint_parent_ids"]
    pose_world = joint_local_to_world(pose_local, joint_parent_ids)

    mirror = PoseMirror_SOMA(soma_rig["joint_names"])
    pose_mirror_world = mirror(pose_world)
    pose_mirror_local = joint_world_to_local(pose_mirror_world, joint_parent_ids)

    skinning = BatchedSkinning(
        joint_parent_ids,
        soma_rig["skinning_weights"],
        soma_rig["bind_pose_world"],
        soma_rig["bind_shape"],
        mode="warp",
    )

    verts_orig_list, verts_pm_list = [], []
    for s in range(0, n_frames, BATCH_SIZE):
        e = min(s + BATCH_SIZE, n_frames)
        verts_orig_list.append(
            skinning.pose(pose_local[s:e, :, :3, :3], pose_local[s:e, 1, :3, 3]).cpu()
        )
        verts_pm_list.append(
            skinning.pose(pose_mirror_local[s:e, :, :3, :3], pose_mirror_local[s:e, 1, :3, 3]).cpu()
        )

    verts_orig = torch.cat(verts_orig_list)
    verts_param_mirror = torch.cat(verts_pm_list)

    mvi = soma_rig["mirror_vert_indices"]
    verts_mesh_mirror = verts_orig.clone()
    verts_mesh_mirror[..., 0] *= -1
    verts_mesh_mirror = verts_mesh_mirror[:, mvi]

    body_mask = soma_rig["body_mask"]
    err = (verts_param_mirror - verts_mesh_mirror).norm(dim=-1)[:, body_mask].numpy()
    mean_err = float(err.mean())
    p99_err = float(np.percentile(err, 99))

    print(f"\n  [SOMA mirror] frames={n_frames}  mean={mean_err:.6f}  p99={p99_err:.6f}")

    assert mean_err < 0.01, f"SOMA mirror mean error {mean_err:.6f} exceeds 0.01 threshold"
    assert p99_err < 0.02, f"SOMA mirror P99 error {p99_err:.6f} exceeds 0.02 threshold"


@torch.no_grad()
def test_mhr_pose_mirror(device, soma_rig):
    """PoseMirror_MHR: param-mirrored mesh should match geometric mesh mirror."""
    import trimesh

    from soma.geometry.barycentric_interp import BarycentricInterpolator
    from soma.geometry.rig_utils import PoseMirror_MHR

    npz_path = _resolve_npz_path("MHR_POSE_NPZ", "MHR_POSE_DIR")
    if npz_path is None:
        pytest.skip("MHR pose data not available (set MHR_POSE_NPZ or MHR_POSE_DIR)")

    pt_path = ASSETS_DIR / "MHR" / "parameter_transform.npz"
    model_path = ASSETS_DIR / "MHR" / "mhr_model_lod1.pt"
    mesh_mhr_path = ASSETS_DIR / "MHR" / "base_body_lod1.obj"
    mesh_soma_path = ASSETS_DIR / "MHR" / "SOMA_wrap_lod1.obj"
    for p in (pt_path, model_path, mesh_mhr_path, mesh_soma_path):
        if not p.is_file():
            pytest.skip(f"Required MHR asset not found: {p}")

    param_names = _load_mhr_param_names(pt_path)

    mhr_npz = np.load(npz_path)
    pp = torch.from_numpy(mhr_npz["pose_params"]).float().to(device)
    n_total = pp.shape[0]
    if pp.shape[1] < 204:
        pp = torch.cat([pp, torch.zeros(n_total, 204 - pp.shape[1], device=device)], 1)
    else:
        pp = pp[:, :204]
    n_frames = min(n_total, MAX_FRAMES)
    pp = pp[:n_frames]

    mirror = PoseMirror_MHR(param_names)
    pp_mirrored = mirror(pp)

    mhr_model = torch.jit.load(str(model_path), map_location=device)

    mesh_mhr = trimesh.load(str(mesh_mhr_path), maintain_order=True, process=False)
    mesh_soma = trimesh.load(str(mesh_soma_path), maintain_order=True, process=False)
    V_mhr = torch.from_numpy(mesh_mhr.vertices).float().to(device)
    F_mhr = torch.from_numpy(mesh_mhr.faces).to(device)
    V_soma = torch.from_numpy(mesh_soma.vertices).float().to(device)
    mhr_to_soma = BarycentricInterpolator(V_mhr, F_mhr, V_soma)

    id_c = torch.zeros(1, 45, device=device)
    fe = torch.zeros(1, 72, device=device)

    def mhr_forward(params_204):
        results = []
        for s in range(0, params_204.shape[0], BATCH_SIZE):
            e = min(s + BATCH_SIZE, params_204.shape[0])
            v, _ = mhr_model(
                id_c.expand(e - s, -1),
                params_204[s:e],
                fe.expand(e - s, -1),
            )
            results.append(mhr_to_soma(v).cpu())
            del v
        return torch.cat(results)

    verts_orig = mhr_forward(pp)
    verts_param_mirror = mhr_forward(pp_mirrored)

    mvi = soma_rig["mirror_vert_indices"]
    verts_mesh_mirror = verts_orig.clone()
    verts_mesh_mirror[..., 0] *= -1
    verts_mesh_mirror = verts_mesh_mirror[:, mvi]

    body_mask = soma_rig["body_mask"]
    err = (verts_param_mirror - verts_mesh_mirror).norm(dim=-1)[:, body_mask].numpy()
    mean_err = float(err.mean())
    p99_err = float(np.percentile(err, 99))

    print(f"\n  [MHR mirror] frames={n_frames}  mean={mean_err:.6f}  p99={p99_err:.6f}")

    assert mean_err < 0.005, f"MHR mirror mean error {mean_err:.6f} exceeds 0.005 threshold"
    assert p99_err < 0.01, f"MHR mirror P99 error {p99_err:.6f} exceeds 0.01 threshold"