Spaces:
Runtime error
Runtime error
File size: 2,034 Bytes
990e2a9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
import os
import shutil
from pathlib import Path
import trimesh
import ray
from trimesh.sample import sample_surface
from tqdm import tqdm
import argparse
import glob
from OCC.Core.BRepCheck import BRepCheck_Analyzer
from OCC.Core.TopAbs import TopAbs_SOLID
from OCC.Extend.DataExchange import read_step_file
from OCC.Core.ShapeFix import ShapeFix_ShapeTolerance
from eval.check_valid import check_step_valid_soild
@ray.remote
def sample(prefix, root, output_root, checkvalid):
# check valid
if checkvalid:
step_file = glob.glob(os.path.join(root / prefix, '*.step'))
if len(step_file) == 0:
return
is_valid = check_step_valid_soild(step_file[0])
if not is_valid:
return
# prefer stl
option_file = glob.glob(os.path.join(root / prefix, '*.stl')) # + glob.glob(os.path.join(root / prefix, '*.ply'))
if len(option_file) == 0:
return
file = option_file[0]
mesh = trimesh.load(str(file), force="mesh")
if mesh.faces.shape[0] == 0:
return
out_pc, _ = sample_surface(mesh, 2000)
trimesh.PointCloud(out_pc).export(str(output_root / (prefix + ".ply")))
pass
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Construct Brep From Data')
parser.add_argument('--data_root', type=str, required=True)
parser.add_argument('--out_root', type=str, required=True)
parser.add_argument('--valid', action='store_true')
root = Path(parser.parse_args().data_root)
output_root = Path(parser.parse_args().out_root)
check_valid = parser.parse_args().valid
if os.path.exists(output_root):
shutil.rmtree(output_root)
output_root.mkdir(exist_ok=False)
print("\nStart Sampling PointCloud...")
ray.init(
# num_cpus=1,
# local_mode=True,
)
tasks = []
for prefix in tqdm(os.listdir(root)):
tasks.append(sample.remote(prefix, root, output_root, check_valid))
ray.get(tasks)
print("Finish Sampling PointCloud\n")
|