# 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 numpy as np import torch from ..render import mesh from ..render import render from ..render import regularizer import kaolin from ..render import util as render_utils import torch.nn.functional as F ############################################################################### # Marching tetrahedrons implementation (differentiable), adapted from # https://github.com/NVIDIAGameWorks/kaolin/blob/master/kaolin/ops/conversions/tetmesh.py ############################################################################### class DMTet: def __init__(self): self.triangle_table = torch.tensor([ [-1, -1, -1, -1, -1, -1], [ 1, 0, 2, -1, -1, -1], [ 4, 0, 3, -1, -1, -1], [ 1, 4, 2, 1, 3, 4], [ 3, 1, 5, -1, -1, -1], [ 2, 3, 0, 2, 5, 3], [ 1, 4, 0, 1, 5, 4], [ 4, 2, 5, -1, -1, -1], [ 4, 5, 2, -1, -1, -1], [ 4, 1, 0, 4, 5, 1], [ 3, 2, 0, 3, 5, 2], [ 1, 3, 5, -1, -1, -1], [ 4, 1, 2, 4, 3, 1], [ 3, 0, 4, -1, -1, -1], [ 2, 0, 1, -1, -1, -1], [-1, -1, -1, -1, -1, -1] ], dtype=torch.long, device='cuda') self.num_triangles_table = torch.tensor([0,1,1,2,1,2,2,1,1,2,2,1,2,1,1,0], dtype=torch.long, device='cuda') self.base_tet_edges = torch.tensor([0,1,0,2,0,3,1,2,1,3,2,3], dtype=torch.long, device='cuda') ############################################################################### # Utility functions ############################################################################### def sort_edges(self, edges_ex2): with torch.no_grad(): order = (edges_ex2[:,0] > edges_ex2[:,1]).long() order = order.unsqueeze(dim=1) a = torch.gather(input=edges_ex2, index=order, dim=1) b = torch.gather(input=edges_ex2, index=1-order, dim=1) return torch.stack([a, b],-1) def map_uv(self, faces, face_gidx, max_idx): N = int(np.ceil(np.sqrt((max_idx+1)//2))) tex_y, tex_x = torch.meshgrid( torch.linspace(0, 1 - (1 / N), N, dtype=torch.float32, device="cuda"), torch.linspace(0, 1 - (1 / N), N, dtype=torch.float32, device="cuda"), indexing='ij' ) pad = 0.9 / N uvs = torch.stack([ tex_x , tex_y, tex_x + pad, tex_y, tex_x + pad, tex_y + pad, tex_x , tex_y + pad ], dim=-1).view(-1, 2) def _idx(tet_idx, N): x = tet_idx % N y = torch.div(tet_idx, N, rounding_mode='trunc') return y * N + x tet_idx = _idx(torch.div(face_gidx, 2, rounding_mode='trunc'), N) tri_idx = face_gidx % 2 uv_idx = torch.stack(( tet_idx * 4, tet_idx * 4 + tri_idx + 1, tet_idx * 4 + tri_idx + 2 ), dim = -1). view(-1, 3) return uvs, uv_idx ############################################################################### # Marching tets implementation ############################################################################### def __call__(self, pos_nx3, sdf_n, tet_fx4, get_tet_gidx=False): with torch.no_grad(): occ_n = sdf_n > 0 occ_fx4 = occ_n[tet_fx4.reshape(-1)].reshape(-1,4) occ_sum = torch.sum(occ_fx4, -1) valid_tets = (occ_sum>0) & (occ_sum<4) occ_sum = occ_sum[valid_tets] # find all vertices all_edges = tet_fx4[valid_tets][:,self.base_tet_edges].reshape(-1,2) all_edges = self.sort_edges(all_edges) unique_edges, idx_map = torch.unique(all_edges,dim=0, return_inverse=True) unique_edges = unique_edges.long() mask_edges = occ_n[unique_edges.reshape(-1)].reshape(-1,2).sum(-1) == 1 mapping = torch.ones((unique_edges.shape[0]), dtype=torch.long, device="cuda") * -1 mapping[mask_edges] = torch.arange(mask_edges.sum(), dtype=torch.long,device="cuda") idx_map = mapping[idx_map] # map edges to verts interp_v = unique_edges[mask_edges] edges_to_interp = pos_nx3[interp_v.reshape(-1)].reshape(-1,2,3) edges_to_interp_sdf = sdf_n[interp_v.reshape(-1)].reshape(-1,2,1) edges_to_interp_sdf[:,-1] *= -1 denominator = edges_to_interp_sdf.sum(1,keepdim = True) edges_to_interp_sdf = torch.flip(edges_to_interp_sdf, [1])/denominator verts = (edges_to_interp * edges_to_interp_sdf).sum(1) idx_map = idx_map.reshape(-1,6) v_id = torch.pow(2, torch.arange(4, dtype=torch.long, device="cuda")) tetindex = (occ_fx4[valid_tets] * v_id.unsqueeze(0)).sum(-1) num_triangles = self.num_triangles_table[tetindex] # Generate triangle indices faces = torch.cat(( torch.gather(input=idx_map[num_triangles == 1], dim=1, index=self.triangle_table[tetindex[num_triangles == 1]][:, :3]).reshape(-1,3), torch.gather(input=idx_map[num_triangles == 2], dim=1, index=self.triangle_table[tetindex[num_triangles == 2]][:, :6]).reshape(-1,3), ), dim=0) # Get global face index (static, does not depend on topology) num_tets = tet_fx4.shape[0] tet_gidx = torch.arange(num_tets, dtype=torch.long, device="cuda")[valid_tets] face_gidx = torch.cat(( tet_gidx[num_triangles == 1]*2, torch.stack((tet_gidx[num_triangles == 2]*2, tet_gidx[num_triangles == 2]*2 + 1), dim=-1).view(-1) ), dim=0) uvs, uv_idx = self.map_uv(faces, face_gidx, num_tets*2) if get_tet_gidx: face_to_valid_tet = torch.cat(( tet_gidx[num_triangles == 1], torch.stack((tet_gidx[num_triangles == 2], tet_gidx[num_triangles == 2]), dim=-1).view(-1) ), dim=0) return verts, faces, uvs, uv_idx, face_to_valid_tet.long() else: return verts, faces, uvs, uv_idx ############################################################################### # Regularizer ############################################################################### def sdf_reg_loss(sdf, all_edges): sdf_f1x6x2 = sdf[all_edges.reshape(-1)].reshape(-1,2) mask = torch.sign(sdf_f1x6x2[...,0]) != torch.sign(sdf_f1x6x2[...,1]) sdf_f1x6x2 = sdf_f1x6x2[mask] sdf_diff = torch.nn.functional.binary_cross_entropy_with_logits(sdf_f1x6x2[...,0], (sdf_f1x6x2[...,1] > 0).float()) + \ torch.nn.functional.binary_cross_entropy_with_logits(sdf_f1x6x2[...,1], (sdf_f1x6x2[...,0] > 0).float()) return sdf_diff ############################################################################### # Geometry interface ############################################################################### class DMTetGeometryFixedTopo(torch.nn.Module): def __init__(self, dmt_geometry, base_mesh, grid_res, scale, FLAGS, deform_scale=1.0, **kwargs): super(DMTetGeometryFixedTopo, self).__init__() self.FLAGS = FLAGS self.grid_res = grid_res self.marching_tets = DMTet() self.initial_guess = base_mesh self.scale = scale self.tanh = False self.deform_scale = deform_scale tets = np.load('./data/tets/{}_tets_cropped.npz'.format(self.grid_res)) self.verts = torch.tensor(tets['vertices'], dtype=torch.float32, device='cuda') * scale self.indices = torch.tensor(tets['indices'], dtype=torch.long, device='cuda') self.generate_edges() self.sdf_sign = torch.nn.Parameter(torch.sign(dmt_geometry.sdf.data + 1e-8).float(), requires_grad=False) self.sdf_sign.data[self.sdf_sign.data == 0] = 1.0 ## Avoid abiguity self.register_parameter('sdf_sign', self.sdf_sign) self.sdf_abs = torch.nn.Parameter(torch.ones_like(dmt_geometry.sdf), requires_grad=False) self.register_parameter('sdf_abs', self.sdf_abs) self.deform = torch.nn.Parameter(dmt_geometry.deform.data, requires_grad=True) self.register_parameter('deform', self.deform) self.sdf_abs_ema = torch.nn.Parameter(self.sdf_abs.clone().detach(), requires_grad=False) self.deform_ema = torch.nn.Parameter(self.deform.clone().detach(), requires_grad=False) def set_init_v_pos(self): with torch.no_grad(): v_deformed = self.get_deformed() verts, faces, uvs, uv_idx = self.marching_tets(v_deformed, self.sdf_sign * self.sdf_abs.abs(), self.indices) self.initial_guess_v_pos = verts def generate_edges(self): with torch.no_grad(): edges = torch.tensor([0,1,0,2,0,3,1,2,1,3,2,3], dtype = torch.long, device = "cuda") all_edges = self.indices[:,edges].reshape(-1,2) all_edges_sorted = torch.sort(all_edges, dim=1)[0] self.all_edges = torch.unique(all_edges_sorted, dim=0) def getAABB(self): return torch.min(self.verts, dim=0).values, torch.max(self.verts, dim=0).values def getVertNNDist(self): raise NotImplementedError v_deformed = (self.verts + 2 / (self.grid_res * 2) * torch.tanh(self.deform)).unsqueeze(0) return (pytorch3d.ops.knn.knn_points(v_deformed, v_deformed, K=2).dists[0, :, -1].detach()) ## K=2 because dist(self, self)=0 def getMesh(self, material): # Run DM tet to get a base mesh v_deformed = self.get_deformed() verts, faces, uvs, uv_idx = self.marching_tets(v_deformed, self.sdf_sign * self.sdf_abs.abs(), self.indices) imesh = mesh.Mesh(verts, faces, v_tex=uvs, t_tex_idx=uv_idx, material=material) # Run mesh operations to generate tangent space imesh = mesh.auto_normals(imesh) imesh = mesh.compute_tangents(imesh) return imesh def getMesh_tet_gidx(self, material): # Run DM tet to get a base mesh v_deformed = self.get_deformed() verts, faces, uvs, uv_idx, tet_gidx = self.marching_tets( v_deformed, self.sdf_sign * self.sdf_abs.abs(), self.indices, get_tet_gidx=True) imesh = mesh.Mesh(verts, faces, v_tex=uvs, t_tex_idx=uv_idx, material=material) # Run mesh operations to generate tangent space imesh = mesh.auto_normals(imesh) imesh = mesh.compute_tangents(imesh) return imesh, tet_gidx def update_ema(self, ema_coeff=0.9): return def get_deformed(self): if self.tanh: v_deformed = self.verts + 2 / (self.grid_res * 2) * torch.tanh(self.deform) * self.deform_scale else: v_deformed = self.verts + 2 / (self.grid_res * 2) * self.deform * self.deform_scale return v_deformed def getValidTetIdx(self): # Run DM tet to get a base mesh v_deformed = self.get_deformed() verts, faces, uvs, uv_idx, tet_gidx = self.marching_tets( v_deformed, self.sdf_sign * self.sdf_abs.abs(), self.indices, get_tet_gidx=True) return tet_gidx.long() def getValidVertsIdx(self): # Run DM tet to get a base mesh v_deformed = self.get_deformed() verts, faces, uvs, uv_idx, tet_gidx = self.marching_tets( v_deformed, self.sdf_sign * self.sdf_abs.abs(), self.indices, get_tet_gidx=True) return self.indices[tet_gidx.long()].unique() def getTetCenters(self): v_deformed = self.get_deformed() # size: N x 3 face_verts = v_deformed[self.indices] # size: M x 4 x 3 face_centers = face_verts.mean(dim=1) # size: M x 3 return face_centers def clamp_deform(self): if not self.tanh: self.deform.data[:] = self.deform.data.clamp(-0.99, 0.99) def render(self, glctx, target, lgt, opt_material, bsdf=None, ema=False, xfm_lgt=None, get_visible_tets=False): opt_mesh = self.getMesh(opt_material) tet_centers = self.getTetCenters() if get_visible_tets else None return render.render_mesh( glctx, opt_mesh, target['mvp'], target['campos'], lgt, target['resolution'], spp=target['spp'], msaa=True, background=target['background'], bsdf=bsdf, xfm_lgt=xfm_lgt, tet_centers=tet_centers) def render_with_mesh(self, glctx, target, lgt, opt_material, bsdf=None, xfm_lgt=None): opt_mesh = self.getMesh(opt_material) return opt_mesh, render.render_mesh( glctx, opt_mesh, target['mvp'], target['campos'], lgt, target['resolution'], spp=target['spp'], num_layers=self.FLAGS.layers, msaa=True, background=target['background'], bsdf=bsdf, xfm_lgt=xfm_lgt) def tick(self, glctx, target, lgt, opt_material, loss_fn, iteration, with_reg=True, xfm_lgt=None, no_depth_thin=True): # ============================================================================================== # Render optimizable object with identical conditions # ============================================================================================== imesh, buffers = self.render_with_mesh(glctx, target, lgt, opt_material, xfm_lgt=xfm_lgt) # ============================================================================================== # Compute loss # ============================================================================================== t_iter = iteration / self.FLAGS.iter # Image-space loss, split into a coverage component and a color component color_ref = target['img'] img_loss = torch.nn.functional.mse_loss(buffers['shaded'][..., 3:], color_ref[..., 3:]) img_loss = img_loss + loss_fn( buffers['shaded'][..., 0:3] * color_ref[..., 3:], color_ref[..., 0:3] * color_ref[..., 3:] ) mask = target['mask'][:, :, :, 0] if no_depth_thin: valid_depth_mask = ( (target['depth_second'] >= 0).float() * ((target['depth_second'] - target['depth']).abs() >= 5e-3).float() ).detach() else: valid_depth_mask = 1.0 depth_diff = (buffers['depth'][:, :, :, :1] - target['depth'][:, :, :, :1]).abs() * mask.unsqueeze(-1) * valid_depth_mask depth_diff = (buffers['depth_second'][:, :, :, :1] - target['depth_second'][:, :, :, :1]).abs() * mask.unsqueeze(-1) * valid_depth_mask * 1e-1 l1_loss_mask = (depth_diff < 1.0).float() img_loss = img_loss + (l1_loss_mask * depth_diff + (1 - l1_loss_mask) * depth_diff.pow(2)).mean() * 100.0 reg_loss = torch.tensor([0], dtype=torch.float32, device="cuda") # Compute regularizer. reg_loss += regularizer.laplace_regularizer_const(imesh.v_pos - self.initial_guess_v_pos, imesh.t_pos_idx) * self.FLAGS.laplace_scale * (1 - t_iter) * 1e-2 ### Chamfer distance for ShapeNet pred_points = kaolin.ops.mesh.sample_points(imesh.v_pos.unsqueeze(0), imesh.t_pos_idx, 50000)[0][0] target_pts = target['spts'] chamfer = kaolin.metrics.pointcloud.chamfer_distance(pred_points.unsqueeze(0), target_pts.unsqueeze(0)).mean() reg_loss += chamfer return img_loss, reg_loss