Spaces:
Build error
Build error
import numpy as np | |
def create_grid(resX, resY, resZ, b_min=np.array([0, 0, 0]), b_max=np.array([1, 1, 1]), transform=None): | |
''' | |
Create a dense grid of given resolution and bounding box | |
:param resX: resolution along X axis | |
:param resY: resolution along Y axis | |
:param resZ: resolution along Z axis | |
:param b_min: vec3 (x_min, y_min, z_min) bounding box corner | |
:param b_max: vec3 (x_max, y_max, z_max) bounding box corner | |
:return: [3, resX, resY, resZ] coordinates of the grid, and transform matrix from mesh index | |
''' | |
coords = np.mgrid[:resX, :resY, :resZ] | |
coords = coords.reshape(3, -1) | |
coords_matrix = np.eye(4) | |
length = b_max - b_min | |
coords_matrix[0, 0] = length[0] / resX | |
coords_matrix[1, 1] = length[1] / resY | |
coords_matrix[2, 2] = length[2] / resZ | |
coords_matrix[0:3, 3] = b_min | |
coords = np.matmul(coords_matrix[:3, :3], coords) + coords_matrix[:3, 3:4] | |
if transform is not None: | |
coords = np.matmul(transform[:3, :3], coords) + transform[:3, 3:4] | |
coords_matrix = np.matmul(transform, coords_matrix) | |
coords = coords.reshape(3, resX, resY, resZ) | |
return coords, coords_matrix | |
def batch_eval(points, eval_func, num_samples=512 * 512 * 512): | |
num_pts = points.shape[1] | |
sdf = np.zeros(num_pts) | |
num_batches = num_pts // num_samples | |
for i in range(num_batches): | |
sdf[i * num_samples:i * num_samples + num_samples] = eval_func( | |
points[:, i * num_samples:i * num_samples + num_samples]) | |
if num_pts % num_samples: | |
sdf[num_batches * num_samples:] = eval_func(points[:, num_batches * num_samples:]) | |
return sdf | |
def eval_grid(coords, eval_func, num_samples=512 * 512 * 512): | |
resolution = coords.shape[1:4] | |
coords = coords.reshape([3, -1]) | |
sdf = batch_eval(coords, eval_func, num_samples=num_samples) | |
return sdf.reshape(resolution) | |
def eval_grid_octree(coords, eval_func, | |
init_resolution=64, threshold=0.01, | |
num_samples=512 * 512 * 512): | |
resolution = coords.shape[1:4] | |
sdf = np.zeros(resolution) | |
dirty = np.ones(resolution, dtype=np.bool) | |
grid_mask = np.zeros(resolution, dtype=np.bool) | |
reso = resolution[0] // init_resolution | |
while reso > 0: | |
# subdivide the grid | |
grid_mask[0:resolution[0]:reso, 0:resolution[1]:reso, 0:resolution[2]:reso] = True | |
# test samples in this iteration | |
test_mask = np.logical_and(grid_mask, dirty) | |
#print('step size:', reso, 'test sample size:', test_mask.sum()) | |
points = coords[:, test_mask] | |
sdf[test_mask] = batch_eval(points, eval_func, num_samples=num_samples) | |
dirty[test_mask] = False | |
# do interpolation | |
if reso <= 1: | |
break | |
for x in range(0, resolution[0] - reso, reso): | |
for y in range(0, resolution[1] - reso, reso): | |
for z in range(0, resolution[2] - reso, reso): | |
# if center marked, return | |
if not dirty[x + reso // 2, y + reso // 2, z + reso // 2]: | |
continue | |
v0 = sdf[x, y, z] | |
v1 = sdf[x, y, z + reso] | |
v2 = sdf[x, y + reso, z] | |
v3 = sdf[x, y + reso, z + reso] | |
v4 = sdf[x + reso, y, z] | |
v5 = sdf[x + reso, y, z + reso] | |
v6 = sdf[x + reso, y + reso, z] | |
v7 = sdf[x + reso, y + reso, z + reso] | |
v = np.array([v0, v1, v2, v3, v4, v5, v6, v7]) | |
v_min = v.min() | |
v_max = v.max() | |
# this cell is all the same | |
if (v_max - v_min) < threshold: | |
sdf[x:x + reso, y:y + reso, z:z + reso] = (v_max + v_min) / 2 | |
dirty[x:x + reso, y:y + reso, z:z + reso] = False | |
reso //= 2 | |
return sdf.reshape(resolution) | |