diff --git a/pixal3d_mlx/mesh.py b/pixal3d_mlx/mesh.py new file mode 100644 index 0000000..e331b6e --- /dev/null +++ b/pixal3d_mlx/mesh.py @@ -0,0 +1,96 @@ +"""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]] + + +def _torch(a): + import torch + return torch.from_numpy(np.asarray(a)) + + +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. Single batch item only, which is all inference ever produces. + """ + from o_voxel.convert import flexible_dual_grid_to_mesh + + 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)