src/models/geometry/rep_3d/flexicubes_geometry.py

# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
import numpy as np
import os
import nvdiffrast.torch as dr
from . import Geometry
from .flexicubes import FlexiCubes # replace later
from .dmtet import sdf_reg_loss_batch
from . import mesh
import torch.nn.functional as F
from src.utils import render

def get_center_boundary_index(grid_res, device):
    v = torch.zeros((grid_res + 1, grid_res + 1, grid_res + 1), dtype=torch.bool, device=device)
    v[grid_res // 2 + 1, grid_res // 2 + 1, grid_res // 2 + 1] = True
    center_indices = torch.nonzero(v.reshape(-1))

    v[grid_res // 2 + 1, grid_res // 2 + 1, grid_res // 2 + 1] = False
    v[:2, ...] = True
    v[-2:, ...] = True
    v[:, :2, ...] = True
    v[:, -2:, ...] = True
    v[:, :, :2] = True
    v[:, :, -2:] = True
    boundary_indices = torch.nonzero(v.reshape(-1))
    return center_indices, boundary_indices

###############################################################################
#  Geometry interface
###############################################################################
class FlexiCubesGeometry(Geometry):
    def __init__(
            self, grid_res=64, scale=2.0, device='cuda', renderer=None,
            render_type='neural_render', args=None):
        super(FlexiCubesGeometry, self).__init__()
        self.grid_res = grid_res
        self.device = device
        self.args = args
        self.fc = FlexiCubes(device, weight_scale=0.5)
        self.verts, self.indices = self.fc.construct_voxel_grid(grid_res)
        if isinstance(scale, list):
            self.verts[:, 0] = self.verts[:, 0] * scale[0]
            self.verts[:, 1] = self.verts[:, 1] * scale[1]
            self.verts[:, 2] = self.verts[:, 2] * scale[1]
        else:
            self.verts = self.verts * scale
            
        all_edges = self.indices[:, self.fc.cube_edges].reshape(-1, 2)
        self.all_edges = torch.unique(all_edges, dim=0)

        # Parameters used for fix boundary sdf
        self.center_indices, self.boundary_indices = get_center_boundary_index(self.grid_res, device)
        self.renderer = renderer
        self.render_type = render_type
        self.ctx = dr.RasterizeCudaContext(device=device)

    def getAABB(self):
        return torch.min(self.verts, dim=0).values, torch.max(self.verts, dim=0).values
    
    @torch.no_grad()
    def map_uv(self, 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")
        )

        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='floor')
            return y * N + x

        tet_idx = _idx(torch.div(face_gidx, N, rounding_mode='floor'), 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
    
    def rotate_x(self, a, device=None):
        s, c = np.sin(a), np.cos(a)
        return torch.tensor([[1, 0, 0, 0], 
                            [0, c,-s, 0], 
                            [0, s, c, 0], 
                         [0, 0, 0, 1]], dtype=torch.float32, device=device)
    def rotate_z(self, a, device=None):
        s, c = np.sin(a), np.cos(a)
        return torch.tensor([[ c, -s, 0, 0],
                            [ s,  c, 0, 0],
                            [ 0,  0, 1, 0],
                            [ 0,  0, 0, 1]], dtype=torch.float32, device=device)
    def rotate_y(self, a, device=None):
        s, c = np.sin(a), np.cos(a)
        return torch.tensor([[ c, 0,  s, 0],
                            [ 0, 1,  0, 0],
                            [-s, 0,  c, 0],
                            [ 0, 0,  0, 1]], dtype=torch.float32, device=device)


    def get_mesh(self, v_deformed_nx3, sdf_n, weight_n=None, with_uv=False, indices=None, is_training=False):
        if indices is None:
            indices = self.indices

        verts, faces, v_reg_loss = self.fc(v_deformed_nx3, sdf_n, indices, self.grid_res,
                                            beta_fx12=weight_n[:, :12], alpha_fx8=weight_n[:, 12:20],
                                            gamma_f=weight_n[:, 20], training=is_training
                                            )
        
        face_gidx = torch.arange(faces.shape[0], dtype=torch.long, device="cuda")
        uvs, uv_idx = self.map_uv(face_gidx, faces.shape[0])
        
        verts = verts @ self.rotate_x(np.pi / 2, device=verts.device)[:3,:3]
        verts = verts @ self.rotate_y(np.pi / 2, device=verts.device)[:3,:3]

        imesh = mesh.Mesh(verts, faces, v_tex=uvs, t_tex_idx=uv_idx)
       
        imesh = mesh.auto_normals(imesh)
        imesh = mesh.compute_tangents(imesh)
        
        return verts, faces, v_reg_loss, imesh


    
    def render_mesh(self, mesh_v_nx3, mesh_f_fx3, mesh, camera_mv_bx4x4, camera_pos, env, planes, kd_fn, materials, resolution=256, hierarchical_mask=False, gt_albedo_map=None, gt_normal_map=None, gt_depth_map=None):
        return_value = dict()
        buffer_dict = render.render_mesh(self.ctx, mesh, camera_mv_bx4x4, camera_pos, env, 
                                         planes, kd_fn, materials, [resolution, resolution], 
                                         spp=1, num_layers=1, msaa=True, background=None, gt_albedo_map=gt_albedo_map)

        return buffer_dict
    

    def render(self, v_deformed_bxnx3=None, sdf_bxn=None, camera_mv_bxnviewx4x4=None, resolution=256):
        # Here I assume a batch of meshes (can be different mesh and geometry), for the other shapes, the batch is 1
        v_list = []
        f_list = []
        n_batch = v_deformed_bxnx3.shape[0]
        all_render_output = []
        for i_batch in range(n_batch):
            verts_nx3, faces_fx3 = self.get_mesh(v_deformed_bxnx3[i_batch], sdf_bxn[i_batch])
            v_list.append(verts_nx3)
            f_list.append(faces_fx3)
            render_output = self.render_mesh(verts_nx3, faces_fx3, camera_mv_bxnviewx4x4[i_batch], resolution)
            all_render_output.append(render_output)

        # Concatenate all render output
        return_keys = all_render_output[0].keys()
        return_value = dict()
        for k in return_keys:
            value = [v[k] for v in all_render_output]
            return_value[k] = value
            # We can do concatenation outside of the render
        return return_value