Spaces:
Sleeping
Sleeping
File size: 4,445 Bytes
98a77e0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
# Copyright (c) 2020-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
import torch
import nvdiffrast.torch as dr
from . import util
from . import mesh
######################################################################################
# Computes the image gradient, useful for kd/ks smoothness losses
######################################################################################
def image_grad(buf, std=0.01):
t, s = torch.meshgrid(torch.linspace(-1.0 + 1.0 / buf.shape[1], 1.0 - 1.0 / buf.shape[1], buf.shape[1], device="cuda"),
torch.linspace(-1.0 + 1.0 / buf.shape[2], 1.0 - 1.0 / buf.shape[2], buf.shape[2], device="cuda"),
indexing='ij')
tc = torch.normal(mean=0, std=std, size=(buf.shape[0], buf.shape[1], buf.shape[2], 2), device="cuda") + torch.stack((s, t), dim=-1)[None, ...]
tap = dr.texture(buf, tc, filter_mode='linear', boundary_mode='clamp')
return torch.abs(tap[..., :-1] - buf[..., :-1]) * tap[..., -1:] * buf[..., -1:]
######################################################################################
# Computes the avergage edge length of a mesh.
# Rough estimate of the tessellation of a mesh. Can be used e.g. to clamp gradients
######################################################################################
def avg_edge_length(v_pos, t_pos_idx):
e_pos_idx = mesh.compute_edges(t_pos_idx)
edge_len = util.length(v_pos[:, e_pos_idx[:, 0]] - v_pos[:, e_pos_idx[:, 1]])
return torch.mean(edge_len)
######################################################################################
# Laplacian regularization using umbrella operator (Fujiwara / Desbrun).
# https://mgarland.org/class/geom04/material/smoothing.pdf
######################################################################################
def laplace_regularizer_const(v_pos, t_pos_idx):
batch_size = v_pos.shape[0]
term = torch.zeros_like(v_pos)
norm = torch.zeros_like(v_pos[..., 0:1])
v0 = v_pos[:, t_pos_idx[0, :, 0], :]
v1 = v_pos[:, t_pos_idx[0, :, 1], :]
v2 = v_pos[:, t_pos_idx[0, :, 2], :]
term.scatter_add_(1, t_pos_idx[..., 0:1].repeat(batch_size, 1, 3), (v1 - v0) + (v2 - v0))
term.scatter_add_(1, t_pos_idx[..., 1:2].repeat(batch_size, 1, 3), (v0 - v1) + (v2 - v1))
term.scatter_add_(1, t_pos_idx[..., 2:3].repeat(batch_size, 1, 3), (v0 - v2) + (v1 - v2))
two = torch.ones_like(v0) * 2.0
# norm.scatter_add_(1, t_pos_idx[..., 0:1].repeat(batch_size, 1, 3), two)
# norm.scatter_add_(1, t_pos_idx[..., 1:2].repeat(batch_size, 1, 3), two)
# norm.scatter_add_(1, t_pos_idx[..., 2:3].repeat(batch_size, 1, 3), two)
norm.scatter_add_(1, t_pos_idx[..., 0:1].repeat(batch_size, 1, 1), two)
norm.scatter_add_(1, t_pos_idx[..., 1:2].repeat(batch_size, 1, 1), two)
norm.scatter_add_(1, t_pos_idx[..., 2:3].repeat(batch_size, 1, 1), two)
term = term / torch.clamp(norm, min=1.0)
return torch.mean(term ** 2)
######################################################################################
# Smooth vertex normals
######################################################################################
def normal_consistency(v_pos, t_pos_idx):
# Compute face normals
v0 = v_pos[:, t_pos_idx[0, :, 0]]
v1 = v_pos[:, t_pos_idx[0, :, 1]]
v2 = v_pos[:, t_pos_idx[0, :, 2]]
face_normals = util.safe_normalize(torch.cross(v1 - v0, v2 - v0, dim=-1))
tris_per_edge = mesh.compute_edge_to_face_mapping(t_pos_idx)
# Fetch normals for both faces sharing an edge
n0 = face_normals[:, tris_per_edge[:, 0], :]
n1 = face_normals[:, tris_per_edge[:, 1], :]
# Compute error metric based on normal difference
term = torch.clamp(util.dot(n0, n1), min=-1.0, max=1.0)
term = (1.0 - term) * 0.5
return torch.mean(torch.abs(term))
def get_edge_length(v_pos, t_pos_idx):
e_pos_idx = mesh.compute_edges(t_pos_idx)
edge_len = util.length(v_pos[:, e_pos_idx[:, 0]] - v_pos[:, e_pos_idx[:, 1]])
return edge_len
|