Mesh export: Flexible Dual Grid -> triangles -> GLB, via o_voxel
The shape decoder's 7 channels are not an SDF — O-Voxel solves a QEF over a Flexible Dual Grid, which is what lets it carry open and non-manifold surfaces: 0:3 vertex offset in-voxel, (1+2m)*sigmoid(v)-m so it may sit OUTSIDE its own cell 3:6 per-axis intersection logits, thresholded at 0 6:7 quad split weight through softplus mesh.py is the MLX->torch boundary for export. o_voxel's convert/postprocess are native (C++/Metal) and deliberately 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, so reusing that build beats reimplementing a QEF solver in MLX. Installed into the shared venv from ~/Documents/trellis-2-mrp-mlx/o-voxel; it needs cv2 and xatlas, and NOT utils3d (which drags in open3d, with no cp312 wheel). Verified against the REAL shape_dec (292/292 params, resolution 256): decoded 5954 voxels x 7ch -> 5954 vertices, 6886 faces -> GLB written Not watertight, correctly: the input was a random latent, and FlexiDualGrid represents open surfaces by design. Vertices land inside the octant of the unit cube matching the sparse coords fed in, which is the check that the grid indexing is right. Still to wire: the SLAT stage itself (sparse latents seeded from the occupancy coords, with proj features gathered at those coords). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
parent
6f40fae0cc
commit
5484d59cb3
96
pixal3d_mlx/mesh.py
Normal file
96
pixal3d_mlx/mesh.py
Normal file
@ -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)
|
||||
Loading…
Reference in New Issue
Block a user