sphereparam.py

"""Small Python bridge for Hoppe/Praun spherical parameterization.

The native optimization is performed by the compiled command-line tools in this
repository. This module keeps the Python side simple: load a triangular mesh,
run the native pipeline, and save/load `.npz` files with vertices, faces,
spherical coordinates, and UV coordinates.
"""

from __future__ import annotations

import os
import re
import subprocess
import sys
import tempfile
from pathlib import Path

import numpy as np
import trimesh


REPO_ROOT = Path(__file__).resolve().parent
BIN_DIR = REPO_ROOT / "bin" / "unix"
MESHTOPM = REPO_ROOT / "demos" / "bin" / "meshtopm.sh"
SPHPARAM = BIN_DIR / "SphereParam"
REVERSELINES = BIN_DIR / "reverselines"
FILTERPROG = BIN_DIR / "Filterprog"
MESHSIMPLIFY = BIN_DIR / "MeshSimplify"

_SPH_RE = re.compile(
    r"^Vertex\s+(\d+)\s+[-+0-9.eE]+\s+[-+0-9.eE]+\s+[-+0-9.eE]+.*"
    r"\{[^}]*\bsph=\(([-+0-9.eE]+)\s+([-+0-9.eE]+)\s+([-+0-9.eE]+)\)"
)


def required_tools() -> dict[str, Path]:
    return {
        "MeshSimplify": MESHSIMPLIFY,
        "reverselines": REVERSELINES,
        "Filterprog": FILTERPROG,
        "SphereParam": SPHPARAM,
        "meshtopm.sh": MESHTOPM,
    }


def check_tools() -> None:
    missing = [f"{name}: {path}" for name, path in required_tools().items() if not path.exists()]
    if missing:
        raise FileNotFoundError(
            "Native tools are missing:\n  "
            + "\n  ".join(missing)
            + "\nRun `./compile.sh` from the repository root."
        )


def load_hoppe_m(path: Path) -> tuple[np.ndarray, np.ndarray]:
    vertices: dict[int, list[float]] = {}
    faces: list[list[int]] = []
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            parts = line.strip().split()
            if not parts:
                continue
            if parts[0] == "Vertex":
                vertices[int(parts[1])] = [float(parts[2]), float(parts[3]), float(parts[4])]
            elif parts[0] == "Face":
                ids = []
                for token in parts[2:]:
                    if token.startswith("{"):
                        break
                    ids.append(int(token))
                if len(ids) == 3:
                    faces.append(ids)
    if not vertices or not faces:
        raise ValueError(f"{path} does not look like a triangular Hoppe .m mesh")
    ordered_ids = sorted(vertices)
    id_to_zero = {vid: i for i, vid in enumerate(ordered_ids)}
    verts = np.array([vertices[vid] for vid in ordered_ids], dtype=np.float64)
    tris = np.array([[id_to_zero[v] for v in face] for face in faces], dtype=np.int64)
    return verts, tris


def load_mesh(path: str | Path) -> tuple[np.ndarray, np.ndarray]:
    path = Path(path)
    if path.suffix.lower() == ".m":
        return load_hoppe_m(path)
    mesh = trimesh.load_mesh(path, process=False)
    if isinstance(mesh, trimesh.Scene):
        mesh = trimesh.util.concatenate(tuple(mesh.geometry.values()))
    if not isinstance(mesh, trimesh.Trimesh):
        raise ValueError(f"Could not load {path} as a triangle mesh")
    if mesh.faces.shape[1] != 3:
        mesh = mesh.triangulate()
    return np.asarray(mesh.vertices, dtype=np.float64), np.asarray(mesh.faces, dtype=np.int64)


def normalize_for_display(vertices: np.ndarray) -> np.ndarray:
    centered = vertices - vertices.mean(axis=0, keepdims=True)
    radius = np.linalg.norm(centered, axis=1).max()
    return centered / max(radius, 1e-12)


def sphere_to_uv(sphere_vertices: np.ndarray) -> np.ndarray:
    xyz = np.asarray(sphere_vertices, dtype=np.float64)
    xyz = xyz / np.maximum(np.linalg.norm(xyz, axis=1, keepdims=True), 1e-12)
    u = (np.arctan2(xyz[:, 1], xyz[:, 0]) + np.pi) / (2.0 * np.pi)
    v = np.arccos(np.clip(xyz[:, 2], -1.0, 1.0)) / np.pi
    return np.column_stack([u, v])


def uv_xy_colors(uv: np.ndarray) -> np.ndarray:
    colors = np.zeros((len(uv), 3), dtype=np.float64)
    colors[:, 0] = uv[:, 0]
    colors[:, 1] = uv[:, 1]
    colors[:, 2] = 1.0 - 0.5 * (uv[:, 0] + uv[:, 1])
    return np.clip(colors, 0.0, 1.0)


def _write_hoppe_m(path: Path, vertices: np.ndarray, faces: np.ndarray) -> None:
    with path.open("w", encoding="utf-8") as f:
        for i, p in enumerate(vertices, start=1):
            f.write(f"Vertex {i}  {p[0]:.9g} {p[1]:.9g} {p[2]:.9g}\n")
        for i, face in enumerate(faces, start=1):
            f.write(f"Face {i}  {int(face[0]) + 1} {int(face[1]) + 1} {int(face[2]) + 1}\n")


def _parse_sphereparam_output(text: str, n_vertices: int) -> np.ndarray:
    sphere = np.full((n_vertices, 3), np.nan, dtype=np.float64)
    for line in text.splitlines():
        match = _SPH_RE.match(line)
        if not match:
            continue
        idx = int(match.group(1)) - 1
        if 0 <= idx < n_vertices:
            sphere[idx] = [float(match.group(2)), float(match.group(3)), float(match.group(4))]
    if not np.isfinite(sphere).all():
        missing = np.where(~np.isfinite(sphere).all(axis=1))[0][:10] + 1
        raise RuntimeError(f"SphereParam output is missing sph coordinates for vertices {missing.tolist()}")
    return sphere / np.maximum(np.linalg.norm(sphere, axis=1, keepdims=True), 1e-12)


def _native_env(threads: int | None) -> dict[str, str]:
    env = os.environ.copy()
    env["PATH"] = f"{BIN_DIR}:{env.get('PATH', '')}"
    env_lib = Path(sys.prefix) / "lib"
    if env_lib.exists():
        env["LD_LIBRARY_PATH"] = f"{env_lib}:{env.get('LD_LIBRARY_PATH', '')}"
    if threads is not None:
        env["OMP_NUM_THREADS"] = str(int(threads))
    return env


def compute_spherical_parameterization(
    vertices: np.ndarray,
    faces: np.ndarray,
    *,
    effort: int = 2,
    base_faces: int | None = None,
    threads: int | None = None,
    verbose: bool = True,
) -> np.ndarray:
    """Return one unit-sphere point per input vertex.

    The input must be a closed triangular genus-0 mesh. Vertex order is preserved.
    """
    check_tools()
    vertices = np.asarray(vertices, dtype=np.float64)
    faces = np.asarray(faces, dtype=np.int64)
    if vertices.ndim != 2 or vertices.shape[1] != 3:
        raise ValueError("vertices must have shape (n, 3)")
    if faces.ndim != 2 or faces.shape[1] != 3:
        raise ValueError("faces must have shape (m, 3)")
    if faces.min(initial=0) < 0 or faces.max(initial=0) >= len(vertices):
        raise ValueError("faces contain vertex indices outside vertices")

    with tempfile.TemporaryDirectory(prefix="sphereparam_") as tmp:
        tmp = Path(tmp)
        mesh_path = tmp / "input.m"
        pm_path = tmp / "input.pm"
        original_indices_path = tmp / "original_indices.txt"
        _write_hoppe_m(mesh_path, vertices, faces)

        env = _native_env(threads)
        meshtopm_cmd = [
            str(MESHTOPM),
            str(mesh_path),
            "-minqem",
            "-vsgeom",
            "-dihallow",
            "-original_indices",
            str(original_indices_path),
        ]
        if base_faces is not None:
            meshtopm_cmd += ["-nf", str(int(base_faces)), "-simp"]
        with pm_path.open("wb") as pm_file:
            subprocess.run(
                meshtopm_cmd,
                stdout=pm_file,
                stderr=None if verbose else subprocess.DEVNULL,
                check=True,
                cwd=REPO_ROOT,
                env=env,
            )

        result = subprocess.run(
            [
                str(SPHPARAM),
                str(pm_path),
                "-original_mesh",
                str(mesh_path),
                "-original_indices",
                str(original_indices_path),
                "-to",
                "original_mesh",
                "-verbose",
                "1" if verbose else "0",
                "-effort",
                str(int(effort)),
            ],
            text=True,
            stdout=subprocess.PIPE,
            stderr=None if verbose else subprocess.DEVNULL,
            check=True,
            cwd=REPO_ROOT,
            env=env,
        )
    return _parse_sphereparam_output(result.stdout, len(vertices))


def parameterize_mesh_file(
    mesh_path: str | Path,
    out_path: str | Path,
    *,
    effort: int = 2,
    base_faces: int | None = None,
    threads: int | None = None,
    verbose: bool = True,
) -> Path:
    mesh_path = Path(mesh_path)
    out_path = Path(out_path)
    vertices, faces = load_mesh(mesh_path)
    sphere_vertices = compute_spherical_parameterization(
        vertices,
        faces,
        effort=effort,
        base_faces=base_faces,
        threads=threads,
        verbose=verbose,
    )
    uv = sphere_to_uv(sphere_vertices)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(
        out_path,
        vertices=vertices,
        faces=faces,
        sphere_vertices=sphere_vertices,
        uv=uv,
        source_mesh=str(mesh_path),
        effort=int(effort),
        base_faces=-1 if base_faces is None else int(base_faces),
    )
    return out_path


def load_parameterization_npz(path: str | Path) -> dict[str, np.ndarray]:
    data = np.load(path, allow_pickle=False)
    required = {"vertices", "faces", "sphere_vertices", "uv"}
    missing = required - set(data.files)
    if missing:
        raise ValueError(f"{path} is missing keys: {sorted(missing)}")
    return {key: data[key] for key in data.files}


def save_registered_mesh_parameterization(
    mesh_path: str | Path,
    out_path: str | Path,
    *,
    template_faces: np.ndarray,
    template_sphere_vertices: np.ndarray,
    template_uv: np.ndarray,
    template_path: str | Path,
    require_same_faces: bool = False,
) -> Path:
    """Save a registered mesh using an already-computed template parameterization.

    This is for datasets where every mesh is in vertexwise correspondence with a
    common template. The mesh geometry changes, but vertex ids and topology are
    assumed to follow the template, so the UV/sphere map is copied from it.
    """
    mesh_path = Path(mesh_path)
    out_path = Path(out_path)
    vertices, faces = load_mesh(mesh_path)
    if len(vertices) != len(template_sphere_vertices):
        raise ValueError(
            f"{mesh_path} has {len(vertices)} vertices, but template has "
            f"{len(template_sphere_vertices)} vertices"
        )
    if require_same_faces and not np.array_equal(faces, template_faces):
        raise ValueError(f"{mesh_path} faces do not exactly match the template faces")
    out_path.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(
        out_path,
        vertices=vertices,
        faces=template_faces,
        sphere_vertices=template_sphere_vertices,
        uv=template_uv,
        source_mesh=str(mesh_path),
        template_mesh=str(template_path),
        registered_to_template=True,
    )
    return out_path