add Bethany tool

Bethany/json2step.py

@@ -0,0 +1,125 @@
+import os
+import glob
+import json
+import numpy as np
+from OCC.Core.BRepCheck import BRepCheck_Analyzer
+from lib.DataExchange import write_step_file
+import argparse
+import sys
+sys.path.append(".")
+from lib.extrude import CADSequence
+from lib.visualize import vec2CADsolid, create_CAD, create_CAD_index
+from lib.file_utils import ensure_dir
+
+
+import traceback
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--src', type=str, required=True, help="source folder")
+parser.add_argument('--idx', type=int, default=0, help="export n files starting from idx.")
+parser.add_argument('--num', type=int, default=-1, help="number of shapes to export. -1 exports all shapes.")
+parser.add_argument('--mode', type=str, default="default", help="mode of generation")
+parser.add_argument('--filter', action="store_true", help="use opencascade analyzer to filter invalid model")
+parser.add_argument('-o', '--outputs', type=str, default=None, help="save folder")
+args = parser.parse_args()
+
+src_dir = args.src
+print(src_dir)
+out_paths = sorted(glob.glob(os.path.join(src_dir, "*.{}".format("json"))))
+if args.num != -1:
+    out_paths = out_paths[args.idx:args.idx+args.num]
+save_dir = args.src + "_step" if args.outputs is None else args.outputs
+ensure_dir(save_dir)
+
+
+from multiprocessing import Process, cpu_count
+
+num_processes = cpu_count()
+
+from OCC.Core.TDocStd import TDocStd_Document
+from OCC.Core.TCollection import TCollection_ExtendedString
+
+from lib.timeout import timeout_decorator
+def main_process(process_id):
+    if args.mode == "segment":
+        for index in range(process_id, len(out_paths), num_processes):
+            path = out_paths[index]
+            @timeout_decorator
+            def main__():
+                with open(path, 'r') as fp:
+                    data = json.load(fp)
+                cad_seq = CADSequence.from_dict(data)
+                
+                for i in range(len(cad_seq.seq)):
+                    doc_name = TCollection_ExtendedString("pythonocc-doc")
+                    doc = TDocStd_Document(doc_name)
+                    out_shape = create_CAD_index(doc, cad_seq, index=i)
+                    
+                    name = path.split("/")[-1].split(".")[0]
+                    sub_folder = os.path.join(save_dir, name)
+                    ensure_dir(sub_folder)
+                    save_path = os.path.join(sub_folder, name + f".{i:02d}.step")
+                    write_step_file(doc, save_path, application_protocol="AP242DIS")
+                
+                doc_name = TCollection_ExtendedString("pythonocc-doc")
+                doc = TDocStd_Document(doc_name)
+                out_shape = create_CAD(doc, cad_seq)
+
+                if args.filter:
+                    analyzer = BRepCheck_Analyzer(out_shape)
+                    if not analyzer.IsValid():
+                        print("detect invalid.")
+                        return
+                name = path.split("/")[-1].split(".")[0]
+                save_path = os.path.join(save_dir, name + f".step")
+                write_step_file(doc, save_path, application_protocol="AP242DIS")
+            try:
+                main__()
+            except Exception as e:
+                print(f"load and create failed. Error: {e}")
+                traceback.print_exc()
+                continue 
+    elif args.mode == "default":
+        for index in range(process_id, len(out_paths), num_processes):
+            path = out_paths[index]
+            @timeout_decorator
+            def main__():
+                name = path.split("/")[-1].split(".")[0]
+                save_path = os.path.join(save_dir, name + f".step")
+                if os.path.isfile(save_path): return
+                with open(path, 'r') as fp:
+                    data = json.load(fp)
+                cad_seq = CADSequence.from_dict(data)
+                doc_name = TCollection_ExtendedString("pythonocc-doc")
+                doc = TDocStd_Document(doc_name)
+                cad = CADSequence(cad_seq)
+                out_shape = create_CAD(doc, cad_seq)
+                if args.filter:
+                    analyzer = BRepCheck_Analyzer(out_shape)
+                    if not analyzer.IsValid():
+                        print("detect invalid.")
+                        return
+                
+                write_step_file(doc, save_path, application_protocol="AP242DIS")
+            try:
+                main__()
+            except Exception as e:
+                print(f"load and create failed. Error: {e}")
+                traceback.print_exc()
+                continue
+        
+
+    ## python export2step.py --src ./ --outputs ./
+
+if __name__ == "__main__":
+    processes = []
+    for i in range(num_processes):
+        process = Process(target=main_process, args=(i,))
+        processes.append(process)
+        process.start()
+
+    # 等待所有进程完成
+    for process in processes:
+        process.join()
+
+    print('任务完成')