402 lines
16 KiB
Python
402 lines
16 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
from typing import *
|
|
from ._helpers import batched
|
|
|
|
|
|
__all__ = [
|
|
'triangulate',
|
|
'compute_face_normal',
|
|
'compute_face_angles',
|
|
'compute_vertex_normal',
|
|
'compute_vertex_normal_weighted',
|
|
'remove_unreferenced_vertices',
|
|
'remove_corrupted_faces',
|
|
'merge_duplicate_vertices',
|
|
'subdivide_mesh_simple',
|
|
'compute_face_tbn',
|
|
'compute_vertex_tbn',
|
|
'laplacian',
|
|
'laplacian_smooth_mesh',
|
|
'taubin_smooth_mesh',
|
|
'laplacian_hc_smooth_mesh',
|
|
]
|
|
|
|
|
|
def triangulate(
|
|
faces: torch.Tensor,
|
|
vertices: torch.Tensor = None,
|
|
backslash: bool = None
|
|
) -> torch.Tensor:
|
|
"""
|
|
Triangulate a polygonal mesh.
|
|
|
|
Args:
|
|
faces (torch.Tensor): [..., L, P] polygonal faces
|
|
vertices (torch.Tensor, optional): [..., N, 3] 3-dimensional vertices.
|
|
If given, the triangulation is performed according to the distance
|
|
between vertices. Defaults to None.
|
|
backslash (torch.Tensor, optional): [..., L] boolean array indicating
|
|
how to triangulate the quad faces. Defaults to None.
|
|
|
|
|
|
Returns:
|
|
(torch.Tensor): [L * (P - 2), 3] triangular faces
|
|
"""
|
|
if faces.shape[-1] == 3:
|
|
return faces
|
|
P = faces.shape[-1]
|
|
if vertices is not None:
|
|
assert faces.shape[-1] == 4, "now only support quad mesh"
|
|
if backslash is None:
|
|
faces_idx = faces.long()
|
|
backslash = torch.norm(vertices[faces_idx[..., 0]] - vertices[faces_idx[..., 2]], p=2, dim=-1) < \
|
|
torch.norm(vertices[faces_idx[..., 1]] - vertices[faces_idx[..., 3]], p=2, dim=-1)
|
|
if backslash is None:
|
|
loop_indice = torch.stack([
|
|
torch.zeros(P - 2, dtype=int),
|
|
torch.arange(1, P - 1, 1, dtype=int),
|
|
torch.arange(2, P, 1, dtype=int)
|
|
], axis=1)
|
|
return faces[:, loop_indice].reshape(-1, 3)
|
|
else:
|
|
assert faces.shape[-1] == 4, "now only support quad mesh"
|
|
if isinstance(backslash, bool):
|
|
if backslash:
|
|
faces = faces[:, [0, 1, 2, 0, 2, 3]].reshape(-1, 3)
|
|
else:
|
|
faces = faces[:, [0, 1, 3, 3, 1, 2]].reshape(-1, 3)
|
|
else:
|
|
faces = torch.where(
|
|
backslash[:, None],
|
|
faces[:, [0, 1, 2, 0, 2, 3]],
|
|
faces[:, [0, 1, 3, 3, 1, 2]]
|
|
).reshape(-1, 3)
|
|
return faces
|
|
|
|
|
|
@batched(2, None)
|
|
def compute_face_normal(
|
|
vertices: torch.Tensor,
|
|
faces: torch.Tensor
|
|
) -> torch.Tensor:
|
|
"""
|
|
Compute face normals of a triangular mesh
|
|
|
|
Args:
|
|
vertices (torch.Tensor): [..., N, 3] 3-dimensional vertices
|
|
faces (torch.Tensor): [..., T, 3] triangular face indices
|
|
|
|
Returns:
|
|
normals (torch.Tensor): [..., T, 3] face normals
|
|
"""
|
|
N = vertices.shape[0]
|
|
index = torch.arange(N)[:, None]
|
|
normal = torch.cross(
|
|
vertices[index, faces[..., 1].long()] - vertices[index, faces[..., 0].long()],
|
|
vertices[index, faces[..., 2].long()] - vertices[index, faces[..., 0].long()],
|
|
dim=-1
|
|
)
|
|
return F.normalize(normal, p=2, dim=-1)
|
|
|
|
|
|
@batched(2, None)
|
|
def compute_face_angles(
|
|
vertices: torch.Tensor,
|
|
faces: torch.Tensor
|
|
) -> torch.Tensor:
|
|
"""
|
|
Compute face angles of a triangular mesh
|
|
|
|
Args:
|
|
vertices (torch.Tensor): [..., N, 3] 3-dimensional vertices
|
|
faces (torch.Tensor): [T, 3] triangular face indices
|
|
|
|
Returns:
|
|
angles (torch.Tensor): [..., T, 3] face angles
|
|
"""
|
|
face_angles = []
|
|
for i in range(3):
|
|
edge1 = torch.index_select(vertices, dim=-2, index=faces[:, (i + 1) % 3]) - torch.index_select(vertices, dim=-2, index=faces[:, i])
|
|
edge2 = torch.index_select(vertices, dim=-2, index=faces[:, (i + 2) % 3]) - torch.index_select(vertices, dim=-2, index=faces[:, i])
|
|
face_angle = torch.arccos(torch.sum(F.normalize(edge1, p=2, dim=-1) * F.normalize(edge2, p=2, dim=-1), dim=-1))
|
|
face_angles.append(face_angle)
|
|
face_angles = torch.stack(face_angles, dim=-1)
|
|
return face_angles
|
|
|
|
|
|
@batched(2, None, 2)
|
|
def compute_vertex_normal(
|
|
vertices: torch.Tensor,
|
|
faces: torch.Tensor,
|
|
face_normal: torch.Tensor = None
|
|
) -> torch.Tensor:
|
|
"""
|
|
Compute vertex normals of a triangular mesh by averaging neightboring face normals
|
|
|
|
Args:
|
|
vertices (torch.Tensor): [..., N, 3] 3-dimensional vertices
|
|
faces (torch.Tensor): [T, 3] triangular face indices
|
|
face_normal (torch.Tensor, optional): [..., T, 3] face normals.
|
|
None to compute face normals from vertices and faces. Defaults to None.
|
|
|
|
Returns:
|
|
normals (torch.Tensor): [..., N, 3] vertex normals
|
|
"""
|
|
N = vertices.shape[0]
|
|
assert faces.shape[-1] == 3, "Only support triangular mesh"
|
|
if face_normal is None:
|
|
face_normal = compute_face_normal(vertices, faces)
|
|
face_normal = face_normal[:, :, None, :].expand(-1, -1, 3, -1).flatten(-3, -2)
|
|
faces = faces.flatten()
|
|
vertex_normal = torch.index_put(torch.zeros_like(vertices), (torch.arange(N)[:, None], faces[None, :]), face_normal, accumulate=True)
|
|
vertex_normal = F.normalize(vertex_normal, p=2, dim=-1)
|
|
return vertex_normal
|
|
|
|
|
|
@batched(2, None, 2)
|
|
def compute_vertex_normal_weighted(
|
|
vertices: torch.Tensor,
|
|
faces: torch.Tensor,
|
|
face_normal: torch.Tensor = None
|
|
) -> torch.Tensor:
|
|
"""
|
|
Compute vertex normals of a triangular mesh by weighted sum of neightboring face normals
|
|
according to the angles
|
|
|
|
Args:
|
|
vertices (torch.Tensor): [..., N, 3] 3-dimensional vertices
|
|
faces (torch.Tensor): [T, 3] triangular face indices
|
|
face_normal (torch.Tensor, optional): [..., T, 3] face normals.
|
|
None to compute face normals from vertices and faces. Defaults to None.
|
|
|
|
Returns:
|
|
normals (torch.Tensor): [..., N, 3] vertex normals
|
|
"""
|
|
N = vertices.shape[0]
|
|
if face_normal is None:
|
|
face_normal = compute_face_normal(vertices, faces)
|
|
face_angle = compute_face_angles(vertices, faces)
|
|
face_normal = face_normal[:, :, None, :].expand(-1, -1, 3, -1) * face_angle[..., None]
|
|
vertex_normal = torch.index_put(torch.zeros_like(vertices), (torch.arange(N)[:, None], faces.view(N, -1)), face_normal.view(N, -1, 3), accumulate=True)
|
|
vertex_normal = F.normalize(vertex_normal, p=2, dim=-1)
|
|
return vertex_normal
|
|
|
|
|
|
def remove_unreferenced_vertices(
|
|
faces: torch.Tensor,
|
|
*vertice_attrs,
|
|
return_indices: bool = False
|
|
) -> Tuple[torch.Tensor, ...]:
|
|
"""
|
|
Remove unreferenced vertices of a mesh.
|
|
Unreferenced vertices are removed, and the face indices are updated accordingly.
|
|
|
|
Args:
|
|
faces (torch.Tensor): [T, P] face indices
|
|
*vertice_attrs: vertex attributes
|
|
|
|
Returns:
|
|
faces (torch.Tensor): [T, P] face indices
|
|
*vertice_attrs: vertex attributes
|
|
indices (torch.Tensor, optional): [N] indices of vertices that are kept. Defaults to None.
|
|
"""
|
|
P = faces.shape[-1]
|
|
fewer_indices, inv_map = torch.unique(faces, return_inverse=True)
|
|
faces = inv_map.to(torch.int32).reshape(-1, P)
|
|
ret = [faces]
|
|
for attr in vertice_attrs:
|
|
ret.append(attr[fewer_indices])
|
|
if return_indices:
|
|
ret.append(fewer_indices)
|
|
return tuple(ret)
|
|
|
|
|
|
def remove_corrupted_faces(
|
|
faces: torch.Tensor
|
|
) -> torch.Tensor:
|
|
"""
|
|
Remove corrupted faces (faces with duplicated vertices)
|
|
|
|
Args:
|
|
faces (torch.Tensor): [T, 3] triangular face indices
|
|
|
|
Returns:
|
|
torch.Tensor: [T_, 3] triangular face indices
|
|
"""
|
|
corrupted = (faces[:, 0] == faces[:, 1]) | (faces[:, 1] == faces[:, 2]) | (faces[:, 2] == faces[:, 0])
|
|
return faces[~corrupted]
|
|
|
|
|
|
def merge_duplicate_vertices(
|
|
vertices: torch.Tensor,
|
|
faces: torch.Tensor,
|
|
tol: float = 1e-6
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Merge duplicate vertices of a triangular mesh.
|
|
Duplicate vertices are merged by selecte one of them, and the face indices are updated accordingly.
|
|
|
|
Args:
|
|
vertices (torch.Tensor): [N, 3] 3-dimensional vertices
|
|
faces (torch.Tensor): [T, 3] triangular face indices
|
|
tol (float, optional): tolerance for merging. Defaults to 1e-6.
|
|
|
|
Returns:
|
|
vertices (torch.Tensor): [N_, 3] 3-dimensional vertices
|
|
faces (torch.Tensor): [T, 3] triangular face indices
|
|
"""
|
|
vertices_round = torch.round(vertices / tol)
|
|
uni, uni_inv = torch.unique(vertices_round, dim=0, return_inverse=True)
|
|
uni[uni_inv] = vertices
|
|
faces = uni_inv[faces]
|
|
return uni, faces
|
|
|
|
|
|
def subdivide_mesh_simple(vertices: torch.Tensor, faces: torch.Tensor, n: int = 1) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Subdivide a triangular mesh by splitting each triangle into 4 smaller triangles.
|
|
NOTE: All original vertices are kept, and new vertices are appended to the end of the vertex list.
|
|
|
|
Args:
|
|
vertices (torch.Tensor): [N, 3] 3-dimensional vertices
|
|
faces (torch.Tensor): [T, 3] triangular face indices
|
|
n (int, optional): number of subdivisions. Defaults to 1.
|
|
|
|
Returns:
|
|
vertices (torch.Tensor): [N_, 3] subdivided 3-dimensional vertices
|
|
faces (torch.Tensor): [4 * T, 3] subdivided triangular face indices
|
|
"""
|
|
for _ in range(n):
|
|
edges = torch.stack([faces[:, [0, 1]], faces[:, [1, 2]], faces[:, [2, 0]]], dim=0)
|
|
edges = torch.sort(edges, dim=2)
|
|
uni_edges, uni_inv = torch.unique(edges, return_inverse=True, dim=0)
|
|
midpoints = (vertices[uni_edges[:, 0]] + vertices[uni_edges[:, 1]]) / 2
|
|
|
|
n_vertices = vertices.shape[0]
|
|
vertices = torch.cat([vertices, midpoints], dim=0)
|
|
faces = torch.cat([
|
|
torch.stack([faces[:, 0], n_vertices + uni_inv[0], n_vertices + uni_inv[2]], axis=1),
|
|
torch.stack([faces[:, 1], n_vertices + uni_inv[1], n_vertices + uni_inv[0]], axis=1),
|
|
torch.stack([faces[:, 2], n_vertices + uni_inv[2], n_vertices + uni_inv[1]], axis=1),
|
|
torch.stack([n_vertices + uni_inv[0], n_vertices + uni_inv[1], n_vertices + uni_inv[2]], axis=1),
|
|
], dim=0)
|
|
return vertices, faces
|
|
|
|
|
|
def compute_face_tbn(pos: torch.Tensor, faces_pos: torch.Tensor, uv: torch.Tensor, faces_uv: torch.Tensor, eps: float = 1e-7) -> torch.Tensor:
|
|
"""compute TBN matrix for each face
|
|
|
|
Args:
|
|
pos (torch.Tensor): shape (..., N_pos, 3), positions
|
|
faces_pos (torch.Tensor): shape(T, 3)
|
|
uv (torch.Tensor): shape (..., N_uv, 3) uv coordinates,
|
|
faces_uv (torch.Tensor): shape(T, 3)
|
|
|
|
Returns:
|
|
torch.Tensor: (..., T, 3, 3) TBN matrix for each face. Note TBN vectors are normalized but not necessarily orthognal
|
|
"""
|
|
e01 = torch.index_select(pos, dim=-2, index=faces_pos[:, 1]) - torch.index_select(pos, dim=-2, index=faces_pos[:, 0])
|
|
e02 = torch.index_select(pos, dim=-2, index=faces_pos[:, 2]) - torch.index_select(pos, dim=-2, index=faces_pos[:, 0])
|
|
uv01 = torch.index_select(uv, dim=-2, index=faces_uv[:, 1]) - torch.index_select(uv, dim=-2, index=faces_uv[:, 0])
|
|
uv02 = torch.index_select(uv, dim=-2, index=faces_uv[:, 2]) - torch.index_select(uv, dim=-2, index=faces_uv[:, 0])
|
|
normal = torch.cross(e01, e02)
|
|
tangent_bitangent = torch.stack([e01, e02], dim=-1) @ torch.inverse(torch.stack([uv01, uv02], dim=-1))
|
|
tbn = torch.cat([tangent_bitangent, normal.unsqueeze(-1)], dim=-1)
|
|
tbn = tbn / (torch.norm(tbn, p=2, dim=-2, keepdim=True) + eps)
|
|
return tbn
|
|
|
|
|
|
def compute_vertex_tbn(faces_topo: torch.Tensor, pos: torch.Tensor, faces_pos: torch.Tensor, uv: torch.Tensor, faces_uv: torch.Tensor) -> torch.Tensor:
|
|
"""compute TBN matrix for each face
|
|
|
|
Args:
|
|
faces_topo (torch.Tensor): (T, 3), face indice of topology
|
|
pos (torch.Tensor): shape (..., N_pos, 3), positions
|
|
faces_pos (torch.Tensor): shape(T, 3)
|
|
uv (torch.Tensor): shape (..., N_uv, 3) uv coordinates,
|
|
faces_uv (torch.Tensor): shape(T, 3)
|
|
|
|
Returns:
|
|
torch.Tensor: (..., V, 3, 3) TBN matrix for each face. Note TBN vectors are normalized but not necessarily orthognal
|
|
"""
|
|
n_vertices = faces_topo.max().item() + 1
|
|
n_tri = faces_topo.shape[-2]
|
|
batch_shape = pos.shape[:-2]
|
|
face_tbn = compute_face_tbn(pos, faces_pos, uv, faces_uv) # (..., T, 3, 3)
|
|
face_tbn = face_tbn[..., :, None, :, :].repeat(*[1] * len(batch_shape), 1, 3, 1, 1).view(*batch_shape, n_tri * 3, 3, 3) # (..., T * 3, 3, 3)
|
|
vertex_tbn = torch.index_add(torch.zeros(*batch_shape, n_vertices, 3, 3).to(face_tbn), dim=-3, index=faces_topo.view(-1), source=face_tbn)
|
|
vertex_tbn = vertex_tbn / (torch.norm(vertex_tbn, p=2, dim=-2, keepdim=True) + 1e-7)
|
|
return vertex_tbn
|
|
|
|
|
|
def laplacian(vertices: torch.Tensor, faces: torch.Tensor, weight: str = 'uniform') -> torch.Tensor:
|
|
"""Laplacian smooth with cotangent weights
|
|
|
|
Args:
|
|
vertices (torch.Tensor): shape (..., N, 3)
|
|
faces (torch.Tensor): shape (T, 3)
|
|
weight (str): 'uniform' or 'cotangent'
|
|
"""
|
|
sum_verts = torch.zeros_like(vertices) # (..., N, 3)
|
|
sum_weights = torch.zeros(*vertices.shape[:-1]).to(vertices) # (..., N)
|
|
face_verts = torch.index_select(vertices, -2, faces.view(-1)).view(*vertices.shape[:-2], *faces.shape, vertices.shape[-1]) # (..., T, 3)
|
|
if weight == 'cotangent':
|
|
for i in range(3):
|
|
e1 = face_verts[..., (i + 1) % 3, :] - face_verts[..., i, :]
|
|
e2 = face_verts[..., (i + 2) % 3, :] - face_verts[..., i, :]
|
|
cot_angle = (e1 * e2).sum(dim=-1) / torch.cross(e1, e2, dim=-1).norm(p=2, dim=-1) # (..., T, 3)
|
|
sum_verts = torch.index_add(sum_verts, -2, faces[:, (i + 1) % 3], face_verts[..., (i + 2) % 3, :] * cot_angle[..., None])
|
|
sum_weights = torch.index_add(sum_weights, -1, faces[:, (i + 1) % 3], cot_angle)
|
|
sum_verts = torch.index_add(sum_verts, -2, faces[:, (i + 2) % 3], face_verts[..., (i + 1) % 3, :] * cot_angle[..., None])
|
|
sum_weights = torch.index_add(sum_weights, -1, faces[:, (i + 2) % 3], cot_angle)
|
|
elif weight == 'uniform':
|
|
for i in range(3):
|
|
sum_verts = torch.index_add(sum_verts, -2, faces[:, i], face_verts[..., (i + 1) % 3, :])
|
|
sum_weights = torch.index_add(sum_weights, -1, faces[:, i], torch.ones_like(face_verts[..., i, 0]))
|
|
else:
|
|
raise NotImplementedError
|
|
return sum_verts / (sum_weights[..., None] + 1e-7)
|
|
|
|
|
|
def laplacian_smooth_mesh(vertices: torch.Tensor, faces: torch.Tensor, weight: str = 'uniform', times: int = 5) -> torch.Tensor:
|
|
"""Laplacian smooth with cotangent weights
|
|
|
|
Args:
|
|
vertices (torch.Tensor): shape (..., N, 3)
|
|
faces (torch.Tensor): shape (T, 3)
|
|
weight (str): 'uniform' or 'cotangent'
|
|
"""
|
|
for _ in range(times):
|
|
vertices = laplacian(vertices, faces, weight)
|
|
return vertices
|
|
|
|
|
|
def taubin_smooth_mesh(vertices: torch.Tensor, faces: torch.Tensor, lambda_: float = 0.5, mu_: float = -0.51) -> torch.Tensor:
|
|
"""Taubin smooth mesh
|
|
|
|
Args:
|
|
vertices (torch.Tensor): _description_
|
|
faces (torch.Tensor): _description_
|
|
lambda_ (float, optional): _description_. Defaults to 0.5.
|
|
mu_ (float, optional): _description_. Defaults to -0.51.
|
|
|
|
Returns:
|
|
torch.Tensor: _description_
|
|
"""
|
|
pt = vertices + lambda_ * laplacian_smooth_mesh(vertices, faces)
|
|
p = pt + mu_ * laplacian_smooth_mesh(pt, faces)
|
|
return p
|
|
|
|
|
|
def laplacian_hc_smooth_mesh(vertices: torch.Tensor, faces: torch.Tensor, times: int = 5, alpha: float = 0.5, beta: float = 0.5, weight: str = 'uniform'):
|
|
"""HC algorithm from Improved Laplacian Smoothing of Noisy Surface Meshes by J.Vollmer et al.
|
|
"""
|
|
p = vertices
|
|
for i in range(times):
|
|
q = p
|
|
p = laplacian_smooth_mesh(vertices, faces, weight)
|
|
b = p - (alpha * vertices + (1 - alpha) * q)
|
|
p = p - (beta * b + (1 - beta) * laplacian_smooth_mesh(b, faces, weight)) * 0.8
|
|
return p
|