LTT/PRM
24
1# Copyright (c) 2020-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2#3# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual4# property and proprietary rights in and to this material, related5# documentation and any modifications thereto. Any use, reproduction, 6# disclosure or distribution of this material and related documentation 7# without an express license agreement from NVIDIA CORPORATION or 8# its affiliates is strictly prohibited.9 10import os11import numpy as np12import torch13 14from . import obj15from src.models.geometry.rep_3d import util16 17######################################################################################18# Base mesh class19######################################################################################20class Mesh:21 def __init__(self, v_pos=None, t_pos_idx=None, v_nrm=None, t_nrm_idx=None, v_tex=None, t_tex_idx=None, v_tng=None, t_tng_idx=None, material=None, base=None):22 self.v_pos = v_pos23 self.v_nrm = v_nrm24 self.v_tex = v_tex25 self.v_tng = v_tng26 self.t_pos_idx = t_pos_idx27 self.t_nrm_idx = t_nrm_idx28 self.t_tex_idx = t_tex_idx29 self.t_tng_idx = t_tng_idx30 self.material = material31 32 if base is not None:33 self.copy_none(base)34 35 def copy_none(self, other):36 if self.v_pos is None:37 self.v_pos = other.v_pos38 if self.t_pos_idx is None:39 self.t_pos_idx = other.t_pos_idx40 if self.v_nrm is None:41 self.v_nrm = other.v_nrm42 if self.t_nrm_idx is None:43 self.t_nrm_idx = other.t_nrm_idx44 if self.v_tex is None:45 self.v_tex = other.v_tex46 if self.t_tex_idx is None:47 self.t_tex_idx = other.t_tex_idx48 if self.v_tng is None:49 self.v_tng = other.v_tng50 if self.t_tng_idx is None:51 self.t_tng_idx = other.t_tng_idx52 if self.material is None:53 self.material = other.material54 55 def clone(self):56 out = Mesh(base=self)57 if out.v_pos is not None:58 out.v_pos = out.v_pos.clone().detach()59 if out.t_pos_idx is not None:60 out.t_pos_idx = out.t_pos_idx.clone().detach()61 if out.v_nrm is not None:62 out.v_nrm = out.v_nrm.clone().detach()63 if out.t_nrm_idx is not None:64 out.t_nrm_idx = out.t_nrm_idx.clone().detach()65 if out.v_tex is not None:66 out.v_tex = out.v_tex.clone().detach()67 if out.t_tex_idx is not None:68 out.t_tex_idx = out.t_tex_idx.clone().detach()69 if out.v_tng is not None:70 out.v_tng = out.v_tng.clone().detach()71 if out.t_tng_idx is not None:72 out.t_tng_idx = out.t_tng_idx.clone().detach()73 return out74 def rotate_x_90(self):75 # 定义绕X轴旋转90度的旋转矩阵76 rotate_x = torch.tensor([[1, 0, 0, 0], 77 [0, 0, 1, 0], 78 [0, -1, 0, 0], 79 [0, 0, 0, 1]], dtype=torch.float32, device=self.v_pos.device)80 81 # 将旋转矩阵应用到顶点坐标82 if self.v_pos is not None:83 v_pos_homo = torch.cat((self.v_pos, torch.ones(self.v_pos.shape[0], 1, device=self.v_pos.device)), dim=1)84 v_pos_rotated = v_pos_homo @ rotate_x.T85 self.v_pos = v_pos_rotated[:, :3]86 87 # 将旋转矩阵应用到法线88 if self.v_nrm is not None:89 v_nrm_homo = torch.cat((self.v_nrm, torch.zeros(self.v_nrm.shape[0], 1, device=self.v_nrm.device)), dim=1)90 v_nrm_rotated = v_nrm_homo @ rotate_x.T91 self.v_nrm = v_nrm_rotated[:, :3]92######################################################################################93# Mesh loeading helper94######################################################################################95 96def load_mesh(filename, mtl_override=None):97 name, ext = os.path.splitext(filename)98 if ext == ".obj":99 return obj.load_obj(filename, clear_ks=True, mtl_override=mtl_override)100 assert False, "Invalid mesh file extension"101 102######################################################################################103# Compute AABB104######################################################################################105def aabb(mesh):106 return torch.min(mesh.v_pos, dim=0).values, torch.max(mesh.v_pos, dim=0).values107 108######################################################################################109# Compute unique edge list from attribute/vertex index list110######################################################################################111def compute_edges(attr_idx, return_inverse=False):112 with torch.no_grad():113 # Create all edges, packed by triangle114 all_edges = torch.cat((115 torch.stack((attr_idx[:, 0], attr_idx[:, 1]), dim=-1),116 torch.stack((attr_idx[:, 1], attr_idx[:, 2]), dim=-1),117 torch.stack((attr_idx[:, 2], attr_idx[:, 0]), dim=-1),118 ), dim=-1).view(-1, 2)119 120 # Swap edge order so min index is always first121 order = (all_edges[:, 0] > all_edges[:, 1]).long().unsqueeze(dim=1)122 sorted_edges = torch.cat((123 torch.gather(all_edges, 1, order),124 torch.gather(all_edges, 1, 1 - order)125 ), dim=-1)126 127 # Eliminate duplicates and return inverse mapping128 return torch.unique(sorted_edges, dim=0, return_inverse=return_inverse)129 130######################################################################################131# Compute unique edge to face mapping from attribute/vertex index list132######################################################################################133def compute_edge_to_face_mapping(attr_idx, return_inverse=False):134 with torch.no_grad():135 # Get unique edges136 # Create all edges, packed by triangle137 all_edges = torch.cat((138 torch.stack((attr_idx[:, 0], attr_idx[:, 1]), dim=-1),139 torch.stack((attr_idx[:, 1], attr_idx[:, 2]), dim=-1),140 torch.stack((attr_idx[:, 2], attr_idx[:, 0]), dim=-1),141 ), dim=-1).view(-1, 2)142 143 # Swap edge order so min index is always first144 order = (all_edges[:, 0] > all_edges[:, 1]).long().unsqueeze(dim=1)145 sorted_edges = torch.cat((146 torch.gather(all_edges, 1, order),147 torch.gather(all_edges, 1, 1 - order)148 ), dim=-1)149 150 # Elliminate duplicates and return inverse mapping151 unique_edges, idx_map = torch.unique(sorted_edges, dim=0, return_inverse=True)152 153 tris = torch.arange(attr_idx.shape[0]).repeat_interleave(3).cuda()154 155 tris_per_edge = torch.zeros((unique_edges.shape[0], 2), dtype=torch.int64).cuda()156 157 # Compute edge to face table158 mask0 = order[:,0] == 0159 mask1 = order[:,0] == 1160 tris_per_edge[idx_map[mask0], 0] = tris[mask0]161 tris_per_edge[idx_map[mask1], 1] = tris[mask1]162 163 return tris_per_edge164 165######################################################################################166# Align base mesh to reference mesh:move & rescale to match bounding boxes.167######################################################################################168def unit_size(mesh):169 with torch.no_grad():170 vmin, vmax = aabb(mesh)171 scale = 2 / torch.max(vmax - vmin).item()172 v_pos = mesh.v_pos - (vmax + vmin) / 2 # Center mesh on origin173 v_pos = v_pos * scale # Rescale to unit size174 175 return Mesh(v_pos, base=mesh)176 177######################################################################################178# Center & scale mesh for rendering179######################################################################################180def center_by_reference(base_mesh, ref_aabb, scale):181 center = (ref_aabb[0] + ref_aabb[1]) * 0.5182 scale = scale / torch.max(ref_aabb[1] - ref_aabb[0]).item()183 v_pos = (base_mesh.v_pos - center[None, ...]) * scale184 return Mesh(v_pos, base=base_mesh)185 186######################################################################################187# Simple smooth vertex normal computation188######################################################################################189def auto_normals(imesh):190 191 i0 = imesh.t_pos_idx[:, 0]192 i1 = imesh.t_pos_idx[:, 1]193 i2 = imesh.t_pos_idx[:, 2]194 195 v0 = imesh.v_pos[i0, :]196 v1 = imesh.v_pos[i1, :]197 v2 = imesh.v_pos[i2, :]198 199 face_normals = torch.cross(v1 - v0, v2 - v0)200 201 # Splat face normals to vertices202 v_nrm = torch.zeros_like(imesh.v_pos)203 v_nrm.scatter_add_(0, i0[:, None].repeat(1,3), face_normals)204 v_nrm.scatter_add_(0, i1[:, None].repeat(1,3), face_normals)205 v_nrm.scatter_add_(0, i2[:, None].repeat(1,3), face_normals)206 207 # Normalize, replace zero (degenerated) normals with some default value208 v_nrm = torch.where(util.dot(v_nrm, v_nrm) > 1e-20, v_nrm, torch.tensor([0.0, 0.0, 1.0], dtype=torch.float32, device='cuda'))209 v_nrm = util.safe_normalize(v_nrm)210 211 if torch.is_anomaly_enabled():212 assert torch.all(torch.isfinite(v_nrm))213 214 return Mesh(v_nrm=v_nrm, t_nrm_idx=imesh.t_pos_idx, base=imesh)215 216######################################################################################217# Compute tangent space from texture map coordinates218# Follows http://www.mikktspace.com/ conventions219######################################################################################220def compute_tangents(imesh):221 vn_idx = [None] * 3222 pos = [None] * 3223 tex = [None] * 3224 for i in range(0,3):225 pos[i] = imesh.v_pos[imesh.t_pos_idx[:, i]]226 tex[i] = imesh.v_tex[imesh.t_tex_idx[:, i]]227 vn_idx[i] = imesh.t_nrm_idx[:, i]228 229 tangents = torch.zeros_like(imesh.v_nrm)230 231 # Compute tangent space for each triangle232 uve1 = tex[1] - tex[0]233 uve2 = tex[2] - tex[0]234 pe1 = pos[1] - pos[0]235 pe2 = pos[2] - pos[0]236 237 nom = (pe1 * uve2[..., 1:2] - pe2 * uve1[..., 1:2])238 denom = (uve1[..., 0:1] * uve2[..., 1:2] - uve1[..., 1:2] * uve2[..., 0:1])239 240 # Avoid division by zero for degenerated texture coordinates241 tang = nom / torch.where(denom > 0.0, torch.clamp(denom, min=1e-6), torch.clamp(denom, max=-1e-6))242 243 # Update all 3 vertices244 for i in range(0,3):245 idx = vn_idx[i][:, None].repeat(1,3)246 tangents.scatter_add_(0, idx, tang) # tangents[n_i] = tangents[n_i] + tang247 248 # Normalize and make sure tangent is perpendicular to normal249 tangents = util.safe_normalize(tangents)250 tangents = util.safe_normalize(tangents - util.dot(tangents, imesh.v_nrm) * imesh.v_nrm)251 252 if torch.is_anomaly_enabled():253 assert torch.all(torch.isfinite(tangents))254 255 return Mesh(v_tng=tangents, t_tng_idx=imesh.t_nrm_idx, base=imesh)256