Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,461 Bytes
476e0f0 |
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 |
#
# 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()
@torch.no_grad()
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 |