346 lines
16 KiB
C++
346 lines
16 KiB
C++
// flexible_dual_grid.h — single-header CPU port of TRELLIS.2's flexible dual
|
||
// grid mesh extraction (o-voxel/o_voxel/convert/flexible_dual_grid.py, eval
|
||
// path). Turns the shape decoder's per-voxel 7-channel output into a triangle
|
||
// mesh, replacing the CUDA hashmap kernel with an std::unordered_map.
|
||
//
|
||
// Per active voxel v at integer coord c (in [0, grid_size)):
|
||
// dual vertex V_v = (c + offset_v) * voxel_size + aabb0 (unit cube here)
|
||
// offset_v = (1 + 2*margin) * sigmoid(feat[0:3]) - margin
|
||
// intersected feat[3:6] > 0, one flag per axis (x, y, z)
|
||
// split_weight softplus(feat[6])
|
||
// For each voxel with an intersected axis, the 4 voxels around that edge
|
||
// (offsets below) contribute their dual vertices as a quad; if all 4 exist the
|
||
// quad is split into 2 triangles along the diagonal chosen by the decoder's
|
||
// learned split_weight = softplus(feat[6]) (reference eval-path tie-break).
|
||
#pragma once
|
||
|
||
#include <algorithm>
|
||
#include <cmath>
|
||
#include <cstdint>
|
||
#include <unordered_map>
|
||
#include <unordered_set>
|
||
#include <vector>
|
||
|
||
namespace fdg {
|
||
|
||
struct Mesh {
|
||
std::vector<float> verts; // 3 per vertex
|
||
std::vector<int> tris; // 3 indices per triangle
|
||
size_t n_verts() const { return verts.size() / 3; }
|
||
size_t n_tris() const { return tris.size() / 3; }
|
||
};
|
||
|
||
namespace detail {
|
||
|
||
inline uint64_t key(int32_t x, int32_t y, int32_t z) {
|
||
return ((uint64_t) (uint32_t) x << 40) |
|
||
((uint64_t) (uint32_t) y << 20) |
|
||
(uint64_t) (uint32_t) z;
|
||
}
|
||
|
||
// The 4 neighbor-voxel offsets around an edge, per axis (matches
|
||
// edge_neighbor_voxel_offset in the reference).
|
||
static const int EDGE_OFF[3][4][3] = {
|
||
{{0,0,0},{0,0,1},{0,1,1},{0,1,0}}, // x-axis edge
|
||
{{0,0,0},{1,0,0},{1,0,1},{0,0,1}}, // y-axis edge
|
||
{{0,0,0},{0,1,0},{1,1,0},{1,0,0}}, // z-axis edge
|
||
};
|
||
|
||
inline void cross(const float * a, const float * b, float * o) {
|
||
o[0] = a[1]*b[2] - a[2]*b[1];
|
||
o[1] = a[2]*b[0] - a[0]*b[2];
|
||
o[2] = a[0]*b[1] - a[1]*b[0];
|
||
}
|
||
|
||
} // namespace detail
|
||
|
||
// feats: [n_voxels * 7], voxel-major (dec output). coords: [n_voxels * 3].
|
||
// grid_size = input_res * decoder upscale (e.g. 32 * 16 = 512). margin 0.5.
|
||
inline Mesh extract(const float * feats, const int32_t * coords, int n,
|
||
int grid_size, float margin = 0.5f) {
|
||
using namespace detail;
|
||
Mesh m;
|
||
if (n <= 0) return m;
|
||
|
||
const float vs = 1.0f / (float) grid_size; // voxel size (aabb span 1)
|
||
const float aabb0 = -0.5f;
|
||
|
||
// dual vertices + hashmap
|
||
std::vector<float> V((size_t) n * 3);
|
||
std::unordered_map<uint64_t, int> idx;
|
||
idx.reserve((size_t) n * 2);
|
||
for (int v = 0; v < n; ++v) {
|
||
const float * f = feats + (size_t) v * 7;
|
||
for (int a = 0; a < 3; ++a) {
|
||
const float s = 1.0f / (1.0f + std::exp(-f[a])); // sigmoid
|
||
const float off = (1.0f + 2.0f * margin) * s - margin;
|
||
V[(size_t) v * 3 + a] = ((float) coords[(size_t) v * 3 + a] + off) * vs + aabb0;
|
||
}
|
||
idx[key(coords[(size_t) v * 3], coords[(size_t) v * 3 + 1], coords[(size_t) v * 3 + 2])] = v;
|
||
}
|
||
m.verts = V;
|
||
|
||
// quads from intersected edges
|
||
for (int v = 0; v < n; ++v) {
|
||
const float * f = feats + (size_t) v * 7;
|
||
const int32_t cx = coords[(size_t) v * 3];
|
||
const int32_t cy = coords[(size_t) v * 3 + 1];
|
||
const int32_t cz = coords[(size_t) v * 3 + 2];
|
||
for (int axis = 0; axis < 3; ++axis) {
|
||
if (f[3 + axis] <= 0.0f) continue; // not intersected on this axis
|
||
int q[4];
|
||
bool ok = true;
|
||
for (int i = 0; i < 4; ++i) {
|
||
const int32_t nx = cx + EDGE_OFF[axis][i][0];
|
||
const int32_t ny = cy + EDGE_OFF[axis][i][1];
|
||
const int32_t nz = cz + EDGE_OFF[axis][i][2];
|
||
auto it = idx.find(key(nx, ny, nz));
|
||
if (it == idx.end()) { ok = false; break; }
|
||
q[i] = it->second;
|
||
}
|
||
if (!ok) continue;
|
||
|
||
// Choose the quad diagonal by the decoder's learned split_weight
|
||
// (softplus of feat[6]), exactly as the reference eval path does
|
||
// (FlexiDualGridVaeDecoder -> flexible_dual_grid_to_mesh, train=False):
|
||
// split 1: (0,1,2)+(0,2,3) when sw0*sw2 > sw1*sw3
|
||
// split 2: (0,1,3)+(3,1,2) otherwise
|
||
// (A geometric best-aligned-normals heuristic is the reference's
|
||
// split_weight=None fallback; the shipped decoder always emits
|
||
// feat[6], so we follow the learned choice.)
|
||
auto sw = [&](int i) {
|
||
const float x = feats[(size_t) q[i] * 7 + 6];
|
||
return x > 20.0f ? x : std::log1p(std::exp(x)); // softplus
|
||
};
|
||
if (sw(0) * sw(2) > sw(1) * sw(3)) {
|
||
m.tris.push_back(q[0]); m.tris.push_back(q[1]); m.tris.push_back(q[2]);
|
||
m.tris.push_back(q[0]); m.tris.push_back(q[2]); m.tris.push_back(q[3]);
|
||
} else {
|
||
m.tris.push_back(q[0]); m.tris.push_back(q[1]); m.tris.push_back(q[3]);
|
||
m.tris.push_back(q[3]); m.tris.push_back(q[1]); m.tris.push_back(q[2]);
|
||
}
|
||
}
|
||
}
|
||
return m;
|
||
}
|
||
|
||
// Per-vertex shading normals for the dual grid's *unoriented*, heavily
|
||
// non-manifold mesh.
|
||
//
|
||
// The reference mesher emits every quad with a fixed vertex order regardless of
|
||
// which way the surface crosses the edge, so a large fraction of faces are wound
|
||
// opposite to their neighbours. Two failure modes follow: (a) a naive
|
||
// area-weighted normal cancels at those seams, and (b) any *sign* fix that leaves
|
||
// stray flips is not harmless — the normal is interpolated across the triangle
|
||
// *before* the fragment shader's abs(dot), so two adjacent vertices with opposite
|
||
// normals make the interpolated normal cross zero mid-face → normalize() explodes
|
||
// → speckled/blocky specular. The mesh is too non-manifold to 2-colour the
|
||
// winding cleanly (edges shared by >2 faces frustrate it), so we don't try to.
|
||
//
|
||
// 1. Recover a smooth, winding-INDEPENDENT normal *direction* per vertex as the
|
||
// dominant eigenvector of the area-weighted structure tensor Σ area·n̂n̂ᵀ
|
||
// (immune to winding sign since n̂n̂ᵀ == (−n̂)(−n̂)ᵀ).
|
||
// 2. Resolve the arbitrary per-vertex *sign* consistently with a parity
|
||
// union-find over mesh edges, so edge-adjacent vertices share a hemisphere
|
||
// and the interpolated normal stays clear of zero. Only genuinely frustrated
|
||
// (odd-cycle / non-manifold) edges are left flipped.
|
||
// Final shading is orientation-independent (viewer uses abs(dot)), so only local
|
||
// smoothness matters, not a globally correct outward sign. (Winding unification
|
||
// and sign diffusion were both tried and are worse on this mesh: 2-colouring the
|
||
// >2-face non-manifold edges frustrates more, and Jacobi diffusion checkerboards.)
|
||
inline std::vector<float> vertex_normals(const Mesh & m) {
|
||
const size_t nv = m.n_verts();
|
||
std::vector<double> A((size_t) nv * 6, 0.0); // sym structure tensor per vertex
|
||
std::vector<float> seed((size_t) nv * 3, 0.0f); // signed area sum: a sign hint
|
||
for (size_t t = 0; t < m.tris.size(); t += 3) {
|
||
const int i0 = m.tris[t], i1 = m.tris[t + 1], i2 = m.tris[t + 2];
|
||
const float * a = &m.verts[(size_t) i0 * 3];
|
||
const float * b = &m.verts[(size_t) i1 * 3];
|
||
const float * c = &m.verts[(size_t) i2 * 3];
|
||
float e1[3], e2[3], fn[3];
|
||
for (int k = 0; k < 3; ++k) { e1[k] = b[k]-a[k]; e2[k] = c[k]-a[k]; }
|
||
detail::cross(e1, e2, fn);
|
||
const double area = std::sqrt((double) fn[0]*fn[0] + (double) fn[1]*fn[1] + (double) fn[2]*fn[2]);
|
||
if (area <= 1e-20) continue;
|
||
const double inv = 1.0 / area;
|
||
const double xx = fn[0]*fn[0]*inv, yy = fn[1]*fn[1]*inv, zz = fn[2]*fn[2]*inv;
|
||
const double xy = fn[0]*fn[1]*inv, xz = fn[0]*fn[2]*inv, yz = fn[1]*fn[2]*inv;
|
||
for (int i : {i0, i1, i2}) {
|
||
double * Av = &A[(size_t) i * 6];
|
||
Av[0]+=xx; Av[1]+=yy; Av[2]+=zz; Av[3]+=xy; Av[4]+=xz; Av[5]+=yz;
|
||
float * sv = &seed[(size_t) i * 3];
|
||
sv[0]+=fn[0]; sv[1]+=fn[1]; sv[2]+=fn[2];
|
||
}
|
||
}
|
||
// (1) smooth, sign-ambiguous direction per vertex
|
||
std::vector<float> dir((size_t) nv * 3, 0.0f);
|
||
for (size_t v = 0; v < nv; ++v) {
|
||
const double * Av = &A[v * 6];
|
||
double x = seed[v*3], y = seed[v*3+1], z = seed[v*3+2];
|
||
double l = std::sqrt(x*x + y*y + z*z);
|
||
if (l < 1e-20) { x = Av[0]; y = Av[3]; z = Av[4]; l = std::sqrt(x*x+y*y+z*z); }
|
||
if (l < 1e-20) { dir[v*3+2] = 1.0f; continue; }
|
||
x/=l; y/=l; z/=l;
|
||
for (int it = 0; it < 8; ++it) {
|
||
const double nx = Av[0]*x + Av[3]*y + Av[4]*z;
|
||
const double ny = Av[3]*x + Av[1]*y + Av[5]*z;
|
||
const double nz = Av[4]*x + Av[5]*y + Av[2]*z;
|
||
const double nl = std::sqrt(nx*nx + ny*ny + nz*nz);
|
||
if (nl < 1e-20) break;
|
||
x = nx/nl; y = ny/nl; z = nz/nl;
|
||
}
|
||
dir[v*3] = (float) x; dir[v*3+1] = (float) y; dir[v*3+2] = (float) z;
|
||
}
|
||
// (2) Resolve the arbitrary per-vertex sign *consistently* via a parity
|
||
// union-find over mesh edges: two edge-adjacent vertices whose directions
|
||
// are anti-aligned must end up with opposite signs (and vice versa), so
|
||
// within each connected component neighbours share a hemisphere. Only
|
||
// genuinely frustrated (odd-cycle / non-manifold) edges are left flipped.
|
||
std::vector<int> ufp(nv), ufr(nv, 0);
|
||
std::vector<uint8_t> ufb(nv, 0); // parity of a vertex relative to its parent
|
||
for (size_t i = 0; i < nv; ++i) ufp[i] = (int) i;
|
||
auto find = [&](int v, int & parity) {
|
||
int p = 0;
|
||
while (ufp[v] != v) { p ^= ufb[v]; v = ufp[v]; }
|
||
parity = p; return v;
|
||
};
|
||
auto join = [&](int a, int b) {
|
||
const float * na = &dir[(size_t) a * 3];
|
||
const float * nb = &dir[(size_t) b * 3];
|
||
const int rel = (na[0]*nb[0] + na[1]*nb[1] + na[2]*nb[2]) < 0.0f ? 1 : 0;
|
||
int pa, pb, ra = find(a, pa), rb = find(b, pb);
|
||
if (ra == rb) return;
|
||
if (ufr[ra] < ufr[rb]) { std::swap(ra, rb); std::swap(pa, pb); }
|
||
ufp[rb] = ra; ufb[rb] = (uint8_t) (pa ^ pb ^ rel);
|
||
if (ufr[ra] == ufr[rb]) ufr[ra]++;
|
||
};
|
||
for (size_t t = 0; t < m.tris.size(); t += 3) {
|
||
const int a = m.tris[t], b = m.tris[t+1], c = m.tris[t+2];
|
||
join(a, b); join(b, c); join(c, a);
|
||
}
|
||
// (3) pick each component's global sign toward the signed-sum seed (so the
|
||
// result is deterministic and roughly outward), then emit oriented normals
|
||
std::unordered_map<int, double> comp_sign;
|
||
for (size_t v = 0; v < nv; ++v) {
|
||
int p, r = find((int) v, p);
|
||
const double s = p ? -1.0 : 1.0;
|
||
const float * sv = &seed[v*3];
|
||
comp_sign[r] += s * (dir[v*3]*sv[0] + dir[v*3+1]*sv[1] + dir[v*3+2]*sv[2]);
|
||
}
|
||
std::vector<float> nrm((size_t) nv * 3, 0.0f);
|
||
for (size_t v = 0; v < nv; ++v) {
|
||
int p, r = find((int) v, p);
|
||
double s = p ? -1.0 : 1.0;
|
||
if (comp_sign[r] < 0.0) s = -s;
|
||
nrm[v*3] = (float) (s * dir[v*3]);
|
||
nrm[v*3+1] = (float) (s * dir[v*3+1]);
|
||
nrm[v*3+2] = (float) (s * dir[v*3+2]);
|
||
}
|
||
return nrm;
|
||
}
|
||
|
||
// Remove triangles that belong to tiny disconnected islands (the floating
|
||
// specks that read as "blemishes"), keeping the vertex array and its indexing
|
||
// intact so a parallel per-vertex attribute array (e.g. baked PBR) stays
|
||
// aligned. Components are face groups connected through shared vertices; a
|
||
// component is dropped when its face count is below min_frac of the total.
|
||
inline void drop_small_components(Mesh & m, float min_frac = 0.0005f) {
|
||
const size_t nt = m.n_tris();
|
||
if (nt == 0) return;
|
||
const size_t nv = m.n_verts();
|
||
std::vector<int> p(nv);
|
||
for (size_t i = 0; i < nv; ++i) p[i] = (int) i;
|
||
auto find = [&](int x) { while (p[x] != x) { p[x] = p[p[x]]; x = p[x]; } return x; };
|
||
auto uni = [&](int a, int b) { a = find(a); b = find(b); if (a != b) p[a] = b; };
|
||
for (size_t t = 0; t < m.tris.size(); t += 3) {
|
||
uni(m.tris[t], m.tris[t+1]); uni(m.tris[t+1], m.tris[t+2]);
|
||
}
|
||
std::unordered_map<int, int> faces;
|
||
for (size_t t = 0; t < m.tris.size(); t += 3) faces[find(m.tris[t])]++;
|
||
const int min_faces = std::max(1, (int) (min_frac * (double) nt));
|
||
std::vector<int> keep;
|
||
keep.reserve(m.tris.size());
|
||
for (size_t t = 0; t < m.tris.size(); t += 3) {
|
||
if (faces[find(m.tris[t])] >= min_faces) {
|
||
keep.push_back(m.tris[t]); keep.push_back(m.tris[t+1]); keep.push_back(m.tris[t+2]);
|
||
}
|
||
}
|
||
m.tris.swap(keep);
|
||
}
|
||
|
||
// Fill the small holes extract() leaves where a dual-grid quad was skipped (a
|
||
// neighbour voxel was missing). A boundary edge is one used by exactly one
|
||
// triangle; boundary edges chain into loops around each hole. Each loop up to
|
||
// `max_loop` vertices is fan-triangulated from its first vertex. Only triangles
|
||
// are added — no new vertices — so a parallel per-vertex attribute array (baked
|
||
// PBR) stays aligned. Large loops (genuine openings) are left unfilled.
|
||
inline void fill_holes(Mesh & m, int max_loop = 64, int max_passes = 4) {
|
||
if (m.n_tris() == 0) return;
|
||
auto key = [](int a, int b) {
|
||
const uint32_t lo = a < b ? a : b, hi = a < b ? b : a;
|
||
return ((uint64_t) lo << 32) | hi;
|
||
};
|
||
// Iterate: the greedy loop walk misses some loops at non-manifold junctions,
|
||
// and each fill can expose newly closeable loops, so repeat until a pass
|
||
// adds nothing (or the pass cap is hit).
|
||
for (int pass = 0; pass < max_passes; ++pass) {
|
||
const size_t before = m.tris.size();
|
||
// run-length count of undirected edges via sort (lighter than a hashmap)
|
||
std::vector<uint64_t> ek;
|
||
ek.reserve(m.tris.size());
|
||
for (size_t t = 0; t < m.tris.size(); t += 3) {
|
||
ek.push_back(key(m.tris[t], m.tris[t+1]));
|
||
ek.push_back(key(m.tris[t+1], m.tris[t+2]));
|
||
ek.push_back(key(m.tris[t+2], m.tris[t]));
|
||
}
|
||
std::sort(ek.begin(), ek.end());
|
||
std::unordered_map<int, std::vector<int>> adj; // boundary-vertex adjacency
|
||
for (size_t i = 0; i < ek.size(); ) {
|
||
size_t j = i + 1;
|
||
while (j < ek.size() && ek[j] == ek[i]) ++j;
|
||
if (j - i == 1) { // used by exactly one triangle -> boundary edge
|
||
const int a = (int) (ek[i] >> 32), b = (int) (ek[i] & 0xffffffffu);
|
||
adj[a].push_back(b); adj[b].push_back(a);
|
||
}
|
||
i = j;
|
||
}
|
||
if (adj.empty()) break;
|
||
|
||
std::unordered_set<uint64_t> used; // consumed boundary edges
|
||
for (const auto & kv : adj) {
|
||
const int start = kv.first;
|
||
for (const int nb0 : kv.second) {
|
||
if (used.count(key(start, nb0))) continue;
|
||
std::vector<int> loop{start};
|
||
int prev = start, cur = nb0;
|
||
used.insert(key(prev, cur));
|
||
bool closed = false;
|
||
while ((int) loop.size() <= max_loop) {
|
||
loop.push_back(cur);
|
||
if (cur == start) { closed = true; break; }
|
||
int next = -1;
|
||
auto it = adj.find(cur);
|
||
if (it != adj.end())
|
||
for (const int c : it->second)
|
||
if (c != prev && !used.count(key(cur, c))) { next = c; break; }
|
||
if (next < 0) break; // open chain / dead end
|
||
used.insert(key(cur, next));
|
||
prev = cur; cur = next;
|
||
}
|
||
if (!closed) continue;
|
||
loop.pop_back(); // drop the repeated start vertex
|
||
const int k = (int) loop.size();
|
||
if (k < 3 || k > max_loop) continue;
|
||
for (int t = 1; t < k - 1; ++t) { // fan-triangulate from loop[0]
|
||
m.tris.push_back(loop[0]);
|
||
m.tris.push_back(loop[t]);
|
||
m.tris.push_back(loop[t + 1]);
|
||
}
|
||
}
|
||
}
|
||
if (m.tris.size() == before) break; // converged
|
||
}
|
||
}
|
||
|
||
} // namespace fdg
|