MeshDiffusion/nvdiffrec/lib/geometry/dmtet_singleview.py

# 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 os
import numpy as np
import torch

from ..render import mesh
from ..render import render
from ..render import regularizer


import kaolin
import pytorch3d.ops
from ..render import util as render_utils

import torch.nn.functional as F

from ..render import renderutils as ru

###############################################################################
# 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):
        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)


        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)

        valid_vert_idx = tet_fx4[tet_gidx[num_triangles > 0]].long().unique()

        return verts, faces, uvs, uv_idx, face_to_valid_tet.long(), valid_vert_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



class Buffer(object):
    def __init__(self, shape, capacity, device) -> None:
        self.len_curr = 0
        self.pointer = 0
        self.capacity = capacity
        self.buffer = torch.zeros((capacity, ) + shape, device=device)
    
    def push(self, x):
        '''
            Push one single data point into the buffer
        '''
        self.buffer[self.pointer] = x
        self.pointer = (self.pointer + 1) % self.capacity
        if self.len_curr < self.capacity:
            self.len_curr += 1
    
    def avg(self):
        return torch.sign(torch.sign(self.buffer[:self.len_curr]).float().mean(dim=0)).float()
        # return self.buffer[:self.len_curr].mean(dim=0)
        # return self.buffer[:self.len_curr][-1]

###############################################################################
#  Geometry interface
###############################################################################

class DMTetGeometry(torch.nn.Module):
    def __init__(self, grid_res, scale, FLAGS, root='./', grid_to_tet=None, deform_scale=2.0, **kwargs):
        super(DMTetGeometry, self).__init__()

        self.FLAGS         = FLAGS
        self.grid_res      = grid_res
        self.marching_tets = DMTet()
        self.cropped       = True
        self.tanh          = False
        self.deform_scale  = deform_scale

        self.grid_to_tet = grid_to_tet

        if self.cropped:
            print("use cropped tets")
            tets = np.load(os.path.join(root, 'data/tets/{}_tets_cropped.npz'.format(self.grid_res)))
        else:
            tets = np.load(os.path.join(root, 'data/tets/{}_tets.npz'.format(self.grid_res)))
        print('tet min and max', tets['vertices'].min() * scale, tets['vertices'].max() * scale)
        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()

        # Random init
        sdf = torch.rand_like(self.verts[:,0]).clamp(-1.0, 1.0) - 0.1
        # sdf = torch.sign(sdf) * 0.1
        # sdf = self.verts.pow(2).sum(dim=-1).sqrt() - 0.5

        self.sdf    = torch.nn.Parameter(sdf.clone().detach(), requires_grad=True)
        self.register_parameter('sdf', self.sdf)

        self.deform = torch.nn.Parameter(torch.zeros_like(self.verts), requires_grad=True)
        self.register_parameter('deform', self.deform)

        self.alpha = None

        self.sdf_ema    = torch.nn.Parameter(sdf.clone().detach(), requires_grad=False)
        self.deform_ema = torch.nn.Parameter(self.deform.clone().detach(), requires_grad=False)

        # self.ema_coeff = 0.7
        self.ema_coeff = 0.9

        self.sdf_buffer = Buffer(sdf.size(), capacity=200, device='cuda')

    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):
        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 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 getValidTetIdx(self):
        # Run DM tet to get a base mesh
        v_deformed = self.get_deformed()
        verts, faces, uvs, uv_idx, tet_gidx, valid_vert_idx = self.marching_tets(v_deformed, self.sdf, self.indices)
        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, valid_vert_idx = self.marching_tets(v_deformed, self.sdf, self.indices)
        return self.indices[tet_gidx.long()].unique()

    def getMesh(self, material, noise=0.0, ema=False):
        # Run DM tet to get a base mesh
        v_deformed = self.get_deformed(ema=ema)

        if ema:
            sdf = self.sdf_ema
        else:
            sdf = self.sdf

        verts, faces, uvs, uv_idx, tet_gidx, valid_vert_idx = self.marching_tets(v_deformed, sdf, self.indices)
        imesh = mesh.Mesh(verts, faces, v_tex=uvs, t_tex_idx=uv_idx, material=material)

        if material is not None:
            # Run mesh operations to generate tangent space
            imesh = mesh.auto_normals(imesh)
            imesh = mesh.compute_tangents(imesh)
        imesh.valid_vert_idx = valid_vert_idx

        return imesh

    def getMesh_no_deform(self, material, noise=0.0, ema=False):
        # Run DM tet to get a base mesh
        if ema:
            # sdf = self.sdf * (1 - self.ema_coeff) + self.sdf_ema.detach() * self.ema_coeff
            sdf = self.sdf_ema
        else:
            sdf = self.sdf

        verts, faces, uvs, uv_idx, tet_gidx, valid_vert_idx = self.marching_tets(self.verts, torch.sign(sdf), 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_no_deform_gd(self, material, noise=0.0, ema=False):
        # Run DM tet to get a base mesh
        v_deformed = self.get_deformed(no_grad=True)


        if ema:
            # sdf = self.sdf * (1 - self.ema_coeff) + self.sdf_ema.detach() * self.ema_coeff
            sdf = self.sdf_ema
        else:
            sdf = self.sdf

        verts, faces, uvs, uv_idx, tet_gidx, valid_vert_idx = self.marching_tets(v_deformed, sdf, 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 get_deformed(self, no_grad=False, ema=False):
        if no_grad:
            deform = self.deform.detach()
        else:
            deform = self.deform

        if self.tanh:
            v_deformed = self.verts + 2 / (self.grid_res * 2) * torch.tanh(deform) * self.deform_scale
        else:
            v_deformed = self.verts + 2 / (self.grid_res * 2) * deform * self.deform_scale
        return v_deformed
    
    def get_angle(self):
        with torch.no_grad():
            comb_list = [
                (0, 1, 2, 3),
                (0, 1, 3, 2),
                (0, 2, 3, 1),
                (1, 2, 3, 0)
            ]

            directions = torch.zeros(self.indices.size(0), 4).cuda()
            dir_vec = torch.zeros(self.indices.size(0), 4, 3).cuda()
            vert_inds = torch.zeros(self.indices.size(0), 4).cuda().long()
            count = 0
            vpos_list = self.get_deformed()
            for comb in comb_list:
                face = self.indices[:, comb[:3]]
                face_pos = vpos_list[face, :]
                face_center = face_pos.mean(1, keepdim=False)
                v = self.indices[:, comb[3]]
                test_vec = vpos_list[v]
                ref_vec = render_utils.safe_normalize(vpos_list[face[:, 0]] - face_center)
                distance_vec = test_vec - render_utils.dot(test_vec, ref_vec) * ref_vec
                directions[:, count] = torch.sign(render_utils.dot(test_vec, distance_vec)[:, 0])
                dir_vec[:, count, :] = distance_vec
                vert_inds[:, count] = v
                count += 1
            return directions, dir_vec, vert_inds


    def clamp_deform(self):
        if not self.tanh:
            self.deform.data[:] = self.deform.data.clamp(-0.99, 0.99)
            self.sdf.data[:] = self.sdf.data.clamp(-1.0, 1.0)
 
    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, ema=ema)
        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, noise=0.0, ema=False, xfm_lgt=None):
        opt_mesh = self.getMesh(opt_material, noise=noise, ema=ema)
        return opt_mesh, 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)
    
    def update_ema(self, ema_coeff=0.9):
        self.sdf_buffer.push(self.sdf)
        self.sdf_ema.data[:] = self.sdf_buffer.avg()
        # self.sdf_ema.data[:] = self.sdf.data[:] * (1 - ema_coeff) + self.sdf_ema.data[:] * ema_coeff
        # self.deform_ema.data[:] = self.deform.data[:] * (1 - ema_coeff) + self.deform_ema.data[:] * ema_coeff
        self.deform_ema.data[:] = self.deform.data[:]


    def render_ema(self, glctx, target, lgt, opt_material, bsdf=None, xfm_lgt=None):
        opt_mesh = self.getMesh(opt_material, ema=True)
        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)

    def init_with_gt_surface(self, gt_verts, surface_faces, campos):
        with torch.no_grad():
            surface_face_verts = gt_verts[surface_faces]
            surface_centers = surface_face_verts.mean(dim=1)
            v_pos = self.get_deformed()
            results = pytorch3d.ops.knn_points(v_pos[None, ...], surface_centers[None, ...])
            dists, nn_idx = results.dists, results.idx
            displacement = (v_pos - surface_centers[nn_idx[0, :, 0]])
            view_dirs = campos - surface_centers
            normals = torch.cross(
                surface_face_verts[:, 0] - surface_face_verts[:, 1], surface_face_verts[:, 0] - surface_face_verts[:, 2])
            mask = ((normals * view_dirs).sum(dim=-1, keepdim=True) >= 0).float()
            normals = normals * mask - normals * (1 - mask)
            outside_verts_idx = ((displacement * normals[nn_idx[0, :, 0]]).sum(dim=-1) > 0)
            self.sdf.data[outside_verts_idx] = 1.0
            

    def tick(self, glctx, target, lgt, opt_material, loss_fn, iteration, with_reg=True, xfm_lgt=None, no_depth_thin=True):

        if iteration < 100:
            self.deform.requires_grad = False
            self.deform_scale = 2.0
        else:
            self.deform.requires_grad = True
            self.deform_scale = 2.0
        
        if iteration > 200 and iteration < 2000 and iteration % 20 == 0:
            with torch.no_grad():
                v_pos = self.get_deformed()
                v_pos_camera_homo = ru.xfm_points(v_pos[None, ...], target['mvp'])
                v_pos_camera = v_pos_camera_homo[:, :, :2] / v_pos_camera_homo[:, :, -1:]
                v_pos_camera_discrete = ((v_pos_camera * 0.5 + 0.5).clip(0, 1) * (target['resolution'][0] - 1)).long()
                target_mask = target['mask_cont'][:, :, :, 0] == 0
                for k in range(target_mask.size(0)):
                    assert v_pos_camera_discrete[k].min() >= 0 and v_pos_camera_discrete[k].max() < target['resolution'][0]
                    v_mask = target_mask[k, v_pos_camera_discrete[k, :, 1], v_pos_camera_discrete[k, :, 0]].view(v_pos.size(0))
                    # print(v_mask.sum())
                    self.sdf.data[v_mask] = self.sdf.data[v_mask].abs().clamp(0.0, 1.0)

        # ==============================================================================================
        #  Render optimizable object with identical conditions
        # ==============================================================================================
        imesh, buffers = self.render_with_mesh(glctx, target, lgt, opt_material, noise=0.0, xfm_lgt=xfm_lgt)

        # ==============================================================================================
        #  Compute loss
        # ==============================================================================================

        # Image-space loss, split into a coverage component and a color component
        color_ref = target['img']
        img_loss = torch.tensor(0.0).cuda()
        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_cont'][:, :, :, 0] == 1.0).float()
        mask_curr = (buffers['mask_cont'][:, :, :, 0] == 1.0).float()

        if iteration % 300 == 0 and iteration < 1790:
            self.deform.data[:] *= 0.4

        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
        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.0).cuda()

        # SDF regularizer
        iter_thres = 0
        sdf_weight = self.FLAGS.sdf_regularizer - (self.FLAGS.sdf_regularizer - 0.01) * min(1.0, 4.0 * ((iteration - iter_thres) / (self.FLAGS.iter - iter_thres)))

        sdf_mask = torch.zeros_like(self.sdf, device=self.sdf.device)
        sdf_mask[imesh.valid_vert_idx] = 1.0
        sdf_masked = self.sdf.detach() * sdf_mask + self.sdf * (1 - sdf_mask)
        reg_loss = sdf_reg_loss(sdf_masked, self.all_edges).mean() * sdf_weight * 2.5 # Dropoff to 0.01


        # Albedo (k_d) smoothnesss regularizer
        reg_loss += torch.mean(buffers['kd_grad'][..., :-1] * buffers['kd_grad'][..., -1:]) * 0.03 * min(1.0, iteration / 500)

        # Visibility regularizer
        reg_loss += torch.mean(buffers['occlusion'][..., :-1] * buffers['occlusion'][..., -1:]) * 1e0 * min(1.0, iteration / 500)

        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