mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-07 23:21:04 +08:00
566 lines
23 KiB
Python
566 lines
23 KiB
Python
"""Chart parameterization: ortho PCA projection, falling back to ABF/LSCM."""
|
|
from __future__ import annotations
|
|
|
|
import warnings
|
|
from typing import List, Tuple
|
|
|
|
import numpy as np
|
|
import scipy.sparse as sp
|
|
import scipy.sparse.linalg as spla
|
|
import torch
|
|
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."""
|
|
At = A.T.tocsr()
|
|
AtA = (At @ A).tocsc()
|
|
Atb = At @ b
|
|
return spla.spsolve(AtA, Atb)
|
|
|
|
|
|
def _triangle_local_2d(verts_3d: np.ndarray, faces: np.ndarray) -> np.ndarray:
|
|
"""Per-triangle 2D coords [F, 3, 2] with v0 at origin, v1 along +x."""
|
|
v0 = verts_3d[faces[:, 0]]
|
|
v1 = verts_3d[faces[:, 1]]
|
|
v2 = verts_3d[faces[:, 2]]
|
|
e01 = v1 - v0
|
|
e02 = v2 - v0
|
|
L01 = np.linalg.norm(e01, axis=1).clip(min=1e-20)
|
|
x_axis = e01 / L01[:, None]
|
|
n = np.cross(e01, e02)
|
|
n /= np.linalg.norm(n, axis=1, keepdims=True).clip(min=1e-20)
|
|
y_axis = np.cross(n, x_axis)
|
|
|
|
out = np.zeros((faces.shape[0], 3, 2), dtype=np.float64)
|
|
out[:, 1, 0] = L01
|
|
out[:, 2, 0] = (e02 * x_axis).sum(axis=1)
|
|
out[:, 2, 1] = (e02 * y_axis).sum(axis=1)
|
|
return out
|
|
|
|
|
|
def _pick_pins(loops: List[List[int]], verts_3d: np.ndarray) -> Tuple[int, int]:
|
|
"""Pick the longest-diameter axis-extremal boundary vertex pair across all boundary verts."""
|
|
if not loops:
|
|
# Closed surface: two far verts via two-pass farthest.
|
|
d2 = np.sum((verts_3d - verts_3d[0]) ** 2, axis=1)
|
|
a = int(np.argmax(d2))
|
|
d2 = np.sum((verts_3d - verts_3d[a]) ** 2, axis=1)
|
|
b = int(np.argmax(d2))
|
|
return a, b
|
|
boundary_verts: List[int] = []
|
|
for loop in loops:
|
|
boundary_verts.extend(loop)
|
|
seen = set()
|
|
uniq = []
|
|
for v in boundary_verts:
|
|
if v not in seen:
|
|
seen.add(v)
|
|
uniq.append(v)
|
|
bv = np.asarray(uniq, dtype=np.int64)
|
|
pts = verts_3d[bv]
|
|
pin_pairs = []
|
|
for axis in range(3):
|
|
i_min = int(bv[int(np.argmin(pts[:, axis]))])
|
|
i_max = int(bv[int(np.argmax(pts[:, axis]))])
|
|
d = float(np.linalg.norm(verts_3d[i_min] - verts_3d[i_max]))
|
|
pin_pairs.append((d, i_min, i_max))
|
|
d0, _, _ = pin_pairs[0]
|
|
d1, _, _ = pin_pairs[1]
|
|
d2, _, _ = pin_pairs[2]
|
|
if d0 > d1 and d0 > d2:
|
|
_, a, b = pin_pairs[0]
|
|
elif d1 > d2:
|
|
_, a, b = pin_pairs[1]
|
|
else:
|
|
_, a, b = pin_pairs[2]
|
|
return a, b
|
|
|
|
|
|
def _ortho_project(verts_3d: np.ndarray) -> np.ndarray:
|
|
"""PCA-fit plane normal, axis-aligned tangent, project verts to 2D."""
|
|
centroid = verts_3d.mean(axis=0)
|
|
pts = verts_3d - centroid
|
|
cov = pts.T @ pts
|
|
_w, ev = np.linalg.eigh(cov)
|
|
normal = ev[:, 0]
|
|
a = np.abs(normal)
|
|
if a[0] < a[1] and a[0] < a[2]:
|
|
t = np.array([1.0, 0.0, 0.0])
|
|
elif a[1] < a[2]:
|
|
t = np.array([0.0, 1.0, 0.0])
|
|
else:
|
|
t = np.array([0.0, 0.0, 1.0])
|
|
t = t - normal * float(np.dot(normal, t))
|
|
t /= max(float(np.linalg.norm(t)), 1e-20)
|
|
b = np.cross(normal, t)
|
|
return np.stack([verts_3d @ t, verts_3d @ b], axis=1)
|
|
|
|
|
|
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]
|
|
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_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]
|
|
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
|
|
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(
|
|
uvs: np.ndarray, faces: np.ndarray, face_face: np.ndarray, eps: float = 1e-9
|
|
) -> bool:
|
|
"""True if any chart-boundary edge pair crosses in 2D (ortho folded the chart)."""
|
|
fi, ei = np.nonzero(face_face < 0)
|
|
n = fi.size
|
|
if n < 2:
|
|
return False
|
|
a = uvs[faces[fi, ei]].astype(np.float64)
|
|
b = uvs[faces[fi, (ei + 1) % 3]].astype(np.float64)
|
|
d = b - a
|
|
# Pairwise segment crossings, row-chunked to bound memory at chunk*n.
|
|
chunk = max(1, min(n, 1_000_000 // max(n, 1)))
|
|
for s in range(0, n, chunk):
|
|
e = min(s + chunk, n)
|
|
d1 = d[s:e, None, :]
|
|
denom = d1[:, :, 0] * d[None, :, 1] - d1[:, :, 1] * d[None, :, 0]
|
|
rx = a[None, :, 0] - a[s:e, None, 0]
|
|
ry = a[None, :, 1] - a[s:e, None, 1]
|
|
with np.errstate(divide="ignore", invalid="ignore"):
|
|
t = (rx * d[None, :, 1] - ry * d[None, :, 0]) / denom
|
|
u = (rx * d1[:, :, 1] - ry * d1[:, :, 0]) / denom
|
|
cross = (
|
|
(np.abs(denom) >= eps)
|
|
& (t > eps) & (t < 1.0 - eps)
|
|
& (u > eps) & (u < 1.0 - eps)
|
|
)
|
|
if bool(cross.any()):
|
|
return True
|
|
return False
|
|
|
|
|
|
def _abf_face_coefficients(
|
|
verts_3d: np.ndarray, faces: np.ndarray
|
|
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
|
|
"""Per-face ABF constraint (largest-sine vertex at local index 2); returns (faces_reordered, cosine, sine, valid_mask) with valid_mask False for degenerate tris."""
|
|
p0 = verts_3d[faces[:, 0]]
|
|
p1 = verts_3d[faces[:, 1]]
|
|
p2 = verts_3d[faces[:, 2]]
|
|
e01 = p1 - p0
|
|
e12 = p2 - p1
|
|
e20 = p0 - p2
|
|
L01 = np.linalg.norm(e01, axis=1).clip(min=1e-20)
|
|
L12 = np.linalg.norm(e12, axis=1).clip(min=1e-20)
|
|
L20 = np.linalg.norm(e20, axis=1).clip(min=1e-20)
|
|
cos_a0 = ((-e20) * e01).sum(axis=1) / (L20 * L01)
|
|
cos_a1 = ((-e01) * e12).sum(axis=1) / (L01 * L12)
|
|
cos_a2 = ((-e12) * e20).sum(axis=1) / (L12 * L20)
|
|
cos_a0 = cos_a0.clip(-1.0, 1.0)
|
|
cos_a1 = cos_a1.clip(-1.0, 1.0)
|
|
cos_a2 = cos_a2.clip(-1.0, 1.0)
|
|
a = np.arccos(cos_a0)
|
|
b_ang = np.arccos(cos_a1)
|
|
c_ang = np.arccos(cos_a2)
|
|
angles = np.stack([a, b_ang, c_ang], axis=1)
|
|
sines = np.stack([np.sin(a), np.sin(b_ang), np.sin(c_ang)], axis=1)
|
|
valid = (angles > 1e-12).all(axis=1)
|
|
ids = faces.astype(np.int64).copy()
|
|
|
|
s0, s1, s2 = sines[:, 0], sines[:, 1], sines[:, 2]
|
|
pattA = (s1 > s0) & (s1 > s2)
|
|
pattB = (~pattA) & (s0 > s1) & (s0 > s2)
|
|
|
|
if pattA.any():
|
|
old_a = angles[pattA].copy()
|
|
old_s = sines[pattA].copy()
|
|
old_id = ids[pattA].copy()
|
|
angles[pattA] = old_a[:, [2, 0, 1]]
|
|
sines[pattA] = old_s[:, [2, 0, 1]]
|
|
ids[pattA] = old_id[:, [2, 0, 1]]
|
|
if pattB.any():
|
|
old_a = angles[pattB].copy()
|
|
old_s = sines[pattB].copy()
|
|
old_id = ids[pattB].copy()
|
|
angles[pattB] = old_a[:, [1, 2, 0]]
|
|
sines[pattB] = old_s[:, [1, 2, 0]]
|
|
ids[pattB] = old_id[:, [1, 2, 0]]
|
|
|
|
a0 = angles[:, 0]
|
|
s0 = sines[:, 0]
|
|
s1 = sines[:, 1]
|
|
s2 = sines[:, 2]
|
|
c0 = np.cos(a0)
|
|
ratio = np.where(s2 > 0.0, s1 / s2.clip(min=1e-20), 1.0)
|
|
cosine = c0 * ratio
|
|
sine = s0 * ratio
|
|
return ids, cosine, sine, valid
|
|
|
|
|
|
def lscm_chart(
|
|
local_verts: Tensor,
|
|
local_faces: Tensor,
|
|
local_face_face: Tensor,
|
|
pin_positions: "np.ndarray | None" = None,
|
|
) -> Tensor:
|
|
"""ABF parameterization on one chart (degenerate faces use plain LSCM rows; two pins fix gauge at pin_positions)."""
|
|
verts_np = local_verts.detach().cpu().numpy().astype(np.float64)
|
|
faces_np = local_faces.detach().cpu().numpy().astype(np.int64)
|
|
Vc = verts_np.shape[0]
|
|
Fc = faces_np.shape[0]
|
|
|
|
if Vc < 3 or Fc == 0:
|
|
return torch.zeros((Vc, 2), dtype=torch.float32, device=local_verts.device)
|
|
|
|
loops = _mesh.chart_boundary_loops(local_faces, local_face_face)
|
|
pin_a, pin_b = _pick_pins(loops, verts_np)
|
|
|
|
if pin_positions is not None and pin_positions.shape == (Vc, 2):
|
|
pa = pin_positions[pin_a]
|
|
pb = pin_positions[pin_b]
|
|
u_a, v_a = float(pa[0]), float(pa[1])
|
|
u_b, v_b = float(pb[0]), float(pb[1])
|
|
else:
|
|
u_a, v_a = 0.0, 0.0
|
|
u_b, v_b = 1.0, 0.0
|
|
|
|
abf_ids, abf_cos, abf_sin, abf_valid = _abf_face_coefficients(verts_np, faces_np)
|
|
|
|
rows_list: List[np.ndarray] = []
|
|
cols_list: List[np.ndarray] = []
|
|
vals_list: List[np.ndarray] = []
|
|
|
|
# ABF rows for valid faces.
|
|
valid_idx = np.nonzero(abf_valid)[0]
|
|
if valid_idx.size:
|
|
Nv = valid_idx.size
|
|
id0 = abf_ids[valid_idx, 0]
|
|
id1 = abf_ids[valid_idx, 1]
|
|
id2 = abf_ids[valid_idx, 2]
|
|
cosf = abf_cos[valid_idx]
|
|
sinf = abf_sin[valid_idx]
|
|
r_real = valid_idx * 2
|
|
r_imag = valid_idx * 2 + 1
|
|
ones = np.ones(Nv, dtype=np.float64)
|
|
rows_list.extend([r_real] * 5)
|
|
cols_list.extend([id0, id0 + Vc, id1, id1 + Vc, id2])
|
|
vals_list.extend([cosf - 1.0, -sinf, -cosf, sinf, ones])
|
|
rows_list.extend([r_imag] * 5)
|
|
cols_list.extend([id0, id0 + Vc, id1, id1 + Vc, id2 + Vc])
|
|
vals_list.extend([sinf, cosf - 1.0, -sinf, -cosf, ones])
|
|
|
|
# Plain-LSCM rows for invalid (degenerate) faces.
|
|
invalid_idx = np.nonzero(~abf_valid)[0]
|
|
if invalid_idx.size:
|
|
tri2d_inv = _triangle_local_2d(verts_np, faces_np[invalid_idx])
|
|
twice_area_inv = (
|
|
tri2d_inv[:, 1, 0] * tri2d_inv[:, 2, 1]
|
|
- tri2d_inv[:, 1, 1] * tri2d_inv[:, 2, 0]
|
|
)
|
|
weight_inv = 1.0 / np.sqrt(2.0 * np.abs(twice_area_inv).clip(min=1e-20))
|
|
r_real_inv = invalid_idx * 2
|
|
r_imag_inv = invalid_idx * 2 + 1
|
|
for j in range(3):
|
|
jp1 = (j + 1) % 3
|
|
jp2 = (j + 2) % 3
|
|
a_j = (tri2d_inv[:, jp1, 0] - tri2d_inv[:, jp2, 0]) * weight_inv
|
|
b_j = (tri2d_inv[:, jp1, 1] - tri2d_inv[:, jp2, 1]) * weight_inv
|
|
v_idx = faces_np[invalid_idx, j]
|
|
rows_list.extend([r_real_inv, r_real_inv, r_imag_inv, r_imag_inv])
|
|
cols_list.extend([v_idx, v_idx + Vc, v_idx, v_idx + Vc])
|
|
vals_list.extend([a_j, -b_j, b_j, a_j])
|
|
|
|
rows = np.concatenate(rows_list) if rows_list else np.empty(0, dtype=np.int64)
|
|
cols = np.concatenate(cols_list) if cols_list else np.empty(0, dtype=np.int64)
|
|
vals = np.concatenate(vals_list) if vals_list else np.empty(0, dtype=np.float64)
|
|
|
|
A_full = sp.csr_matrix((vals, (rows, cols)), shape=(2 * Fc, 2 * Vc))
|
|
|
|
pin_cols = np.array([pin_a, pin_b, pin_a + Vc, pin_b + Vc], dtype=np.int64)
|
|
pin_vals = np.array([u_a, u_b, v_a, v_b], dtype=np.float64)
|
|
|
|
free_mask = np.ones(2 * Vc, dtype=bool)
|
|
free_mask[pin_cols] = False
|
|
free_cols = np.nonzero(free_mask)[0]
|
|
|
|
A_pinned = A_full[:, pin_cols]
|
|
A_free = A_full[:, free_cols]
|
|
b = -(A_pinned @ pin_vals)
|
|
|
|
# Singular system (under-constrained chart) falls back to ortho.
|
|
fallback_to_ortho = False
|
|
try:
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("error", category=sp.linalg.MatrixRankWarning)
|
|
x_free = solve_least_squares(A_free, b)
|
|
if not np.all(np.isfinite(x_free)):
|
|
fallback_to_ortho = True
|
|
except (sp.linalg.MatrixRankWarning, RuntimeError):
|
|
fallback_to_ortho = True # singular / under-constrained system
|
|
|
|
if fallback_to_ortho:
|
|
if pin_positions is not None and pin_positions.shape == (Vc, 2):
|
|
uvs = pin_positions.astype(np.float32)
|
|
else:
|
|
uvs = _ortho_project(verts_np).astype(np.float32)
|
|
return torch.from_numpy(uvs).to(local_verts.device)
|
|
|
|
full = np.zeros(2 * Vc, dtype=np.float64)
|
|
full[free_cols] = x_free
|
|
full[pin_cols] = pin_vals
|
|
uvs = np.stack([full[:Vc], full[Vc:]], axis=1).astype(np.float32)
|
|
if not np.all(np.isfinite(uvs)):
|
|
if pin_positions is not None and pin_positions.shape == (Vc, 2):
|
|
uvs = pin_positions.astype(np.float32)
|
|
else:
|
|
uvs = _ortho_project(verts_np).astype(np.float32)
|
|
|
|
return torch.from_numpy(uvs).to(local_verts.device)
|