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

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")