Spaces:
Running
on
Zero
Running
on
Zero
File size: 5,234 Bytes
7f51798 |
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 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
# 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
from . import Geometry
from .flexicubes import FlexiCubes # replace later
from .dmtet import sdf_reg_loss_batch
import torch.nn.functional as F
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
def getAABB(self):
return torch.min(self.verts, dim=0).values, torch.max(self.verts, dim=0).values
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
)
return verts, faces, v_reg_loss
def render_mesh(self, mesh_v_nx3, mesh_f_fx3, camera_mv_bx4x4, resolution=256, hierarchical_mask=False):
return_value = dict()
if self.render_type == 'neural_render':
tex_pos, mask, hard_mask, rast, v_pos_clip, mask_pyramid, depth, normal = self.renderer.render_mesh(
mesh_v_nx3.unsqueeze(dim=0),
mesh_f_fx3.int(),
camera_mv_bx4x4,
mesh_v_nx3.unsqueeze(dim=0),
resolution=resolution,
device=self.device,
hierarchical_mask=hierarchical_mask
)
return_value['tex_pos'] = tex_pos
return_value['mask'] = mask
return_value['hard_mask'] = hard_mask
return_value['rast'] = rast
return_value['v_pos_clip'] = v_pos_clip
return_value['mask_pyramid'] = mask_pyramid
return_value['depth'] = depth
return_value['normal'] = normal
else:
raise NotImplementedError
return return_value
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
|