def extract_mesh(self, occ_hat, z, c=None, stats_dict=dict()):
''' Extracts the mesh from the predicted occupancy grid.
Args:
occ_hat (tensor): value grid of occupancies
z (tensor): latent code z
c (tensor): latent conditioned code c
stats_dict (dict): stats dictionary
'''
# Some short hands
n_x, n_y, n_z = occ_hat.shape
box_size = 1 + self.padding
threshold = np.log(self.threshold) - np.log(1. - self.threshold)
# Make sure that mesh is watertight
t0 = time.time()
occ_hat_padded = np.pad(occ_hat, 1, 'constant', constant_values=-1e6)
vertices, triangles = libmcubes.marching_cubes(occ_hat_padded,
threshold)
stats_dict['time (marching cubes)'] = time.time() - t0
# Strange behaviour in libmcubes: vertices are shifted by 0.5
vertices -= 0.5
# Undo padding
vertices -= 1
# Normalize to bounding box
vertices /= np.array([n_x - 1, n_y - 1, n_z - 1])
vertices = box_size * (vertices - 0.5)
# Estimate normals if needed
if self.with_normals and not vertices.shape[0] == 0:
t0 = time.time()
normals = self.estimate_normals(vertices, z, c)
stats_dict['time (normals)'] = time.time() - t0
else:
normals = None
# Create mesh
mesh = trimesh.Trimesh(vertices,
triangles,
vertex_normals=normals,
process=False)
# Directly return if mesh is empty
if vertices.shape[0] == 0:
return mesh
# TODO: normals are lost here
if self.simplify_nfaces is not None:
t0 = time.time()
mesh = simplify_mesh(mesh, self.simplify_nfaces, 5.)
stats_dict['time (simplify)'] = time.time() - t0
# Refine mesh
if self.refinement_step > 0:
t0 = time.time()
self.refine_mesh(mesh, occ_hat, z, c)
stats_dict['time (refine)'] = time.time() - t0
return mesh
def extract_mesh(self,
occ_hat,
label_hat,
p_hat,
c,
stats_dict=dict(),
**kwargs):
''' Extracts the mesh from the predicted occupancy grid.
Args:
occ_hat (tensor): value grid of occupancies
label_hat (tensor): value grid of predicted part labels
p_hat (tensor): value grid of predicted locations in the A-pose
c (tensor): latent conditioned code c
stats_dict (dict): stats dictionary
'''
# Some short hands
n_x, n_y, n_z = occ_hat.shape
box_size = 1 + self.padding
threshold = np.log(self.threshold) - np.log(1. - self.threshold)
# Make sure that mesh is watertight
t0 = time.time()
occ_hat_padded = np.pad(occ_hat, 1, 'constant', constant_values=-1e6)
vertices, triangles = libmcubes.marching_cubes(occ_hat_padded,
threshold)
stats_dict['time (marching cubes)'] = time.time() - t0
# Strange behaviour in libmcubes: vertices are shifted by 0.5
vertices -= 0.5
# Undo padding
vertices -= 1
# Construct part labels and A-pose vertices by sampling
# occupancy grid and translation grid
r_verts = np.round(vertices).astype('int32')
labels = label_hat[r_verts[:, 0], r_verts[:, 1], r_verts[:, 2]]
colors = self.colors[labels]
# Normalize to bounding box
vertices /= np.array([n_x - 1, n_y - 1, n_z - 1])
if p_hat is not None:
with torch.no_grad():
v = torch.from_numpy(vertices).to(self.device).float()
v = v * 2 - 1 # range: [-1, 1]
v = v.unsqueeze(0).unsqueeze(1).unsqueeze(
1) # 1 x 1 x 1 x n_pts x 3
p_hat = torch.from_numpy(p_hat).to(self.device).float()
p_hat = p_hat.permute(3, 2, 1,
0).unsqueeze(0) # 1 X C x D x H x W
# v_rest is in [-1, 1]
v_rest = torch.nn.functional.grid_sample(p_hat,
v,
align_corners=True)
v_rest = v_rest.squeeze(0).squeeze(1).squeeze(1).transpose(
0, 1) / 1.5 * kwargs['scale'] # + kwargs['loc']
vertices_rest = v_rest.detach().cpu().numpy()
else:
vertices_rest = None
# vertices = box_size * (vertices - 0.5)
vertices = 4 / 3 * kwargs['scale'].item() * (
vertices - 0.5) + kwargs['loc'].cpu().numpy()
# Estimate normals if needed
if self.with_normals and not vertices.shape[0] == 0:
t0 = time.time()
normals = self.estimate_normals(vertices, c)
stats_dict['time (normals)'] = time.time() - t0
else:
normals = None
# Create mesh
mesh = {}
mesh['part_labels'] = labels
mesh['posed'] = trimesh.Trimesh(vertices,
triangles,
vertex_normals=normals,
vertex_colors=colors,
process=False)
if vertices_rest is not None:
mesh['unposed'] = trimesh.Trimesh(vertices_rest,
triangles,
vertex_normals=normals,
vertex_colors=colors,
process=False)
# Directly return if mesh is empty
if vertices.shape[0] == 0:
return mesh
# TODO: normals are lost here
if self.simplify_nfaces is not None:
t0 = time.time()
mesh = simplify_mesh(mesh, self.simplify_nfaces, 5.)
stats_dict['time (simplify)'] = time.time() - t0
# Refine mesh
if self.refinement_step > 0:
t0 = time.time()
self.refine_mesh(mesh, occ_hat, c)
stats_dict['time (refine)'] = time.time() - t0
return mesh