Spaces:
Running
on
Zero
Running
on
Zero
| # | |
| # Copyright (C) 2024, ShanghaiTech | |
| # SVIP research group, https://github.com/svip-lab | |
| # All rights reserved. | |
| # | |
| # This software is free for non-commercial, research and evaluation use | |
| # under the terms of the LICENSE.md file. | |
| # | |
| # For inquiries contact huangbb@shanghaitech.edu.cn | |
| # | |
| import numpy as np | |
| import torch | |
| import trimesh | |
| from skimage import measure | |
| # modified from here https://github.com/autonomousvision/sdfstudio/blob/370902a10dbef08cb3fe4391bd3ed1e227b5c165/nerfstudio/utils/marching_cubes.py#L201 | |
| def marching_cubes_with_contraction( | |
| sdf, | |
| resolution=512, | |
| bounding_box_min=(-1.0, -1.0, -1.0), | |
| bounding_box_max=(1.0, 1.0, 1.0), | |
| return_mesh=False, | |
| level=0, | |
| simplify_mesh=True, | |
| inv_contraction=None, | |
| max_range=32.0, | |
| ): | |
| assert resolution % 512 == 0 | |
| resN = resolution | |
| cropN = 512 | |
| level = 0 | |
| N = resN // cropN | |
| grid_min = bounding_box_min | |
| grid_max = bounding_box_max | |
| xs = np.linspace(grid_min[0], grid_max[0], N + 1) | |
| ys = np.linspace(grid_min[1], grid_max[1], N + 1) | |
| zs = np.linspace(grid_min[2], grid_max[2], N + 1) | |
| meshes = [] | |
| for i in range(N): | |
| for j in range(N): | |
| for k in range(N): | |
| print(i, j, k) | |
| x_min, x_max = xs[i], xs[i + 1] | |
| y_min, y_max = ys[j], ys[j + 1] | |
| z_min, z_max = zs[k], zs[k + 1] | |
| x = np.linspace(x_min, x_max, cropN) | |
| y = np.linspace(y_min, y_max, cropN) | |
| z = np.linspace(z_min, z_max, cropN) | |
| xx, yy, zz = np.meshgrid(x, y, z, indexing="ij") | |
| points = torch.tensor(np.vstack([xx.ravel(), yy.ravel(), zz.ravel()]).T, dtype=torch.float).cuda() | |
| def evaluate(points): | |
| z = [] | |
| for _, pnts in enumerate(torch.split(points, 256**3, dim=0)): | |
| z.append(sdf(pnts)) | |
| z = torch.cat(z, axis=0) | |
| return z | |
| # construct point pyramids | |
| points = points.reshape(cropN, cropN, cropN, 3) | |
| points = points.reshape(-1, 3) | |
| pts_sdf = evaluate(points.contiguous()) | |
| z = pts_sdf.detach().cpu().numpy() | |
| if not (np.min(z) > level or np.max(z) < level): | |
| z = z.astype(np.float32) | |
| verts, faces, normals, _ = measure.marching_cubes( | |
| volume=z.reshape(cropN, cropN, cropN), | |
| level=level, | |
| spacing=( | |
| (x_max - x_min) / (cropN - 1), | |
| (y_max - y_min) / (cropN - 1), | |
| (z_max - z_min) / (cropN - 1), | |
| ), | |
| ) | |
| verts = verts + np.array([x_min, y_min, z_min]) | |
| meshcrop = trimesh.Trimesh(verts, faces, normals) | |
| meshes.append(meshcrop) | |
| print("finished one block") | |
| combined = trimesh.util.concatenate(meshes) | |
| combined.merge_vertices(digits_vertex=6) | |
| # inverse contraction and clipping the points range | |
| if inv_contraction is not None: | |
| combined.vertices = inv_contraction(torch.from_numpy(combined.vertices).float().cuda()).cpu().numpy() | |
| combined.vertices = np.clip(combined.vertices, -max_range, max_range) | |
| return combined |