wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-07-08 15:37:46 +08:00
Code Issues Actions 15 Packages Projects Releases Wiki Activity

Refactor/optimize UV unwrap

Browse Source
This commit is contained in:
kijai 2026-07-03 20:23:26 +03:00
parent 3b34e177cb
commit 5067ac461e
4 changed files with 1196 additions and 986 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1201
comfy_extras/mesh3d/uv_unwrap/pack.py
View File

File diff suppressed because it is too large Load Diff

316
comfy_extras/mesh3d/uv_unwrap/parameterize.py
View File

@ -12,6 +12,8 @@ from torch import Tensor
from . import mesh as _mesh
LSCM_BATCH_MAX_VERTS = 256 # charts above this solve per-chart sparse (lscm_chart)
def solve_least_squares(A: sp.csr_matrix, b: np.ndarray) -> np.ndarray:
"""Solve ||Ax - b||^2 by factorizing AtA."""
@ -99,54 +101,258 @@ def _ortho_project(verts_3d: np.ndarray) -> np.ndarray:
return np.stack([verts_3d @ t, verts_3d @ b], axis=1)
def _stretch_metrics(verts_3d: np.ndarray, uvs: np.ndarray, faces: np.ndarray) -> Tuple[float, float, int, int]:
"""Sander's stretch metric. Returns (rms, max, n_flipped, n_zero_area)."""
p = verts_3d[faces]
def ortho_project_concat(verts: np.ndarray, chart_of_vert: np.ndarray, n_charts: int) -> np.ndarray:
"""_ortho_project for every chart at once over concatenated per-chart vertices."""
cnt = np.bincount(chart_of_vert, minlength=n_charts).clip(min=1).astype(np.float64)
cen = np.stack([np.bincount(chart_of_vert, weights=verts[:, i], minlength=n_charts)
for i in range(3)], axis=1) / cnt[:, None]
d = verts - cen[chart_of_vert]
cov = np.zeros((n_charts, 3, 3), dtype=np.float64)
for i in range(3):
for j in range(i, 3):
s = np.bincount(chart_of_vert, weights=d[:, i] * d[:, j], minlength=n_charts)
cov[:, i, j] = s
cov[:, j, i] = s
_w, ev = np.linalg.eigh(cov)
normal = ev[:, :, 0]
t = np.eye(3, dtype=np.float64)[np.argmin(np.abs(normal), axis=1)]
t = t - normal * (normal * t).sum(axis=1, keepdims=True)
t /= np.linalg.norm(t, axis=1, keepdims=True).clip(min=1e-20)
b = np.cross(normal, t)
tt, bb = t[chart_of_vert], b[chart_of_vert]
return np.stack([(verts * tt).sum(1), (verts * bb).sum(1)], axis=1)
def stretch_metrics_concat(
verts: np.ndarray, uvs: np.ndarray, faces: np.ndarray,
chart_of_face: np.ndarray, n_charts: int,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""Per-chart Sander stretch metrics (rms, max, n_flipped, n_zero_area); rms/max inf where undefined."""
p = verts[faces]
t = uvs[faces]
parametric_area = 0.5 * (
pa_signed = 0.5 * (
(t[:, 1, 1] - t[:, 0, 1]) * (t[:, 2, 0] - t[:, 0, 0])
- (t[:, 2, 1] - t[:, 0, 1]) * (t[:, 1, 0] - t[:, 0, 0])
)
n_flipped = int((parametric_area < -1e-12).sum())
n_zero = int((np.abs(parametric_area) < 1e-12).sum())
pa = np.abs(parametric_area).clip(min=1e-20)
geom_area = 0.5 * np.linalg.norm(
np.cross(p[:, 1] - p[:, 0], p[:, 2] - p[:, 0]), axis=1
)
keep = (geom_area > 1e-12) & (np.abs(parametric_area) > 1e-12)
if not keep.any():
return float("inf"), float("inf"), n_flipped, n_zero
t1 = t[:, 0, 0]
s1 = t[:, 0, 1]
t2 = t[:, 1, 0]
s2 = t[:, 1, 1]
t3 = t[:, 2, 0]
s3 = t[:, 2, 1]
- (t[:, 2, 1] - t[:, 0, 1]) * (t[:, 1, 0] - t[:, 0, 0]))
n_flip = np.bincount(chart_of_face[pa_signed < -1e-12], minlength=n_charts)
n_zero = np.bincount(chart_of_face[np.abs(pa_signed) < 1e-12], minlength=n_charts)
pa = np.abs(pa_signed).clip(min=1e-20)
ga = 0.5 * np.linalg.norm(np.cross(p[:, 1] - p[:, 0], p[:, 2] - p[:, 0]), axis=1)
keep = (ga > 1e-12) & (np.abs(pa_signed) > 1e-12)
t1, s1 = t[:, 0, 0], t[:, 0, 1]
t2, s2 = t[:, 1, 0], t[:, 1, 1]
t3, s3 = t[:, 2, 0], t[:, 2, 1]
inv_2pa = 1.0 / (2.0 * pa)
Ss = (
p[:, 0] * (t2 - t3)[:, None]
+ p[:, 1] * (t3 - t1)[:, None]
+ p[:, 2] * (t1 - t2)[:, None]
) * inv_2pa[:, None]
St = (
p[:, 0] * (s3 - s2)[:, None]
+ p[:, 1] * (s1 - s3)[:, None]
+ p[:, 2] * (s2 - s1)[:, None]
) * inv_2pa[:, None]
Ss = (p[:, 0] * (t2 - t3)[:, None] + p[:, 1] * (t3 - t1)[:, None]
+ p[:, 2] * (t1 - t2)[:, None]) * inv_2pa[:, None]
St = (p[:, 0] * (s3 - s2)[:, None] + p[:, 1] * (s1 - s3)[:, None]
+ p[:, 2] * (s2 - s1)[:, None]) * inv_2pa[:, None]
a = (Ss * Ss).sum(axis=1)
bb = (Ss * St).sum(axis=1)
c = (St * St).sum(axis=1)
sigma2_sq = 0.5 * (a + c + np.sqrt(np.maximum(0.0, (a - c) ** 2 + 4 * bb ** 2)))
rms_sq = (a + c) * 0.5
rms_stretch_sq_sum = float((rms_sq[keep] * geom_area[keep]).sum())
total_geom = float(geom_area[keep].sum())
total_param = float(pa[keep].sum())
if total_geom <= 0.0:
return float("inf"), float("inf"), n_flipped, n_zero
norm_factor = np.sqrt(total_param / total_geom)
rms_stretch = float(np.sqrt(rms_stretch_sq_sum / total_geom)) * norm_factor
max_stretch = float(np.sqrt(sigma2_sq[keep].max())) * norm_factor
return rms_stretch, max_stretch, n_flipped, n_zero
cf = chart_of_face[keep]
tg = np.bincount(cf, weights=ga[keep], minlength=n_charts)
tp = np.bincount(cf, weights=pa[keep], minlength=n_charts)
rs = np.bincount(cf, weights=(rms_sq * ga)[keep], minlength=n_charts)
smax = np.zeros(n_charts, dtype=np.float64)
np.maximum.at(smax, cf, sigma2_sq[keep])
ok = tg > 0.0
tg_safe = np.where(ok, tg, 1.0)
norm = np.sqrt(tp / tg_safe)
rms = np.where(ok, np.sqrt(rs / tg_safe) * norm, np.inf)
mx = np.where(ok, np.sqrt(smax) * norm, np.inf)
return rms, mx, n_flip, n_zero
def _segment_argmax(vals: np.ndarray, seg: np.ndarray, n: int) -> np.ndarray:
"""Index of the (first) max element per segment; -1 for empty segments."""
amax = np.full(n, -np.inf)
np.maximum.at(amax, seg, vals)
hit = vals == amax[seg]
out = np.full(n, np.iinfo(np.int64).max, dtype=np.int64)
np.minimum.at(out, seg[hit], np.nonzero(hit)[0])
return np.where(out == np.iinfo(np.int64).max, -1, out)
def lscm_charts_batch(
verts: np.ndarray, # (sumV, 3) float64, per-chart concatenated
uv_pins: np.ndarray, # (sumV, 2) float64, ortho UVs (pin values + fallback)
faces_gl: np.ndarray, # (sumF, 3) global-local ids into verts
face_pos: np.ndarray, # (sumF,) row index of each face within its chart
chart_of_face: np.ndarray, # (sumF,)
chart_of_vert: np.ndarray, # (sumV,)
vert_offsets: np.ndarray, # (n_charts+1,)
chart_ids: np.ndarray, # charts to solve (each with >=3 verts, >=1 face)
n_charts: int,
max_bucket_verts: int = LSCM_BATCH_MAX_VERTS,
device: "torch.device | None" = None,
) -> dict:
"""Batched dense ABF/LSCM; returns {chart_id: (Vc, 2) float32}. Charts larger than
max_bucket_verts are left out (the caller solves those sparse)."""
out: dict = {}
if chart_ids.size == 0:
return out
sel = np.zeros(n_charts, dtype=bool)
sel[chart_ids] = True
vcounts = np.diff(vert_offsets)
# ABF coefficients for all selected faces in one shot
fmask = sel[chart_of_face]
f_ids = np.nonzero(fmask)[0]
abf_ids, abf_cos, abf_sin, abf_valid = _abf_face_coefficients(verts, faces_gl[f_ids])
# farthest-point pin pair per chart (two passes)
vmask = sel[chart_of_vert]
v_ids = np.nonzero(vmask)[0]
cv = chart_of_vert[v_ids]
first = vert_offsets[:-1]
d0 = ((verts[v_ids] - verts[first[cv]]) ** 2).sum(1)
pin_a = _segment_argmax(d0, cv, n_charts) # global vert index (into v_ids space)
pin_a = np.where(pin_a >= 0, v_ids[pin_a.clip(min=0)], -1)
d1 = ((verts[v_ids] - verts[pin_a.clip(min=0)[cv]]) ** 2).sum(1)
pin_b = _segment_argmax(d1, cv, n_charts)
pin_b = np.where(pin_b >= 0, v_ids[pin_b.clip(min=0)], -1)
# degenerate (all verts coincide): any distinct vert within the chart (Vc >= 3 guaranteed)
alt = np.where(pin_a == first, first + 1, first)
pin_b = np.where(pin_a == pin_b, alt, pin_b)
fcounts = np.bincount(chart_of_face[f_ids], minlength=n_charts)
# size-sorted chunks padded to their own max, bounded by an element budget so one
# face-heavy chart can't inflate a whole chunk
small = chart_ids[vcounts[chart_ids] <= max_bucket_verts]
sorted_ids = small[np.argsort(vcounts[small], kind="stable")]
budget = (96 << 20) // 8 # float64 elements in a chunk's A
chunks = []
cs = 0
fmax_r = vmax_r = 0
for idx in range(sorted_ids.size):
c2 = sorted_ids[idx]
fm2 = max(fmax_r, int(fcounts[c2]))
vm2 = max(vmax_r, int(vcounts[c2]))
nb = idx - cs + 1
if nb > 1 and (nb > 128 or nb * 4 * fm2 * vm2 > budget):
chunks.append((cs, idx))
cs = idx
fmax_r, vmax_r = int(fcounts[c2]), int(vcounts[c2])
else:
fmax_r, vmax_r = fm2, vm2
if sorted_ids.size:
chunks.append((cs, sorted_ids.size))
for s, e in chunks:
cids = sorted_ids[s:e]
B = cids.size
Vmax = int(vcounts[cids].max())
Fmax = int(fcounts[cids].max())
N = 2 * Vmax
R = 2 * Fmax
compact = np.full(n_charts, -1, dtype=np.int64)
compact[cids] = np.arange(B)
fm = compact[chart_of_face[f_ids]] >= 0
fi = f_ids[fm] # face rows for this chunk
bi = compact[chart_of_face[fi]] # chart slot per face
frow = face_pos[fi]
v0 = vert_offsets[chart_of_face[fi]] # local id = global-local - v0
am = fm.nonzero()[0] # index into abf_* arrays
pieces_i: list = []
pieces_v: list = []
def scatter(rows, cols, vals, bsel):
pieces_i.append((bsel * R + rows) * N + cols)
pieces_v.append(vals)
val = abf_valid[am]
ii = am[val]
ids = abf_ids[ii] - v0[val, None] # local vert ids, reordered
cosf, sinf = abf_cos[ii], abf_sin[ii]
rr, bsel = frow[val] * 2, bi[val]
ones = np.ones(ii.size)
for cc2, vv in ((ids[:, 0], cosf - 1.0), (ids[:, 0] + Vmax, -sinf),
(ids[:, 1], -cosf), (ids[:, 1] + Vmax, sinf), (ids[:, 2], ones)):
scatter(rr, cc2, vv, bsel)
for cc2, vv in ((ids[:, 0], sinf), (ids[:, 0] + Vmax, cosf - 1.0),
(ids[:, 1], -sinf), (ids[:, 1] + Vmax, -cosf), (ids[:, 2] + Vmax, ones)):
scatter(rr + 1, cc2, vv, bsel)
inv = ~val
if inv.any():
jj = fi[inv]
tri2d = _triangle_local_2d(verts, faces_gl[jj])
twice = tri2d[:, 1, 0] * tri2d[:, 2, 1] - tri2d[:, 1, 1] * tri2d[:, 2, 0]
w = 1.0 / np.sqrt(2.0 * np.abs(twice).clip(min=1e-20))
rr2, bs2 = frow[inv] * 2, bi[inv]
lids = faces_gl[jj] - v0[inv, None]
for j in range(3):
jp1, jp2 = (j + 1) % 3, (j + 2) % 3
aj = (tri2d[:, jp1, 0] - tri2d[:, jp2, 0]) * w
bj = (tri2d[:, jp1, 1] - tri2d[:, jp2, 1]) * w
vc2 = lids[:, j]
scatter(rr2, vc2, aj, bs2)
scatter(rr2, vc2 + Vmax, -bj, bs2)
scatter(rr2 + 1, vc2, bj, bs2)
scatter(rr2 + 1, vc2 + Vmax, aj, bs2)
flat = np.concatenate(pieces_i)
A = np.bincount(flat, weights=np.concatenate(pieces_v),
minlength=B * R * N).reshape(B, R, N)
# pins: move their columns to the RHS, then constrain via identity rows
voff = vert_offsets[cids]
pa_l = pin_a[cids] - voff
pb_l = pin_b[cids] - voff
pin_cols = np.stack([pa_l, pb_l, pa_l + Vmax, pb_l + Vmax], 1) # (B,4)
pin_vals = np.stack([uv_pins[pin_a[cids], 0], uv_pins[pin_b[cids], 0],
uv_pins[pin_a[cids], 1], uv_pins[pin_b[cids], 1]], 1)
rhs = np.zeros((B, R), dtype=np.float64)
barange = np.arange(B)
for k in range(4):
rhs -= A[barange, :, pin_cols[:, k]] * pin_vals[:, k, None]
A[barange, :, pin_cols[:, k]] = 0.0
# constrained columns: the 4 pins + padding beyond each chart's vert count
vcs = vcounts[cids]
padm = np.arange(Vmax)[None, :] >= vcs[:, None]
con = np.concatenate([padm, padm], axis=1) # (B,N)
np.put_along_axis(con, pin_cols, True, axis=1)
cval = np.zeros((B, N), dtype=np.float64)
np.put_along_axis(cval, pin_cols, pin_vals, axis=1)
# normal equations + batched solve; the fp64 dense algebra goes to the GPU when available
use_gpu = device is not None and device.type == "cuda"
if use_gpu:
A_t = torch.from_numpy(A).to(device)
At = A_t.transpose(1, 2)
AtA = At @ A_t
Atb = (At @ torch.from_numpy(rhs).to(device).unsqueeze(2)).squeeze(2)
con_t = torch.from_numpy(con).to(device)
free2 = (~con_t[:, :, None]) & (~con_t[:, None, :])
AtA = AtA * free2
diag = torch.diagonal(AtA, dim1=1, dim2=2)
# median (not max) positive diagonal: a degenerate face's ~1e19 squared row
# weight would blow a max-scaled eps past the unit ABF rows
dpos = torch.where(diag > 0, diag, torch.full_like(diag, float("nan")))
dsc = 1e-12 * torch.nan_to_num(dpos.nanmedian(dim=1).values, nan=1e-8).clamp_min(1e-20)
diag += torch.where(con_t, torch.ones_like(diag), dsc[:, None].expand_as(diag))
Atb = torch.where(con_t, torch.from_numpy(cval).to(device), Atb)
x = torch.linalg.solve(AtA, Atb).cpu().numpy()
else:
At = A.transpose(0, 2, 1)
AtA = At @ A # batched BLAS dgemm
Atb = (At @ rhs[:, :, None])[:, :, 0]
AtA *= (~con[:, :, None]) & (~con[:, None, :])
dg = AtA.reshape(B, -1)[:, ::N + 1]
# median positive diagonal (see GPU branch): robust to degenerate-face weights
dpos = np.where(dg > 0, dg, np.nan)
with np.errstate(all="ignore"):
dsc = 1e-12 * np.nan_to_num(np.nanmedian(dpos, axis=1), nan=1e-8).clip(min=1e-20)
dg += np.where(con, 1.0, dsc[:, None])
Atb2 = np.where(con, cval, Atb)
x = np.linalg.solve(AtA, Atb2)
for i2, c2 in enumerate(cids):
vc3 = int(vcs[i2])
out[int(c2)] = np.stack([x[i2, :vc3], x[i2, Vmax:Vmax + vc3]], 1).astype(np.float32)
return out
def _uv_boundary_self_intersects(
@ -181,36 +387,6 @@ def _uv_boundary_self_intersects(
return False
def parametrize_chart(
local_verts: Tensor, local_faces: Tensor, local_face_face: Tensor
) -> Tensor:
"""Parameterize one chart: ortho first, ABF/LSCM fallback; charts <=5 faces stay ortho."""
verts_np = local_verts.detach().cpu().numpy().astype(np.float64)
faces_np = local_faces.detach().cpu().numpy().astype(np.int64)
if verts_np.shape[0] < 3 or faces_np.shape[0] == 0:
return torch.zeros((verts_np.shape[0], 2), dtype=torch.float32, device=local_verts.device)
ortho = _ortho_project(verts_np)
n_faces = faces_np.shape[0]
if n_faces <= 5:
return torch.from_numpy(ortho.astype(np.float32)).to(local_verts.device)
rms, mx, n_flip, n_zero = _stretch_metrics(verts_np, ortho, faces_np)
flip_ok = n_flip == 0 or n_flip == n_faces
if flip_ok and n_zero == 0 and rms <= 1.5 and mx <= 2.0:
ff_np = local_face_face.detach().cpu().numpy().astype(np.int64)
if not _uv_boundary_self_intersects(ortho, faces_np, ff_np):
return torch.from_numpy(ortho.astype(np.float32)).to(local_verts.device)
uvs_t = lscm_chart(local_verts, local_faces, local_face_face, pin_positions=ortho)
# Collapsed UV island (aspect > 100:1) blows up packing scale; fall back to ortho.
uvs_np = uvs_t.detach().cpu().numpy()
bbox = uvs_np.max(axis=0) - uvs_np.min(axis=0)
bbox_max = float(max(bbox[0], bbox[1], 1e-12))
bbox_min = float(max(min(bbox[0], bbox[1]), 1e-12))
if bbox_max / bbox_min > 100.0:
return torch.from_numpy(ortho.astype(np.float32)).to(local_verts.device)
return uvs_t
def _abf_face_coefficients(
verts_3d: np.ndarray, faces: np.ndarray
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:

499
comfy_extras/mesh3d/uv_unwrap/segment.py
View File

@ -1,22 +1,11 @@
"""Adaptive cost-grow chart segmentation (CPU); numba optional, numpy path is nd-only."""
"""Adaptive cost-grow chart segmentation (vectorized torch, CPU or GPU)."""
from __future__ import annotations
from typing import List, Tuple
from typing import Tuple
import numpy as np
import torch
from torch import Tensor
try:
from numba import njit
_HAVE_NUMBA = True
except ImportError:
_HAVE_NUMBA = False
def njit(*args, **kwargs): # noqa: ARG001
def deco(fn):
return fn
return deco if not args else args[0]
from tqdm import tqdm
from .mesh import MeshData, face_edge_lengths
@ -28,129 +17,63 @@ DEFAULT_MAX_COST = 2.0
NORMAL_DEVIATION_HARD_CUTOFF = 0.707 # ~75°
@njit(cache=True, fastmath=False)
def _cost_grow_iter_jit(
face_chart: np.ndarray, face_face: np.ndarray, face_normal: np.ndarray,
face_area: np.ndarray, face_edge_len: np.ndarray,
chart_basis: np.ndarray, chart_normal_sum: np.ndarray,
chart_area: np.ndarray, chart_perim: np.ndarray,
nd_cutoff: float, max_cost: float,
w_nd: float, w_round: float, w_straight: float,
):
"""One grow iter: each unassigned face joins its lowest-cost adjacent chart if cost<max_cost."""
F = face_chart.shape[0]
best_chart_per_face = np.full(F, -1, dtype=np.int64)
best_cost_per_face = np.full(F, np.inf, dtype=np.float32)
def _grow_iter(face_chart, frontier, ff, fn, fa, fel, basis, nsum, area, perim, K,
nd_cutoff, tau, w_nd, w_round, w_straight):
"""One grow pass: each frontier face joins its lowest-cost adjacent chart if cost <= tau;
returns the number of faces assigned."""
u = frontier.nonzero(as_tuple=True)[0]
if u.numel() == 0:
return 0
nb = ff[u] # (U,3) neighbor face ids
nbc = torch.where(nb >= 0, face_chart[nb.clamp_min(0)], nb.new_full((), -1))
valid = nbc >= 0
d = (fn[u][:, None, :] * basis[nbc.clamp_min(0)]).sum(-1)
nd = (1.0 - d).clamp(0.0, 1.0)
valid &= nd < nd_cutoff
el = fel[u] # (U,3)
# l_in per candidate chart j: edge k counts if its (assigned) neighbor is in chart j
inm = (nbc[:, :, None] == nbc[:, None, :]) & valid[:, None, :]
l_in = (el[:, None, :] * inm).sum(-1) # (U,3)
tot = el.sum(-1, keepdim=True)
l_out = tot - l_in
ca = area[nbc.clamp_min(0)]
cp = perim[nbc.clamp_min(0)]
new_perim = cp - l_in + l_out
new_r = new_perim * new_perim / (ca + fa[u][:, None]).clamp_min(1e-20)
round_cost = torch.where((cp <= 1e-20) | (ca <= 1e-20) | (new_r <= 1e-20),
torch.zeros_like(new_r),
1.0 - (cp * cp / ca.clamp_min(1e-20)) / new_r.clamp_min(1e-20))
straight_cost = ((l_out - l_in) / tot.clamp_min(1e-20)).clamp(max=0.0)
cost = w_nd * nd + w_round * round_cost + w_straight * straight_cost
cost = torch.where(valid, cost, cost.new_full((), float("inf")))
best_cost, best_j = cost.min(1)
acc = best_cost <= tau
n_acc = int(acc.sum())
if n_acc == 0:
return 0
for f in range(F):
if face_chart[f] != -1:
continue
nx = face_normal[f, 0]
ny = face_normal[f, 1]
nz = face_normal[f, 2]
af = face_area[f]
for e0 in range(3):
nb0 = face_face[f, e0]
if nb0 < 0:
continue
c = face_chart[nb0]
if c < 0:
continue
d = (nx * chart_basis[c, 0] + ny * chart_basis[c, 1] + nz * chart_basis[c, 2])
nd = np.float32(1.0) - d
if nd > np.float32(1.0):
nd = np.float32(1.0)
if nd < np.float32(0.0):
nd = np.float32(0.0)
if nd >= nd_cutoff:
continue
l_in = np.float32(0.0)
l_out = np.float32(0.0)
for e1 in range(3):
nb1 = face_face[f, e1]
el = face_edge_len[f, e1]
if nb1 < 0:
l_out += el
elif face_chart[nb1] == c:
l_in += el
else:
l_out += el
ca = chart_area[c]
cp = chart_perim[c]
new_perim = cp - l_in + l_out
new_area = ca + af
if cp <= np.float32(1e-20) or ca <= np.float32(1e-20):
round_cost = np.float32(0.0)
else:
old_r = (cp * cp) / ca
new_r = (new_perim * new_perim) / new_area
if new_r <= np.float32(1e-20):
round_cost = np.float32(0.0)
else:
round_cost = np.float32(1.0) - old_r / new_r
denom = l_out + l_in
if denom <= np.float32(1e-20):
straight_cost = np.float32(0.0)
else:
ratio = (l_out - l_in) / denom
if ratio < np.float32(0.0):
straight_cost = ratio
else:
straight_cost = np.float32(0.0)
cost = (w_nd * nd + w_round * round_cost + w_straight * straight_cost)
if cost < best_cost_per_face[f]:
best_cost_per_face[f] = cost
best_chart_per_face[f] = c
n_assigned = 0
for f in range(F):
if face_chart[f] != -1:
continue
if best_chart_per_face[f] < 0:
continue
if best_cost_per_face[f] > max_cost:
continue
c = best_chart_per_face[f]
l_in = np.float32(0.0)
l_out = np.float32(0.0)
for e1 in range(3):
nb1 = face_face[f, e1]
el = face_edge_len[f, e1]
if nb1 < 0:
l_out += el
elif face_chart[nb1] == c:
l_in += el
else:
l_out += el
af = face_area[f]
face_chart[f] = c
chart_normal_sum[c, 0] += face_normal[f, 0] * af
chart_normal_sum[c, 1] += face_normal[f, 1] * af
chart_normal_sum[c, 2] += face_normal[f, 2] * af
chart_area[c] += af
chart_perim[c] = chart_perim[c] - l_in + l_out
nx = chart_normal_sum[c, 0]
ny = chart_normal_sum[c, 1]
nz = chart_normal_sum[c, 2]
nlen = np.sqrt(nx * nx + ny * ny + nz * nz)
if nlen > np.float32(1e-20):
chart_basis[c, 0] = nx / nlen
chart_basis[c, 1] = ny / nlen
chart_basis[c, 2] = nz / nlen
n_assigned += 1
return n_assigned
def _renumber(face_chart: np.ndarray, device) -> Tensor:
unique = np.unique(face_chart[face_chart >= 0])
if unique.size == 0:
return torch.from_numpy(face_chart).to(device)
remap = np.full(int(unique.max()) + 1, -1, dtype=np.int64)
remap[unique] = np.arange(unique.size)
out = face_chart.copy()
mask = out >= 0
out[mask] = remap[out[mask]]
return torch.from_numpy(out).to(device)
f_acc = u[acc]
c_acc = nbc.gather(1, best_j[:, None]).squeeze(1)[acc]
nbc_old = nbc[acc] # neighbor charts before this commit
face_chart[f_acc] = c_acc
nb_acc = nb[acc]
nbs_acc = nb_acc.clamp_min(0)
nbc_post = torch.where(nb_acc >= 0, face_chart[nbs_acc], nb_acc.new_full((), -1))
# frontier update: committed faces leave; their still-unassigned neighbors enter
frontier[f_acc] = False
grow_nb = nbs_acc[(nb_acc >= 0) & (nbc_post < 0)]
frontier[grow_nb] = True
el_acc = el[acc]
cx = c_acc[:, None]
dper = torch.where(nbc_old == cx, -el_acc, # was member: edge turns interior
torch.where(nbc_post == cx, torch.zeros_like(el_acc), # co-committer
el_acc)).sum(1) # boundary / other chart
perim.scatter_add_(0, c_acc, dper)
area.scatter_add_(0, c_acc, fa[f_acc])
nsum.index_add_(0, c_acc, fn[f_acc] * fa[f_acc, None])
nl = nsum[:K].norm(dim=1, keepdim=True)
basis[:K] = torch.where(nl > 1e-20, nsum[:K] / nl.clamp_min(1e-20), basis[:K])
return n_acc
def segment_charts(
@ -166,162 +89,108 @@ def segment_charts(
if F == 0:
return torch.zeros(0, dtype=torch.long, device=device)
face_normal = mesh.face_normal.detach().cpu().numpy().astype(np.float32)
face_area = mesh.face_area.detach().cpu().numpy().astype(np.float32)
face_centroid = mesh.face_centroid.detach().cpu().numpy().astype(np.float32)
face_face = mesh.face_face.detach().cpu().numpy()
fn = mesh.face_normal.detach().to(torch.float32)
fa = mesh.face_area.detach().to(torch.float32)
fc = mesh.face_centroid.detach().to(torch.float32)
ff = mesh.face_face.detach().long()
fel = face_edge_lengths(mesh.vertices, mesh.faces).detach().to(torch.float32)
nd_cutoff = NORMAL_DEVIATION_HARD_CUTOFF
face_chart = np.full(F, -1, dtype=np.int64)
nd_cutoff = np.float32(NORMAL_DEVIATION_HARD_CUTOFF)
nd_threshold = np.float32(min(max_cost / max(w_normal_deviation, 1e-6),
NORMAL_DEVIATION_HARD_CUTOFF * 0.99))
component = (mesh.component.detach().cpu().numpy()
if hasattr(mesh.component, "detach") else np.asarray(mesh.component))
if component.size:
_, first_idx = np.unique(component, return_index=True)
initial_seeds = first_idx.astype(np.int64)
# one seed per connected component (first face of each)
comp = mesh.component.detach().long().to(device)
ncomp = int(comp.max()) + 1 if comp.numel() else 0
if ncomp:
seeds = torch.full((ncomp,), F, dtype=torch.long, device=device)
seeds.scatter_reduce_(0, comp, torch.arange(F, device=device), reduce="amin")
else:
initial_seeds = np.empty(0, dtype=np.int64)
seeds = torch.zeros(1, dtype=torch.long, device=device)
K = seeds.shape[0]
seed_faces: List[int] = [int(s) for s in initial_seeds.tolist()]
if not seed_faces:
seed_faces = [0]
max_total_charts = max(F, 8000)
cap = K + F + 1 # every re-seed assigns a face, so K < K0 + F
face_chart = torch.full((F,), -1, dtype=torch.long, device=device)
basis = torch.zeros(cap, 3, dtype=torch.float32, device=device)
nsum = torch.zeros(cap, 3, dtype=torch.float32, device=device)
area = torch.zeros(cap, dtype=torch.float32, device=device)
perim = torch.zeros(cap, dtype=torch.float32, device=device)
face_chart[seeds] = torch.arange(K, device=device)
basis[:K] = fn[seeds]
nsum[:K] = fn[seeds] * fa[seeds, None]
area[:K] = fa[seeds]
perim[:K] = fel[seeds].sum(1)
frontier = torch.zeros(F, dtype=torch.bool, device=device)
seed_nb = ff[seeds]
seed_nb = seed_nb[seed_nb >= 0]
frontier[seed_nb] = True
frontier &= face_chart < 0
K = len(seed_faces)
chart_basis = np.zeros((K, 3), dtype=np.float32)
chart_normal_sum = np.zeros((K, 3), dtype=np.float32)
chart_area = np.zeros(K, dtype=np.float32)
chart_perim = np.zeros(K, dtype=np.float32)
face_edge_len = (
face_edge_lengths(mesh.vertices, mesh.faces)
.detach().cpu().numpy()
)
for cid, sf in enumerate(seed_faces):
face_chart[sf] = cid
n = face_normal[sf]
a = face_area[sf]
chart_basis[cid] = n.astype(np.float32)
chart_normal_sum[cid] = (n * a).astype(np.float32)
chart_area[cid] = float(a)
chart_perim[cid] = float(face_edge_len[sf].sum())
min_d2 = torch.full((F,), float("inf"), dtype=torch.float32, device=device)
for i in range(0, K, 32): # chunked: (F, <=32, 3) stays small
d2 = ((fc[:, None, :] - fc[seeds[i:i + 32]][None, :, :]) ** 2).sum(-1)
min_d2 = torch.minimum(min_d2, d2.amin(1))
if K == 0:
return _renumber(face_chart, device)
# Multi-pass threshold schedule (low-cost first); tau cap 0.5 keeps cones ~30deg.
tau_final = min(max_cost * 0.25, 0.5)
thresholds = [t for t in (0.05, 0.1, 0.25) if t < tau_final] + [tau_final]
max_inner = max(64, int(F ** 0.5) * 2)
outer_iter = 0
tq = tqdm(total=F, desc="unwrap: segment (adaptive)", unit="face", leave=False)
while True:
outer_iter += 1
if outer_iter > F + 16:
break
for tau in thresholds:
for _ in range(max_inner):
n_added = _grow_iter(face_chart, frontier, ff, fn, fa, fel, basis, nsum,
area, perim, K, nd_cutoff, tau, w_normal_deviation,
w_roundness, w_straightness)
if n_added == 0:
break
tq.update(n_added)
unassigned = face_chart < 0
if int(unassigned.sum()) == 0:
break
if K >= max_total_charts:
break
# re-seed at the unassigned face farthest from every existing seed
new_seed = int(torch.where(unassigned, min_d2,
min_d2.new_full((), float("-inf"))).argmax())
face_chart[new_seed] = K
basis[K] = fn[new_seed]
nsum[K] = fn[new_seed] * fa[new_seed]
area[K] = fa[new_seed]
perim[K] = fel[new_seed].sum()
K += 1
min_d2 = torch.minimum(min_d2, ((fc - fc[new_seed]) ** 2).sum(-1))
tq.update(1)
frontier[new_seed] = False
ns_nb = ff[new_seed]
ns_nb = ns_nb[ns_nb >= 0]
frontier[ns_nb[face_chart[ns_nb] < 0]] = True
min_dist_to_seed = np.full(F, np.inf, dtype=np.float32)
for sf in seed_faces:
d = ((face_centroid - face_centroid[sf]) ** 2).sum(axis=-1)
min_dist_to_seed = np.minimum(min_dist_to_seed, d)
if _HAVE_NUMBA:
# Multi-pass threshold schedule (low-cost first); tau cap 0.5 keeps cones ~30deg.
tau_final = min(max_cost * 0.25, 0.5)
thresholds = [t for t in (0.05, 0.1, 0.25) if t < tau_final] + [tau_final]
max_inner = max(64, int(np.sqrt(F)) * 2)
max_total_charts = max(F, 8000)
outer_iter = 0
while True:
outer_iter += 1
if outer_iter > F + 16:
break
for tau in thresholds:
for _ in range(max_inner):
n_added = _cost_grow_iter_jit(
face_chart, face_face, face_normal, face_area, face_edge_len,
chart_basis, chart_normal_sum, chart_area, chart_perim,
nd_cutoff, np.float32(tau),
np.float32(w_normal_deviation),
np.float32(w_roundness),
np.float32(w_straightness),
)
if n_added == 0:
break
if (face_chart == -1).sum() == 0:
break
if chart_basis.shape[0] >= max_total_charts:
break
unassigned_mask = face_chart == -1
cand = np.where(unassigned_mask, min_dist_to_seed, np.float32(-np.inf))
new_seed = int(np.argmax(cand))
n = face_normal[new_seed]
a = face_area[new_seed]
chart_basis = np.vstack([chart_basis, n[None, :].astype(np.float32)])
chart_normal_sum = np.vstack(
[chart_normal_sum, (n * a)[None, :].astype(np.float32)]
)
chart_area = np.concatenate([chart_area, np.array([a], dtype=np.float32)])
chart_perim = np.concatenate(
[chart_perim, np.array([face_edge_len[new_seed].sum()], dtype=np.float32)]
)
face_chart[new_seed] = chart_basis.shape[0] - 1
new_d = ((face_centroid - face_centroid[new_seed]) ** 2).sum(axis=-1)
min_dist_to_seed = np.minimum(min_dist_to_seed, new_d)
else:
# Numpy fallback: nd-only adaptive grow.
for _ in range(max(64, int(np.sqrt(F)) + 32)):
unassigned = face_chart == -1
if not unassigned.any():
break
u_idx = np.nonzero(unassigned)[0]
nbs = face_face[u_idx]
nbs_safe = np.where(nbs >= 0, nbs, 0)
nb_charts = np.where(nbs >= 0, face_chart[nbs_safe], -1)
valid = (nb_charts >= 0)
if not valid.any():
break
nb_charts_safe = np.where(valid, nb_charts, 0)
nb_basis = chart_basis[nb_charts_safe]
d = (face_normal[u_idx][:, None, :] * nb_basis).sum(axis=-1)
nd = np.where(valid, np.float32(1.0) - d, np.inf).clip(max=1.0)
nd = np.where(nd >= nd_cutoff, np.inf, nd)
best_e = np.argmin(nd, axis=1)
best_cost = nd[np.arange(u_idx.size), best_e]
best_c = nb_charts_safe[np.arange(u_idx.size), best_e]
accept = (best_cost <= nd_threshold) & np.isfinite(best_cost)
if not accept.any():
break
pick_u = u_idx[accept]
pick_c = best_c[accept]
face_chart[pick_u] = pick_c
for f, c in zip(pick_u, pick_c):
chart_normal_sum[c] += face_normal[f] * face_area[f]
chart_area[c] += face_area[f]
tq.close()
# Orphan cleanup: leftover faces join their best-matching neighbor's chart.
if (face_chart == -1).any() and chart_basis.shape[0] > 0:
while True:
orphans = np.nonzero(face_chart == -1)[0]
if orphans.size == 0:
break
nbs = face_face[orphans]
nbs_safe = np.where(nbs >= 0, nbs, 0)
nb_charts = np.where(nbs >= 0, face_chart[nbs_safe], -1)
valid = (nb_charts >= 0)
if not valid.any():
break
nb_charts_safe = np.where(valid, nb_charts, 0)
nb_basis = chart_basis[nb_charts_safe]
d = (face_normal[orphans][:, None, :] * nb_basis).sum(axis=-1)
nd = np.where(valid, np.float32(1.0) - d, np.inf)
best_e = np.argmin(nd, axis=1)
best_c = nb_charts_safe[np.arange(orphans.size), best_e]
assignable = valid.any(axis=1)
if not assignable.any():
break
assign_idx = orphans[assignable]
assign_c = best_c[assignable]
face_chart[assign_idx] = assign_c
if (face_chart == -1).any():
new_singletons = np.nonzero(face_chart == -1)[0]
for f in new_singletons:
face_chart[int(f)] = chart_basis.shape[0]
chart_basis = np.concatenate(
[chart_basis, face_normal[int(f)].astype(np.float32)[None, :]],
axis=0,
)
while True:
orphans = (face_chart < 0).nonzero(as_tuple=True)[0]
if orphans.numel() == 0:
break
nb = ff[orphans]
nbc = torch.where(nb >= 0, face_chart[nb.clamp_min(0)], nb.new_full((), -1))
valid = nbc >= 0
assignable = valid.any(1)
if not bool(assignable.any()):
break
d = (fn[orphans][:, None, :] * basis[nbc.clamp_min(0)]).sum(-1)
ndv = torch.where(valid, 1.0 - d, d.new_full((), float("inf")))
best_c = nbc.gather(1, ndv.argmin(1, keepdim=True)).squeeze(1)
face_chart[orphans[assignable]] = best_c[assignable]
leftover = (face_chart < 0).nonzero(as_tuple=True)[0]
if leftover.numel(): # isolated faces become singleton charts
face_chart[leftover] = K + torch.arange(leftover.numel(), device=device)
return _renumber(face_chart, device)
_, inverse = torch.unique(face_chart, sorted=True, return_inverse=True)
return inverse
# Parallel edge-collapse (PEC) chart clustering (GPU)
@ -400,12 +269,71 @@ def _build_chart_edges(
return chart_pairs, reduced_el
def _merge_small_charts(
chart_id: Tensor, face_normal: Tensor, face_area: Tensor,
face_face: Tensor, face_edge_len: Tensor,
min_faces: int, cost_cap: float,
) -> Tensor:
"""Absorb charts under min_faces faces into their lowest-cone-cost neighbor (capped at cost_cap)."""
if chart_id.numel() == 0:
return chart_id
device = chart_id.device
for _ in range(16):
N = int(chart_id.max().item()) + 1
sizes = torch.bincount(chart_id, minlength=N)
# recompute cones from scratch: area-weighted mean axis, max deviation as half-angle
axis = torch.zeros(N, 3, dtype=torch.float32, device=device)
axis.index_add_(0, chart_id, face_normal * face_area[:, None])
axis = axis / axis.norm(dim=1, keepdim=True).clamp_min(1e-12)
dev = torch.acos((face_normal * axis[chart_id]).sum(1).clamp(-1.0, 1.0))
half = torch.zeros(N, dtype=torch.float32, device=device)
half.scatter_reduce_(0, chart_id, dev, reduce="amax")
edges, _ = _build_chart_edges(face_face, chart_id, face_edge_len)
if edges.shape[0] == 0:
break
a, b = edges[:, 0], edges[:, 1]
_, new_half, _ = _combine_normal_cones(axis[a], half[a], axis[b], half[b])
ok = new_half <= cost_cap
E = edges.shape[0]
key = (torch.clamp(new_half * 1e6, max=2e9).to(torch.int64) << 32) \
| torch.arange(E, dtype=torch.long, device=device)
best = torch.full((N,), 1 << 62, dtype=torch.long, device=device)
va = (sizes[a] < min_faces) & ok
vb = (sizes[b] < min_faces) & ok
best.scatter_reduce_(0, a[va], key[va], reduce="amin")
best.scatter_reduce_(0, b[vb], key[vb], reduce="amin")
src = (best < (1 << 62)).nonzero(as_tuple=True)[0]
if src.numel() == 0:
break
eid = best[src] & 0xFFFFFFFF
ea, eb = a[eid], b[eid]
tgt = torch.where(ea == src, eb, ea)
# cycle break: keep src->tgt only if tgt merges nowhere itself or src > tgt;
# the kept graph is then a DAG, so the pointer-doubling below terminates
prop = torch.arange(N, dtype=torch.long, device=device)
prop[src] = tgt
keepm = (prop[tgt] == tgt) | (src > tgt)
remap = torch.arange(N, dtype=torch.long, device=device)
remap[src[keepm]] = tgt[keepm]
for _ in range(32):
nr = remap[remap]
if torch.equal(nr, remap):
break
remap = nr
chart_id = remap[chart_id]
_, chart_id = torch.unique(chart_id, return_inverse=True)
return chart_id
def cluster_charts_pec(
mesh: MeshData,
max_cost: float = 0.7,
max_iters: int = 1024,
min_faces: int = 8,
) -> Tensor:
"""Parallel edge-collapse clustering; returns face_chart [F]. max_cost is the per-merge cutoff (~0.7 rad ~ 40deg)."""
"""Parallel edge-collapse clustering; returns face_chart [F]. max_cost is the per-merge
cutoff (~0.7 rad ~ 40deg); charts under min_faces are then absorbed at a relaxed 2x cutoff."""
device = mesh.faces.device
F = mesh.faces.shape[0]
faces = mesh.faces.to(torch.long)
@ -471,5 +399,8 @@ def cluster_charts_pec(
remap[win_b] = win_a
chart_id = remap[chart_id]
if min_faces > 1:
chart_id = _merge_small_charts(chart_id, face_normal, mesh.face_area.to(torch.float32),
face_face, face_edge_len, min_faces, 2.0 * max_cost)
_, inverse = torch.unique(chart_id, sorted=True, return_inverse=True)
return inverse

166
comfy_extras/nodes_mesh_postprocess.py
View File

@ -15,6 +15,7 @@ from comfy_extras.mesh3d.uv_unwrap import segment as _uv_seg
from comfy_extras.mesh3d.uv_unwrap import parameterize as _uv_param
from comfy_extras.mesh3d.uv_unwrap import pack as _uv_pack
import logging
import time
from tqdm import tqdm
from scipy.sparse import csr_matrix
from scipy.sparse.csgraph import connected_components
@ -2454,6 +2455,7 @@ def _uv_weld_vertices(v, f, weld_distance):
def _uv_unwrap(positions, indices, segmenter, resolution, padding, weld_distance):
"""UV-unwrap a single mesh; returns (vmapping, indices, uvs); vmapping maps each output
vertex to an input vertex (seam verts duplicated)."""
t_start = time.perf_counter()
v_in = positions.to(torch.float32)
f_in = indices.to(torch.long).reshape(-1, 3)
v_in, f_in, welded_to_orig = _uv_weld_vertices(v_in, f_in, weld_distance)
@ -2479,72 +2481,138 @@ def _uv_unwrap(positions, indices, segmenter, resolution, padding, weld_distance
n_charts = int(face_chart.max().item()) + 1 if face_chart.numel() else 0
areas_cpu = _uv_mesh.chart_3d_areas(mesh.face_area, face_chart, n_charts).detach().cpu()
# per-chart loop on CPU/numpy to avoid per-chart GPU sync
if n_charts == 0:
return (np.empty(0, dtype=np.int64), np.zeros((0, 3), dtype=np.int64),
np.empty((0, 2), dtype=np.float32))
# vectorized chart extraction: one global sort/unique replaces per-chart unique/searchsorted
face_chart_np = face_chart.cpu().numpy()
faces_np = mesh.faces.cpu().numpy()
vertices_np = mesh.vertices.cpu().numpy()
face_face_np = mesh.face_face.cpu().numpy()
sorted_face_idx_np = np.argsort(face_chart_np, kind="stable")
order = np.argsort(face_chart_np, kind="stable")
chart_counts_np = np.bincount(face_chart_np, minlength=n_charts)
chart_offsets_np = np.empty(n_charts + 1, dtype=np.int64)
chart_offsets_np[0] = 0
chart_offsets_np = np.zeros(n_charts + 1, dtype=np.int64)
np.cumsum(chart_counts_np, out=chart_offsets_np[1:])
faces_sorted = faces_np[order]
chart_sorted = face_chart_np[order]
n_verts_in = max(vertices_np.shape[0], 1)
chart_of_slot = np.repeat(chart_sorted, 3)
uniq_keys, local_flat = np.unique(chart_of_slot * n_verts_in + faces_sorted.reshape(-1),
return_inverse=True)
used_verts_all = uniq_keys % n_verts_in # per-chart sorted unique verts, concatenated
vert_counts = np.bincount(uniq_keys // n_verts_in, minlength=n_charts)
vert_offsets = np.zeros(n_charts + 1, dtype=np.int64)
np.cumsum(vert_counts, out=vert_offsets[1:])
local_faces_all = (local_flat - vert_offsets[chart_of_slot]).reshape(-1, 3)
pos_in_chart = np.empty(order.size, dtype=np.int64)
pos_in_chart[order] = np.arange(order.size) - chart_offsets_np[chart_sorted]
ff_sorted = face_face_np[order]
ff_safe = np.maximum(ff_sorted, 0)
keep = (ff_sorted >= 0) & (face_chart_np[ff_safe] == chart_sorted[:, None])
local_ff_all = np.where(keep, pos_in_chart[ff_safe], -1)
all_chart_uvs, all_chart_3d_areas, all_chart_uv_areas, all_chart_faces = [], [], [], []
chart_records = []
# progress: n_charts units for parameterize + 2*n_charts for pack (prepare + place)
pbar = comfy.utils.ProgressBar(3 * n_charts)
# parameterize (batched): ortho-project every chart at once, batched stretch metrics
# decide acceptance, rejected charts solve ABF/LSCM in dense per-size-bucket batches
chart_of_vert = (uniq_keys // n_verts_in).astype(np.int64)
verts_concat = vertices_np[used_verts_all].astype(np.float64)
gl_faces = local_faces_all + vert_offsets[chart_sorted][:, None]
face_pos = pos_in_chart[order] # row of each (sorted) face in its chart
uv0 = _uv_param.ortho_project_concat(verts_concat, chart_of_vert, n_charts)
rms, mx, n_flip, n_zero = _uv_param.stretch_metrics_concat(
verts_concat, uv0, gl_faces, chart_sorted, n_charts)
valid_chart = (vert_counts >= 3) & (chart_counts_np > 0)
auto = valid_chart & (chart_counts_np <= 5) # tiny charts always keep ortho
flip_ok = (n_flip == 0) | (n_flip == chart_counts_np)
cand = valid_chart & ~auto & flip_ok & (n_zero == 0) & (rms <= 1.5) & (mx <= 2.0)
pbar.update(int(auto.sum()))
ortho_ok = auto.copy()
cand_ids = np.nonzero(cand)[0]
for c in tqdm(cand_ids, desc="unwrap: ortho checks", unit="chart", leave=False):
f0, f1 = chart_offsets_np[c], chart_offsets_np[c + 1]
v0, v1 = vert_offsets[c], vert_offsets[c + 1]
if not _uv_param._uv_boundary_self_intersects(
uv0[v0:v1], local_faces_all[f0:f1], local_ff_all[f0:f1]):
ortho_ok[c] = True
pbar.update(1)
lscm_mask = valid_chart & ~ortho_ok
batchable = vert_counts <= _uv_param.LSCM_BATCH_MAX_VERTS
lscm_ids = np.nonzero(lscm_mask & batchable)[0]
big_ids = np.nonzero(lscm_mask & ~batchable)[0]
lscm_uv = _uv_param.lscm_charts_batch(
verts_concat, uv0, gl_faces, face_pos, chart_sorted, chart_of_vert,
vert_offsets, lscm_ids, n_charts,
device=comfy.model_management.get_torch_device())
pbar.update(int(lscm_ids.size))
uvs_np_list: list = [None] * n_charts
uv0_f32 = uv0.astype(np.float32)
for c in tqdm(big_ids, desc="unwrap: LSCM (large charts)", unit="chart", leave=False):
f0, f1 = chart_offsets_np[c], chart_offsets_np[c + 1]
v0, v1 = vert_offsets[c], vert_offsets[c + 1]
uvs_t = _uv_param.lscm_chart(
torch.from_numpy(verts_concat[v0:v1]),
torch.from_numpy(local_faces_all[f0:f1]),
torch.from_numpy(local_ff_all[f0:f1]), pin_positions=uv0[v0:v1])
lscm_uv[int(c)] = uvs_t.detach().cpu().numpy().astype(np.float32)
pbar.update(1)
for c in range(n_charts):
gfi_np = sorted_face_idx_np[chart_offsets_np[c]:chart_offsets_np[c + 1]]
chart_faces_global = faces_np[gfi_np]
used_verts_np = np.unique(chart_faces_global)
local_faces_np = np.searchsorted(used_verts_np, chart_faces_global)
local_verts_np = vertices_np[used_verts_np]
ff_global = face_face_np[gfi_np]
ff_safe = np.maximum(ff_global, 0)
nb_chart = np.where(ff_global >= 0, face_chart_np[ff_safe], -1)
keep = (ff_global >= 0) & (nb_chart == c)
local_neighbor = np.searchsorted(gfi_np, ff_safe)
local_ff_np = np.where(keep, local_neighbor, -1)
v0, v1 = vert_offsets[c], vert_offsets[c + 1]
if ortho_ok[c]:
uvs_np_list[c] = uv0_f32[v0:v1]
continue
u = lscm_uv.get(int(c))
if u is not None and np.all(np.isfinite(u)) and u.size:
# collapsed UV island (aspect > 100:1) blows up packing scale; keep ortho instead
bbox = u.max(axis=0) - u.min(axis=0)
if max(float(bbox.max()), 1e-12) / max(float(bbox.min()), 1e-12) <= 100.0:
uvs_np_list[c] = u
continue
uvs_np_list[c] = (uv0_f32[v0:v1] if valid_chart[c]
else np.zeros((v1 - v0, 2), dtype=np.float32))
lf = torch.from_numpy(local_faces_np)
uvs = _uv_param.parametrize_chart(
torch.from_numpy(local_verts_np), lf, torch.from_numpy(local_ff_np))
ua, ub, uc = uvs[lf[:, 0]], uvs[lf[:, 1]], uvs[lf[:, 2]]
uv_area_sum = float(0.5 * (
(ub[:, 0] - ua[:, 0]) * (uc[:, 1] - ua[:, 1])
- (uc[:, 0] - ua[:, 0]) * (ub[:, 1] - ua[:, 1])).abs().sum().item())
chart_records.append({"local_faces": lf, "vmap": torch.from_numpy(used_verts_np),
"global_face_idx": torch.from_numpy(gfi_np)})
all_chart_uvs.append(uvs)
all_chart_3d_areas.append(float(areas_cpu[c].item()))
all_chart_uv_areas.append(uv_area_sum)
all_chart_faces.append(lf)
# per-chart UV areas in one pass over all faces
uvs_all_np = np.concatenate(uvs_np_list)
ua, ub, uc = uvs_all_np[gl_faces[:, 0]], uvs_all_np[gl_faces[:, 1]], uvs_all_np[gl_faces[:, 2]]
tri_uv_area = 0.5 * np.abs(
(ub[:, 0] - ua[:, 0]) * (uc[:, 1] - ua[:, 1])
- (uc[:, 0] - ua[:, 0]) * (ub[:, 1] - ua[:, 1]))
uv_area_np = np.bincount(chart_sorted, weights=tri_uv_area.astype(np.float64),
minlength=n_charts)
areas_3d_np = areas_cpu.numpy().astype(np.float64)
# auto-tune texel density toward `resolution` (~0.62 pack fill)
total_3d_area = sum(all_chart_3d_areas) or 1.0
total_3d_area = float(areas_3d_np.sum()) or 1.0
target_dim = float(resolution) if resolution > 0 else 1024.0
tex_per_unit = math.sqrt((target_dim * target_dim) * 0.62 / total_3d_area)
placements, atlas_w, atlas_h = _uv_pack.pack_bitmap(
all_chart_uvs, all_chart_3d_areas, all_chart_uv_areas, all_chart_faces,
texels_per_unit=tex_per_unit, padding_texels=padding)
placed = _uv_pack.apply_placements(all_chart_uvs, placements, atlas_w, atlas_h)
with tqdm(total=2 * n_charts, desc="unwrap: pack", unit="chart", leave=False) as tq_pack:
def _pack_progress(done, total):
tq_pack.update(done - tq_pack.n)
pbar.update_absolute(n_charts + done, 3 * n_charts)
p_x, p_y, p_sw, p_th, p_sc, p_chh, atlas_w, atlas_h = _uv_pack.pack_bitmap_concat(
uvs_all_np, vert_offsets, local_faces_all, chart_offsets_np,
areas_3d_np, uv_area_np,
texels_per_unit=tex_per_unit, padding_texels=padding,
progress_callback=_pack_progress)
pbar.update_absolute(3 * n_charts, 3 * n_charts)
# assembly: output verts are the per-chart used-vert lists concatenated in chart order,
# so vert_offsets doubles as the output vertex cursor
n_in_faces = mesh.faces.shape[0]
out_indices = np.zeros((n_in_faces, 3), dtype=np.int64)
out_uvs_list, out_vmap_list, v_cursor = [], [], 0
for c, rec in enumerate(chart_records):
vmap_np = rec["vmap"].cpu().numpy()
local_faces_np = rec["local_faces"].cpu().numpy()
global_face_idx = rec["global_face_idx"].cpu().numpy()
out_uvs_list.append(placed[c].cpu().numpy())
if welded_to_orig is not None:
vmap_np = welded_to_orig[vmap_np]
out_vmap_list.append(vmap_np)
out_indices[global_face_idx] = local_faces_np + v_cursor
v_cursor += vmap_np.shape[0]
vmapping_out = np.concatenate(out_vmap_list) if out_vmap_list else np.empty(0, dtype=np.int64)
uvs_out = np.concatenate(out_uvs_list) if out_uvs_list else np.empty((0, 2), dtype=np.float32)
out_indices[order] = gl_faces
vmapping_out = used_verts_all if welded_to_orig is None else welded_to_orig[used_verts_all]
uvs_out = _uv_pack.apply_placements_concat(
uvs_all_np, vert_offsets, p_x, p_y, p_sw, p_th, p_sc, p_chh, atlas_w, atlas_h)
logging.info(f"[uv_unwrap] {mesh.faces.shape[0]} faces -> {n_charts} charts, "
f"atlas {atlas_w}x{atlas_h}, {time.perf_counter() - t_start:.1f}s")
return vmapping_out, out_indices, uvs_out