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.
118 lines
3.4 KiB
118 lines
3.4 KiB
import os
|
|
import trimesh
|
|
import numpy as np
|
|
from scipy.spatial import KDTree
|
|
from tqdm import tqdm
|
|
import mesh2sdf
|
|
|
|
# 加载BREP模型并转换为三角网格
|
|
def load_brep_to_mesh(file_path: str) -> trimesh.Trimesh:
|
|
mesh = trimesh.load_mesh(file_path)
|
|
return mesh
|
|
|
|
def process_mesh_to_sdf(
|
|
mesh: trimesh.Trimesh,
|
|
filename:str,
|
|
save_dir:str,
|
|
idx: str,
|
|
mesh_scale:int = 0.8,
|
|
size:int= 128,
|
|
level: int = None,
|
|
) -> str:
|
|
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 create_grid(x_range, y_range, z_range, resolution):
|
|
x = np.linspace(x_range[0], x_range[1], resolution)
|
|
y = np.linspace(y_range[0], y_range[1], resolution)
|
|
z = np.linspace(z_range[0], z_range[1], resolution)
|
|
grid = np.meshgrid(x, y, z)
|
|
points = np.vstack(list(map(np.ravel, grid))).T
|
|
return points
|
|
'''
|
|
|
|
|
|
|
|
def test():
|
|
# 参数设置
|
|
input_file = "/mnt/disk2/dataset/furniture/step/furniture_dataset_step/train/bathtub_0004.step"
|
|
output_file = "tmp_data"
|
|
|
|
# 加载BREP模型并转换为三角网格
|
|
mesh = load_brep_to_mesh(input_file)
|
|
|
|
process_mesh_to_sdf(
|
|
mesh=mesh,
|
|
filename=input_file,
|
|
save_dir=output_file,
|
|
idx="0000000",
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# 真的执行转换
|
|
def main():
|
|
|
|
INPUT = '/mnt/disk2/dataset/furniture/step/furniture_dataset_step' # 下面train，val和test
|
|
OUTPUT = '/app/data/furniture_sdf'
|
|
valid = 0
|
|
for set_ in ['train', 'val', 'test']:
|
|
with ProcessPoolExecutor(max_workers=os.cpu_count() // 2) as executor:
|
|
futures = {}
|
|
for step_folder in glob.glob(os.path.join(input_folder, set_, '*.step')):
|
|
future = executor.submit(process, step_folder, timeout=300)
|
|
futures[future] = step_folder
|
|
|
|
for future in tqdm(as_completed(futures), total=len(step_dirs)):
|
|
try:
|
|
status = future.result(timeout=300)
|
|
valid += status
|
|
except TimeoutError:
|
|
print(f"Timeout occurred while processing {futures[future]}")
|
|
except Exception as e:
|
|
print(f"An error occurred while processing {futures[future]}: {e}")
|
|
|
|
print(f'Done... Data Converted Ratio {100.0*valid/len(step_dirs)}%')
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test()
|
|
|