pixal3d_mrp_mlx/pixal3d_mlx/mesh.py

"""Flexible Dual Grid -> triangle mesh -> GLB, via o_voxel.

The shape decoder emits **7 channels per occupied voxel**, and they are not a
signed-distance field — O-Voxel's Flexible Dual Grid solves a QEF instead, which is
what lets it carry open and non-manifold surfaces that marching cubes cannot:

    0:3  vertex offset inside the voxel, `(1+2m)*sigmoid(v) - m` so it may sit
         slightly OUTSIDE its own cell (m = voxel_margin = 0.5)
    3:6  per-axis intersection flags — logits at inference, thresholded at 0
    6:7  quad split weight, through softplus

`o_voxel.convert.flexible_dual_grid_to_mesh` turns those into vertices and faces, and
`o_voxel.postprocess.to_glb` does UV unwrap plus texture baking. Both are native
(C++/Metal) and are NOT ported: o-voxel builds a CPU CppExtension when CUDA is absent,
and the trellis-2 lane on this fleet already runs it with a Metal baker. Reusing that
build is strictly better than reimplementing a QEF solver in MLX.

o_voxel speaks torch, so this module is the MLX->torch boundary for the export path.
"""

from __future__ import annotations

from typing import Tuple

import mlx.core as mx
import numpy as np

# Upstream fixes both: the model always works in a unit cube centred on the origin.
AABB = [[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]]

# How the texture decoder's 6 channels map to PBR slots (pipeline's pbr_attr_layout).
# o_voxel indexes this dict by name and raises KeyError on any missing slot, so a
# partial layout fails deep inside the baker rather than at the call.
PBR_ATTR_LAYOUT = {
    "base_color": slice(0, 3),
    "metallic": slice(3, 4),
    "roughness": slice(4, 5),
    "alpha": slice(5, 6),
}


def _torch(a):
    import torch
    return torch.from_numpy(np.asarray(a))


def output_resolution(h, upsample_factor: int = 16) -> int:
    """The decoder's OUTPUT grid size, which is what o_voxel needs.

    The shape decoder applies four 2x upsamples, so a resolution-64 latent decodes into
    a 1024^3 grid. The `resolution` field in the checkpoint config is the decoder's
    configured default (256) — upstream overrides it per run via `set_resolution`, so
    reading it off the config gives the wrong grid and o_voxel's hashmap then raises an
    opaque out-of-bounds deep inside `insert`.
    """
    return int(mx.max(h.coords[:, 1:]).item()) // upsample_factor * upsample_factor + upsample_factor


def fdg_to_mesh(h, resolution: int, voxel_margin: float = 0.5) -> Tuple:
    """Shape-decoder output -> (vertices, faces) as torch tensors.

    `h` is the decoder's SparseTensor: `h.feats` [N,7], `h.coords` [N,4] with the batch
    index in column 0. `resolution` is the OUTPUT grid size (see `output_resolution`),
    not the decoder's configured one. Single batch item only, which is all inference
    ever produces.
    """
    from o_voxel.convert import flexible_dual_grid_to_mesh

    hi = int(mx.max(h.coords[:, 1:]).item())
    if hi >= resolution:
        raise ValueError(
            f"coords reach {hi} but grid_size={resolution}; pass the decoder's OUTPUT "
            f"resolution (input_res * 16), not its configured default"
        )

    feats = h.feats
    m = voxel_margin
    vertices = (1 + 2 * m) * mx.sigmoid(feats[..., 0:3]) - m
    intersected = feats[..., 3:6] > 0          # logits -> bool at inference
    quad_lerp = mx.logaddexp(feats[..., 6:7], mx.zeros_like(feats[..., 6:7]))  # softplus

    v, f = flexible_dual_grid_to_mesh(
        _torch(h.coords[:, 1:]).int(),
        _torch(vertices).float(),
        _torch(intersected).bool(),
        _torch(quad_lerp).float(),
        aabb=AABB,
        grid_size=resolution,
        train=False,
    )
    return v, f


def bake_vertex_colors(mesh, tex_voxels, resolution: int,
                       attr_layout: dict | None = None):
    """Sample the PBR attribute volume at each vertex -> COLOR_0. Seconds, not minutes.

    The UV path (`to_glb`) runs o_voxel's unwrap+bake, which on this CPU/Metal build
    burned >20 minutes of CPU on a 214k-face mesh even with remesh disabled — xatlas
    scales badly and 214k is the decimation floor, so it cannot simply be fed less.
    The trellis-2 lane reached the same conclusion and ships `--baker vertex` as its
    fast path for exactly this reason.

    Vertex colours lose the metallic/roughness maps — base colour only — but they are
    correct, immediate, and enough to see the asset. Positions map to voxel indices by
    the same linear aabb relation `fdg_to_mesh` used, so no resampling is involved.
    """
    import trimesh

    layout = attr_layout or PBR_ATTR_LAYOUT
    coords = np.asarray(tex_voxels.coords[:, 1:])
    attrs = np.asarray(tex_voxels.feats)

    lo, hi = np.array(AABB[0]), np.array(AABB[1])
    v = np.asarray(mesh.vertices)
    idx = np.floor((v - lo) / (hi - lo) * resolution).astype(np.int64)
    idx = np.clip(idx, 0, resolution - 1)

    # hash voxel coords -> row, then look each vertex up; unmatched vertices keep grey
    key = (coords[:, 0].astype(np.int64) * resolution + coords[:, 1]) * resolution + coords[:, 2]
    order = np.argsort(key)
    key_sorted = key[order]
    q = (idx[:, 0] * resolution + idx[:, 1]) * resolution + idx[:, 2]
    pos = np.searchsorted(key_sorted, q)
    pos = np.clip(pos, 0, len(key_sorted) - 1)
    hit = key_sorted[pos] == q

    base = layout["base_color"]
    rgb = np.full((len(v), 3), 0.5, np.float32)
    rgb[hit] = attrs[order[pos[hit]], base]
    colors = np.concatenate([np.clip(rgb, 0, 1), np.ones((len(v), 1), np.float32)], 1)

    out = trimesh.Trimesh(mesh.vertices, mesh.faces, process=False)
    out.visual = trimesh.visual.ColorVisuals(out, vertex_colors=(colors * 255).astype(np.uint8))
    return out, {"vertices_coloured": int(hit.sum()), "vertices_total": len(v),
                 "hit_rate": round(float(hit.mean()), 4)}


def to_glb(vertices, faces, tex_voxels, attr_layout: dict, resolution: int,
           texture_size: int = 4096, decimation_target: int = 1_000_000,
           prefer_metal: bool = True, remesh: bool = False):
    """Bake the texture voxels onto the mesh and return a trimesh GLB scene.

    `tex_voxels` is the texture decoder's SparseTensor (attrs in `.feats`, positions in
    `.coords`). `attr_layout` maps PBR channel names to slices of that feature vector.
    """
    try:
        if not prefer_metal:
            raise ImportError
        from o_voxel import postprocess as pp
    except ImportError:
        from o_voxel import postprocess_cpu as pp

    return pp.to_glb(
        vertices=vertices,
        faces=faces,
        attr_volume=_torch(tex_voxels.feats).float(),
        coords=_torch(tex_voxels.coords[:, 1:]).int(),
        attr_layout=attr_layout,
        grid_size=resolution,
        aabb=AABB,
        decimation_target=decimation_target,
        texture_size=texture_size,
        # `remesh` defaults OFF here, unlike upstream. Upstream runs on CUDA; this
        # build is the CPU/Metal one and its remesher took >20 minutes on a 214k-face
        # mesh before being killed. We also hand it an already-welded, floater-free,
        # decimated mesh, so the remesh has much less to fix than it would upstream.
        remesh=remesh, remesh_band=1, remesh_project=0,
    )


# Upstream rotates the asset out of its internal frame on the way out (inference.py).
EXPORT_ROTATION = np.array([[-1, 0, 0, 0],
                            [0, 0, -1, 0],
                            [0, -1, 0, 0],
                            [0, 0, 0, 1]], dtype=np.float64)


def to_camera_frame(vertices):
    """Mesh vertices -> the frame `proj.project_points` expects.

    THE GOTCHA: o_voxel returns vertices in the VOXEL GRID's frame — a linear map from
    integer coords into the aabb. `ProjGrid` rotates its lattice by `_BLENDER_ROT`
    BEFORE projecting, so mesh vertices must be rotated the same way to be compared
    against the source image. Skipping this does not throw; it silently reprojects a
    rotated object, which reads as a plausible-looking blob with a halo. It cost a
    wrong diagnosis here: a correct 0.969 silhouette IoU measured as 0.640.
    """
    from .proj import _BLENDER_ROT
    return np.asarray(vertices) @ _BLENDER_ROT.T