"""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