| from typing import Union, Tuple, List |
|
|
| import numpy as np |
| import torch |
| from skimage import measure |
|
|
|
|
| class MeshExtractResult: |
| def __init__(self, verts, faces, vertex_attrs=None, res=64): |
| self.verts = verts |
| self.faces = faces.long() |
| self.vertex_attrs = vertex_attrs |
| self.face_normal = self.comput_face_normals() |
| self.vert_normal = self.comput_v_normals() |
| self.res = res |
| self.success = verts.shape[0] != 0 and faces.shape[0] != 0 |
|
|
| |
| self.tsdf_v = None |
| self.tsdf_s = None |
| self.reg_loss = None |
|
|
| def comput_face_normals(self): |
| i0 = self.faces[..., 0].long() |
| i1 = self.faces[..., 1].long() |
| i2 = self.faces[..., 2].long() |
|
|
| v0 = self.verts[i0, :] |
| v1 = self.verts[i1, :] |
| v2 = self.verts[i2, :] |
| face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1) |
| face_normals = torch.nn.functional.normalize(face_normals, dim=1) |
| return face_normals[:, None, :].repeat(1, 3, 1) |
|
|
| def comput_v_normals(self): |
| i0 = self.faces[..., 0].long() |
| i1 = self.faces[..., 1].long() |
| i2 = self.faces[..., 2].long() |
|
|
| v0 = self.verts[i0, :] |
| v1 = self.verts[i1, :] |
| v2 = self.verts[i2, :] |
| face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1) |
| v_normals = torch.zeros_like(self.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 |
|
|
|
|
| def center_vertices(vertices): |
| """Translate the vertices so that bounding box is centered at zero.""" |
| vert_min = vertices.min(dim=0)[0] |
| vert_max = vertices.max(dim=0)[0] |
| vert_center = 0.5 * (vert_min + vert_max) |
| return vertices - vert_center |
|
|
|
|
| class SurfaceExtractor: |
| def _compute_box_stat( |
| self, bounds: Union[Tuple[float], List[float], float], octree_resolution: int |
| ): |
| if isinstance(bounds, float): |
| bounds = [-bounds, -bounds, -bounds, bounds, bounds, bounds] |
|
|
| bbox_min, bbox_max = np.array(bounds[0:3]), np.array(bounds[3:6]) |
| bbox_size = bbox_max - bbox_min |
| grid_size = [ |
| int(octree_resolution) + 1, |
| int(octree_resolution) + 1, |
| int(octree_resolution) + 1, |
| ] |
| return grid_size, bbox_min, bbox_size |
|
|
| def run(self, *args, **kwargs): |
| return NotImplementedError |
|
|
| def __call__(self, grid_logits, **kwargs): |
| outputs = [] |
| for i in range(grid_logits.shape[0]): |
| try: |
| verts, faces = self.run(grid_logits[i], **kwargs) |
| outputs.append( |
| MeshExtractResult( |
| verts=verts.float(), |
| faces=faces, |
| res=kwargs["octree_resolution"], |
| ) |
| ) |
|
|
| except Exception: |
| import traceback |
|
|
| traceback.print_exc() |
| outputs.append(None) |
|
|
| return outputs |
|
|
|
|
| class MCSurfaceExtractor(SurfaceExtractor): |
| def run(self, grid_logit, *, mc_level, bounds, octree_resolution, **kwargs): |
| verts, faces, normals, _ = measure.marching_cubes( |
| grid_logit.float().cpu().numpy(), mc_level, method="lewiner" |
| ) |
| grid_size, bbox_min, bbox_size = self._compute_box_stat( |
| bounds, octree_resolution |
| ) |
| verts = verts / grid_size * bbox_size + bbox_min |
| verts = torch.tensor(verts, device=grid_logit.device, dtype=torch.float32) |
| faces = torch.tensor( |
| np.ascontiguousarray(faces), device=grid_logit.device, dtype=torch.long |
| ) |
| faces = faces[:, [2, 1, 0]] |
| return verts, faces |
|
|
|
|
| class DMCSurfaceExtractor(SurfaceExtractor): |
| def run(self, grid_logit, *, octree_resolution, **kwargs): |
| device = grid_logit.device |
| if not hasattr(self, "dmc"): |
| try: |
| from diso import DiffDMC |
| except: |
| raise ImportError( |
| "Please install diso via `pip install diso`, or set mc_algo to 'mc'" |
| ) |
| self.dmc = DiffDMC(dtype=torch.float32).to(device) |
| sdf = -grid_logit / octree_resolution |
| sdf = sdf.to(torch.float32).contiguous() |
| verts, faces = self.dmc(sdf, deform=None, return_quads=False, normalize=True) |
| grid_size, bbox_min, bbox_size = self._compute_box_stat( |
| kwargs["bounds"], octree_resolution |
| ) |
| verts = verts * kwargs["bounds"] * 2 - kwargs["bounds"] |
| return verts, faces |
|
|
