180 lines
6.1 KiB
Python
180 lines
6.1 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
import os
|
|
|
|
import mcubes
|
|
import numpy as np
|
|
import torch
|
|
|
|
|
|
def save_obj_mesh_with_color(mesh_path, verts, faces, colors):
|
|
file = open(mesh_path, 'w')
|
|
for idx, v in enumerate(verts):
|
|
c = colors[idx]
|
|
file.write('v %.4f %.4f %.4f %.4f %.4f %.4f\n' %
|
|
(v[0], v[1], v[2], c[0], c[1], c[2]))
|
|
for f in faces:
|
|
f_plus = f + 1
|
|
file.write('f %d %d %d\n' % (f_plus[0], f_plus[2], f_plus[1]))
|
|
file.close()
|
|
|
|
|
|
def save_obj_mesh(mesh_path, verts, faces):
|
|
file = open(mesh_path, 'w')
|
|
for idx, v in enumerate(verts):
|
|
file.write('v %.4f %.4f %.4f\n' % (v[0], v[1], v[2]))
|
|
for f in faces:
|
|
f_plus = f + 1
|
|
file.write('f %d %d %d\n' % (f_plus[0], f_plus[2], f_plus[1]))
|
|
file.close()
|
|
|
|
|
|
def to_tensor(img):
|
|
if len(img.shape) == 2:
|
|
img = img[:, :, np.newaxis]
|
|
img = torch.from_numpy(img.transpose(2, 0, 1)).float()
|
|
img = img / 255.
|
|
return img
|
|
|
|
|
|
def reconstruction(net, calib_tensor, coords, mat, num_samples=50000):
|
|
|
|
def eval_func(points):
|
|
points = np.expand_dims(points, axis=0)
|
|
points = np.repeat(points, 1, axis=0)
|
|
samples = torch.from_numpy(points).cuda().float()
|
|
net.query(samples, calib_tensor)
|
|
pred = net.get_preds()
|
|
pred = pred[0]
|
|
return pred.detach().cpu().numpy()
|
|
|
|
sdf = eval_grid(coords, eval_func, num_samples=num_samples)
|
|
vertices, faces = mcubes.marching_cubes(sdf, 0.5)
|
|
verts = np.matmul(mat[:3, :3], vertices.T) + mat[:3, 3:4]
|
|
verts = verts.T
|
|
return verts, faces
|
|
|
|
|
|
def keep_largest(mesh_big):
|
|
mesh_lst = mesh_big.split(only_watertight=False)
|
|
keep_mesh = mesh_lst[0]
|
|
for mesh in mesh_lst:
|
|
if mesh.vertices.shape[0] > keep_mesh.vertices.shape[0]:
|
|
keep_mesh = mesh
|
|
return keep_mesh
|
|
|
|
|
|
def eval_grid(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=bool)
|
|
grid_mask = np.zeros(resolution, dtype=bool)
|
|
reso = resolution[0] // init_resolution
|
|
|
|
while reso > 0:
|
|
grid_mask[0:resolution[0]:reso, 0:resolution[1]:reso,
|
|
0:resolution[2]:reso] = True
|
|
test_mask = np.logical_and(grid_mask, dirty)
|
|
points = coords[:, test_mask]
|
|
|
|
sdf[test_mask] = batch_eval(points, eval_func, num_samples=num_samples)
|
|
dirty[test_mask] = False
|
|
|
|
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 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()
|
|
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)
|
|
|
|
|
|
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 create_grid(res,
|
|
b_min=np.array([0, 0, 0]),
|
|
b_max=np.array([1, 1, 1]),
|
|
transform=None):
|
|
coords = np.mgrid[:res, :res, :res]
|
|
|
|
coords = coords.reshape(3, -1)
|
|
coords_matrix = np.eye(4)
|
|
length = b_max - b_min
|
|
|
|
coords_matrix[0, 0] = length[0] / res
|
|
coords_matrix[1, 1] = length[1] / res
|
|
coords_matrix[2, 2] = length[2] / res
|
|
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, res, res, res)
|
|
return coords, coords_matrix
|
|
|
|
|
|
def get_submesh(verts,
|
|
faces,
|
|
color,
|
|
verts_retained=None,
|
|
faces_retained=None,
|
|
min_vert_in_face=2):
|
|
verts = verts
|
|
faces = faces
|
|
colors = color
|
|
if verts_retained is not None:
|
|
if verts_retained.dtype != 'bool':
|
|
vert_mask = np.zeros(len(verts), dtype=bool)
|
|
vert_mask[verts_retained] = True
|
|
else:
|
|
vert_mask = verts_retained
|
|
bool_faces = np.sum(
|
|
vert_mask[faces.ravel()].reshape(-1, 3), axis=1) > min_vert_in_face
|
|
elif faces_retained is not None:
|
|
if faces_retained.dtype != 'bool':
|
|
bool_faces = np.zeros(len(faces_retained), dtype=bool)
|
|
else:
|
|
bool_faces = faces_retained
|
|
new_faces = faces[bool_faces]
|
|
vertex_ids = list(set(new_faces.ravel()))
|
|
oldtonew = -1 * np.ones([len(verts)])
|
|
oldtonew[vertex_ids] = range(0, len(vertex_ids))
|
|
new_verts = verts[vertex_ids]
|
|
new_colors = colors[vertex_ids]
|
|
new_faces = oldtonew[new_faces].astype('int32')
|
|
return (new_verts, new_faces, new_colors, bool_faces, vertex_ids)
|