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 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):
    """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=True, 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