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.
 
 

89 lines
2.8 KiB

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