You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
90 lines
2.8 KiB
90 lines
2.8 KiB
7 months ago
|
import os
|
||
|
|
import sys
|
||
|
|
import trimesh
|
||
|
|
import mesh2sdf
|
||
|
|
import numpy as np
|
||
|
|
import time
|
||
|
|
from tqdm import tqdm
|
||
|
|
from concurrent.futures import ProcessPoolExecutor
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
def process_mesh_to_sdf(
|
||
|
|
filename:str,
|
||
|
|
save_dir:str,
|
||
|
|
idx: str,
|
||
|
|
mesh_scale:int = 0.8
|
||
|
|
size:int= 128
|
||
|
|
level: int = None
|
||
|
|
):
|
||
|
|
if not level:
|
||
|
|
level = 2 / size
|
||
|
|
|
||
|
|
mesh = trimesh.load(filename, force='mesh')
|
||
|
|
|
||
|
|
# normalize mesh
|
||
|
|
vertices = mesh.vertices
|
||
|
|
bbmin = vertices.min(0)
|
||
|
|
bbmax = vertices.max(0)
|
||
|
|
center = (bbmin + bbmax) * 0.5
|
||
|
|
scale = 2.0 * mesh_scale / (bbmax - bbmin).max()
|
||
|
|
vertices = (vertices - center) * scale
|
||
|
|
|
||
|
|
# fix mesh
|
||
|
|
t0 = time.time()
|
||
|
|
sdf, mesh = mesh2sdf.compute(
|
||
|
|
vertices, mesh.faces, size, fix=True, level=level, return_mesh=True)
|
||
|
|
t1 = time.time()
|
||
|
|
|
||
|
|
# sdf to x,y,z,sdf
|
||
|
|
xyzsdf = np.zeros((size, size, size, 4))
|
||
|
|
for x in range(size):
|
||
|
|
for y in range(size):
|
||
|
|
for z in range(size):
|
||
|
|
xyzsdf[x, y, z] = [x, y, z, sdf[x, y, z]]
|
||
|
|
# output
|
||
|
|
mesh.vertices = mesh.vertices / scale + center
|
||
|
|
name = filename.split('/')[-1][:-4]
|
||
|
|
mesh.export(os.path.join(save_dir, filename[:-4] + '.fixed.obj'))
|
||
|
|
np.save(os.path.join(save_dir, filename[:-4] + '.xyzsdf.npy'), xyzsdf)
|
||
|
|
#mesh.export(filename[:-4] + '.fixed.obj')
|
||
|
|
#np.save(filename[:-4] + '.npy', sdf)
|
||
|
|
#print('It takes %.4f seconds to process %s' % (t1-t0, filename))
|
||
|
|
return idx
|
||
|
|
|
||
|
|
|
||
|
|
def mesh_to_sdf():
|
||
|
|
dataset_dir = "/home/wch/data/abc"
|
||
|
|
valid_id = []
|
||
|
|
'''
|
||
|
|
for i in tqdm(range(4025,10000)):
|
||
|
|
internal_dir_name = f"{i:08d}"
|
||
|
|
data_dir = os.path.join(dataset_dir, internal_dir_name)
|
||
|
|
for name in os.listdir(data_dir):
|
||
|
|
file = os.path.join(data_dir,name)
|
||
|
|
if os.path.isfile(file) and name.endswith(".obj") and not name.endswith(".fixed.obj"):
|
||
|
|
#process(file,data_dir)
|
||
|
|
executor.submit(process, file, data_dir)
|
||
|
|
valid_id.append(internal_dir_name)
|
||
|
|
'''
|
||
|
|
# 创建一个进程池
|
||
|
|
with ProcessPoolExecutor(max_workers=5) as executor:
|
||
|
|
futures = []
|
||
|
|
for i in tqdm(range(4025, 10000)):
|
||
|
|
internal_dir_name = f"{i:08d}"
|
||
|
|
data_dir = os.path.join(dataset_dir, internal_dir_name)
|
||
|
|
for name in os.listdir(data_dir):
|
||
|
|
file = os.path.join(data_dir, name)
|
||
|
|
if os.path.isfile(file) and name.endswith(".obj") and not name.endswith(".fixed.obj"):
|
||
|
|
future = executor.submit(process, file, data_dir,internal_dir_name)
|
||
|
|
futures.append(future)
|
||
|
|
|
||
|
|
# 等待所有任务完成，并收集结果
|
||
|
|
for future in tqdm(futures):
|
||
|
|
valid_id.append(future.result())
|
||
|
|
|
||
|
|
with open("/home/wch/data/abc/valid.txt", "w") as f:
|
||
|
|
for idx in valid_id:
|
||
|
|
f.write(idx+"\n")
|