File size: 4,017 Bytes
91daf98 | 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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | import os
import argparse
from pathlib import Path
from tqdm import tqdm
from multiprocessing import Pool
from glob import glob
from utils.obj_reconverter import OBJReconverter
from OCC.Core.BRepCheck import BRepCheck_Analyzer
from geometry.obj_parser import OBJParser
from utils.util import write_stl_file
from OCC.Extend.DataExchange import write_step_file
import signal
from contextlib import contextmanager
@contextmanager
def timeout(time):
# Register a function to raise a TimeoutError on the signal.
signal.signal(signal.SIGALRM, raise_timeout)
# Schedule the signal to be sent after ``time``.
signal.alarm(time)
try:
yield
except TimeoutError:
raise Exception("time out")
finally:
# Unregister the signal so it won't be triggered
# if the timeout is not reached.
signal.signal(signal.SIGALRM, signal.SIG_IGN)
def raise_timeout(signum, frame):
raise TimeoutError
NUM_TRHEADS = 36
def find_files(folder, extension):
return sorted([Path(os.path.join(folder, f)) for f in os.listdir(folder) if f.endswith(extension)])
def run_parallel(project_folder):
output_folder = project_folder
param_objs = find_files(project_folder, 'param.obj')
cur_solid = None
extrude_idx = 0
for obj in param_objs:
try:
with timeout(30):
parser = OBJParser(obj)
_, faces, meta_info = parser.parse_file(1.0)
converter = OBJReconverter()
ext_solid, _, _ = converter.parse_obj(faces, meta_info)
set_op = meta_info["set_op"]
if set_op == "NewBodyFeatureOperation" or set_op == "JoinFeatureOperation":
if cur_solid is None:
cur_solid = ext_solid
else:
cur_solid = converter.my_op(cur_solid, ext_solid, 'fuse')
elif set_op == "CutFeatureOperation":
cur_solid = converter.my_op(cur_solid, ext_solid, 'cut')
elif set_op == "IntersectFeatureOperation":
cur_solid = converter.my_op(cur_solid, ext_solid, 'common')
else:
raise Exception("Unknown operation type")
analyzer = BRepCheck_Analyzer(cur_solid)
if not analyzer.IsValid():
raise Exception("brep check failed")
extrude_idx += 1
except Exception as ex:
print(ex)
msg = [project_folder, str(ex)[:100]]
return None
try:
with timeout(30):
stl_name = Path(output_folder).stem + '_'+ str(extrude_idx).zfill(3) + "_final.stl"
output_path = os.path.join(output_folder, stl_name)
write_stl_file(cur_solid, output_path, linear_deflection=0.001, angular_deflection=0.5)
step_name = Path(output_folder).stem + '_'+ str(extrude_idx).zfill(3) + "_final.step"
output_path = os.path.join(output_folder, step_name)
write_step_file(cur_solid, output_path)
except Exception as ex:
print(ex)
msg = [project_folder, str(ex)[:500]]
return None
return cur_solid
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--data_folder", type=str, required=True)
parser.add_argument("--single-file", action='store_true', default=False)
args = parser.parse_args()
if args.single_file:
# If single file, just run the function on that file
run_parallel(args.data_folder)
exit(0)
else:
solids = []
# cad_folders = sorted(glob(args.data_folder+'/*'))[50000:] # why after 50000?
cad_folders = sorted(glob(args.data_folder+'/*'))
# print("len of cad_folder:", len(cad_folders))
convert_iter = Pool(NUM_TRHEADS).imap(run_parallel, cad_folders)
for solid in tqdm(convert_iter, total=len(cad_folders)):
pass
|