# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# pyre-strict
from typing import Dict, List, Optional, Tuple, Union
import torch as th
import torch.nn.functional as thf
from drtk.utils.indexing import index
from torch import Tensor
eps = 1e-8
def face_dpdt(
v: th.Tensor, vt: th.Tensor, vi: th.Tensor, vti: th.Tensor
) -> Tuple[th.Tensor, th.Tensor]:
"""
This function calculates the transposed Jacobian matrix (∂p/∂t)^T for each triangle.
Where:
- p represents the 3D coordinates of a point on the plane of the triangle,
- t denotes the UV coordinates assiciated with the point.
Args:
v: vertex position tensor
N x V x 3
vt: vertex uv tensor
N x T x 2
vi: face vertex position index list tensor
F x 3
vti: face vertex uv index list tensor
F x 3
Jacobian is computed as:
∂p/∂t = ∂p / ∂b * (∂t / ∂b)^-1
Where b - barycentric coordinates
However the implementation computes a transposed Jacobian (purely from
practical perspective - fewer permutations are needed), so the above
becomes:
(∂p/∂t)^T = ((∂t / ∂b)^T)^-1 * (∂p / ∂b)^T
Returns:
dpdt - transposed Jacobian (∂p/∂t)^T. Shape: N x F x 2 x 3
Where ∂p∂t[..., i, j] = ∂p[..., j] / ∂t[..., i]
v012 - vertex positions per triangle. Shape: N x F x 3
where: N - batch size; F - number of triangles
"""
if v.ndim != 3:
raise ValueError(f"Expected v to be 3D, got {v.ndim}D")
if vt.ndim != 3:
raise ValueError(f"Expected vt to be 3D, got {vt.ndim}D")
if vt.shape[0] != v.shape[0]:
raise ValueError(
f"Expected vt to have the same batch size as v, got {vt.shape[0]} and {v.shape[0]}"
)
v012 = v[:, vi]
vt012 = vt[:, vti]
dpdb_t = v012[:, :, 1:3] - v012[:, :, 0:1]
dtdb_t = vt012[:, :, 1:3] - vt012[:, :, 0:1]
# (db / dt)^T = ((dt / db)^T)^-1
dbdt_t = th.inverse(dtdb_t)
# (dp / dt)^T = (db / dt)^T) * (dp / db)^T
dpdt_t = dbdt_t @ dpdb_t
return dpdt_t, v012
def face_attribute_to_vert(v: th.Tensor, vi: th.Tensor, attr: th.Tensor) -> Tensor:
"""
For each vertex, computes a summation of the face attributes to which the
vertex belongs.
Args:
v: vertex positions, shape [N, V, *]
vi: face vertex indices, shape [F, 3] (shared topology, broadcast
across the batch) or [B, F, 3] where B is 1 or N
attr: per-face attribute, shape [N, F, A]
Returns:
Per-vertex sum, shape [N, V, A]
"""
attr = (
attr[:, :, None]
.expand(-1, -1, 3, -1)
.reshape(attr.shape[0], -1, attr.shape[-1])
)
n = v.shape[0]
if vi.dim() == 2:
# Shared topology: flatten [F, 3] -> [F*3] and broadcast across batch.
vi_flat = vi.reshape(-1).unsqueeze(0).expand(n, -1)
elif vi.dim() == 3:
# Batched topology: B must be 1 (broadcasts) or equal to N.
vi_flat = vi.reshape(vi.shape[0], -1).expand(n, -1)
else:
raise ValueError(
f"Expected vi to be 2D [F, 3] or 3D [B, F, 3], got {vi.dim()}D"
)
vattr = th.zeros(v.shape[:-1], dtype=v.dtype, device=v.device)
vattr = th.stack(
[vattr.scatter_add(1, vi_flat, attr[..., i]) for i in range(attr.shape[-1])],
dim=-1,
)
return vattr
[docs]
def face_info(
v: th.Tensor, vi: th.Tensor, to_compute: Optional[List[str]] = None
) -> Union[th.Tensor, Dict[str, th.Tensor]]:
"""Given a set of vertices ``v`` and indices ``vi`` indexing into ``v``
defining a set of faces, compute face information (normals, edges, face
areas) for each face.
Args:
v: Vertex positions, shape [batch_size, n_vertices, 3]
vi: Vertex indices, shape [n_faces, 3]
to_compute: list of desired information. Any of: {normals, edges, areas}, defaults to all.
Returns:
Dict: Face information in the following format::
{
"normals": shape [batch_size, n_faces, 3]
"edges": shape [batch_size, n_faces, 3, 3]
"areas": shape [batch_size, n_faces, 1]
}
or just one of the above values not in a Dict if only one is
requested.
"""
if to_compute is None:
to_compute = ["normals", "edges", "areas"]
b = v.shape[0]
vi = vi.expand(b, -1, -1)
p0 = th.stack([index(v[i], vi[i, :, 0], 0) for i in range(b)])
p1 = th.stack([index(v[i], vi[i, :, 1], 0) for i in range(b)])
p2 = th.stack([index(v[i], vi[i, :, 2], 0) for i in range(b)])
v0 = p1 - p0
v1 = p0 - p2
need_normals = "normals" in to_compute
need_areas = "areas" in to_compute
need_edges = "edges" in to_compute
output = {}
if need_normals or need_areas:
normals = th.cross(v1, v0, dim=-1)
norm = th.linalg.vector_norm(normals, dim=-1, keepdim=True)
if need_areas:
output["areas"] = 0.5 * norm
if need_normals:
output["normals"] = normals / norm.clamp(min=eps)
if need_edges:
v2 = p2 - p1
output["edges"] = th.stack([v0, v1, v2], dim=2)
if len(to_compute) == 1:
return output[to_compute[0]]
else:
return output
[docs]
def vert_binormals(v: Tensor, vt: Tensor, vi: Tensor, vti: Tensor) -> Tensor:
# Compute (dp/dt)^T
dpdt_t, vf = face_dpdt(v, vt, vi, vti)
# Take the dp/dt.u part. Produces u vector in 3D world-space which we use for binormal vector
fbnorms = dpdt_t[:, :, 0, :]
vbnorms = face_attribute_to_vert(v, vi, fbnorms)
return thf.normalize(vbnorms, dim=-1)
[docs]
def vert_normals(
v: th.Tensor, vi: th.Tensor, fnorms: Optional[th.Tensor] = None
) -> th.Tensor:
"""Given a set of vertices ``v`` and indices ``vi`` indexing into ``v``
defining a set of faces, compute normals for each vertex by averaging the
face normals for each face which includes that vertex.
Args:
v: Vertex positions, shape [batch_size, n_vertices, 3]
vi: Vertex indices, shape [batch_size, n_faces, 3]
fnorms: Face normals. Optional, provide them if available, otherwise they will be computed
from `v` and `vi`. Shape [n_faces, 3]
Returns:
th.Tensor: Vertex normals, shape [batch_size, n_vertices, 3]
"""
if fnorms is None:
fnorms = face_info(v, vi, ["normals"])
assert isinstance(fnorms, th.Tensor)
vnorms = face_attribute_to_vert(v, vi, fnorms)
return thf.normalize(vnorms, dim=-1)