Spaces:
Runtime error
Runtime error
| # Copyright (c) 2023, 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 | |
| from ...modules.sparse import SparseTensor | |
| from easydict import EasyDict as edict | |
| from .utils_cube import * | |
| try: | |
| from .flexicube import FlexiCubes | |
| except: | |
| print("Please install kaolin and diso to use the mesh extractor.") | |
| class MeshExtractResult: | |
| def __init__(self, | |
| vertices, | |
| faces, | |
| vertex_attrs=None, | |
| res=64 | |
| ): | |
| self.vertices = vertices | |
| self.faces = faces.long() | |
| self.vertex_attrs = vertex_attrs | |
| self.face_normal = self.comput_face_normals(vertices, faces) | |
| self.res = res | |
| self.success = (vertices.shape[0] != 0 and faces.shape[0] != 0) | |
| # training only | |
| self.tsdf_v = None | |
| self.tsdf_s = None | |
| self.reg_loss = None | |
| def comput_face_normals(self, verts, faces): | |
| i0 = faces[..., 0].long() | |
| i1 = faces[..., 1].long() | |
| i2 = faces[..., 2].long() | |
| v0 = verts[i0, :] | |
| v1 = verts[i1, :] | |
| v2 = verts[i2, :] | |
| face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1) | |
| face_normals = torch.nn.functional.normalize(face_normals, dim=1) | |
| # print(face_normals.min(), face_normals.max(), face_normals.shape) | |
| return face_normals[:, None, :].repeat(1, 3, 1) | |
| def comput_v_normals(self, verts, faces): | |
| i0 = faces[..., 0].long() | |
| i1 = faces[..., 1].long() | |
| i2 = faces[..., 2].long() | |
| v0 = verts[i0, :] | |
| v1 = verts[i1, :] | |
| v2 = verts[i2, :] | |
| face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1) | |
| v_normals = torch.zeros_like(verts) | |
| v_normals.scatter_add_(0, i0[..., None].repeat(1, 3), face_normals) | |
| v_normals.scatter_add_(0, i1[..., None].repeat(1, 3), face_normals) | |
| v_normals.scatter_add_(0, i2[..., None].repeat(1, 3), face_normals) | |
| v_normals = torch.nn.functional.normalize(v_normals, dim=1) | |
| return v_normals | |
| class SparseFeatures2Mesh: | |
| def __init__(self, device="cuda", res=64, use_color=True): | |
| ''' | |
| a model to generate a mesh from sparse features structures using flexicube | |
| ''' | |
| super().__init__() | |
| self.device=device | |
| self.res = res | |
| self.mesh_extractor = FlexiCubes(device=device) | |
| self.sdf_bias = -1.0 / res | |
| verts, cube = construct_dense_grid(self.res, self.device) | |
| self.reg_c = cube.to(self.device) | |
| self.reg_v = verts.to(self.device) | |
| self.use_color = use_color | |
| self._calc_layout() | |
| def _calc_layout(self): | |
| LAYOUTS = { | |
| 'sdf': {'shape': (8, 1), 'size': 8}, | |
| 'deform': {'shape': (8, 3), 'size': 8 * 3}, | |
| 'weights': {'shape': (21,), 'size': 21} | |
| } | |
| if self.use_color: | |
| ''' | |
| 6 channel color including normal map | |
| ''' | |
| LAYOUTS['color'] = {'shape': (8, 6,), 'size': 8 * 6} | |
| self.layouts = edict(LAYOUTS) | |
| start = 0 | |
| for k, v in self.layouts.items(): | |
| v['range'] = (start, start + v['size']) | |
| start += v['size'] | |
| self.feats_channels = start | |
| def get_layout(self, feats : torch.Tensor, name : str): | |
| if name not in self.layouts: | |
| return None | |
| return feats[:, self.layouts[name]['range'][0]:self.layouts[name]['range'][1]].reshape(-1, *self.layouts[name]['shape']) | |
| def __call__(self, cubefeats : SparseTensor, training=False): | |
| """ | |
| Generates a mesh based on the specified sparse voxel structures. | |
| Args: | |
| cube_attrs [Nx21] : Sparse Tensor attrs about cube weights | |
| verts_attrs [Nx10] : [0:1] SDF [1:4] deform [4:7] color [7:10] normal | |
| Returns: | |
| return the success tag and ni you loss, | |
| """ | |
| # add sdf bias to verts_attrs | |
| coords = cubefeats.coords[:, 1:] | |
| feats = cubefeats.feats | |
| sdf, deform, color, weights = [self.get_layout(feats, name) for name in ['sdf', 'deform', 'color', 'weights']] | |
| sdf += self.sdf_bias | |
| v_attrs = [sdf, deform, color] if self.use_color else [sdf, deform] | |
| v_pos, v_attrs, reg_loss = sparse_cube2verts(coords, torch.cat(v_attrs, dim=-1), training=training) | |
| v_attrs_d = get_dense_attrs(v_pos, v_attrs, res=self.res+1, sdf_init=True) | |
| weights_d = get_dense_attrs(coords, weights, res=self.res, sdf_init=False) | |
| if self.use_color: | |
| sdf_d, deform_d, colors_d = v_attrs_d[..., 0], v_attrs_d[..., 1:4], v_attrs_d[..., 4:] | |
| else: | |
| sdf_d, deform_d = v_attrs_d[..., 0], v_attrs_d[..., 1:4] | |
| colors_d = None | |
| x_nx3 = get_defomed_verts(self.reg_v, deform_d, self.res) | |
| vertices, faces, L_dev, colors = self.mesh_extractor( | |
| voxelgrid_vertices=x_nx3, | |
| scalar_field=sdf_d, | |
| cube_idx=self.reg_c, | |
| resolution=self.res, | |
| beta=weights_d[:, :12], | |
| alpha=weights_d[:, 12:20], | |
| gamma_f=weights_d[:, 20], | |
| voxelgrid_colors=colors_d, | |
| training=training) | |
| mesh = MeshExtractResult(vertices=vertices, faces=faces, vertex_attrs=colors, res=self.res) | |
| if training: | |
| if mesh.success: | |
| reg_loss += L_dev.mean() * 0.5 | |
| reg_loss += (weights[:,:20]).abs().mean() * 0.2 | |
| mesh.reg_loss = reg_loss | |
| mesh.tsdf_v = get_defomed_verts(v_pos, v_attrs[:, 1:4], self.res) | |
| mesh.tsdf_s = v_attrs[:, 0] | |
| return mesh | |