import numpy as np from collections import OrderedDict import re from pathlib import Path import argparse import os import json import math # hyperparameters from SkexGen project SKETCH_R = 1 RADIUS_R = 1 EXTRUDE_R = 1.0 SCALE_R = 1.4 OFFSET_R = 0.9 PIX_PAD = 4 CMD_PAD = 3 COORD_PAD = 4 EXT_PAD = 1 EXTRA_PAD = 1 R_PAD = 2 class CADparser: """Parse CAD sequence to CAD object.""" def __init__(self, bit): self.vertex_dict = OrderedDict() self.bit = bit def perform(self, cad_seq): # divide into sketch and extrude sketches, extrudes = self.get_SE(cad_seq) if sketches is None or extrudes is None: return None # sequentially parse each pair of SE into obj se_datas = [] for sketch, extrude in zip(sketches, extrudes): extrude_param, scale, offset = self.parse_extrude(extrude) if extrude_param is None or scale is None or offset is None: return None vertex_str, se_str = self.parse_sketch(sketch, scale, offset) if vertex_str is None or se_str is None: return None se_datas.append( {"vertex": vertex_str, "curve": se_str, "extrude": extrude_param} ) self.vertex_dict.clear() return se_datas def parse_sketch(self, sketch, scale, offset): faces = self.get_faces(sketch) if len(faces) == 0: return None, None se_str = "" for face_idx, face in enumerate(faces): # each face face_str = "face\n" loops = self.get_loops(face) if len(loops) == 0: return None, None for loop_idx, loop in enumerate(loops): # each loop curves = self.get_curves(loop) if len(curves) == 0: return None, None next_curves = curves[1:] next_curves += curves[:1] cur_str = [] for curve, next_curve in zip(curves, next_curves): # each curve if not self.obj_curve(curve, next_curve, cur_str, scale, offset): return None, None loop_str = "" for c in cur_str: loop_str += f"{c}\n" if loop_idx == 0: face_str += f"out\n{loop_str}\n" else: face_str += f"in\n{loop_str}\n" se_str += face_str vertex_str = self.convert_vertices() return vertex_str, se_str def parse_extrude(self, extrude): ext = extrude.split(",") if len(ext) != 18: return None, None, None # operation str to int ext_op = {"add": 1, "cut": 2, "intersect": 3}.get(ext[0], None) if ext_op is None: return None, None, None # dequantize ext_v, ext_T, scale and offset ext_v, ext_T, scale, offset = self.dequantize_extrude_params(ext) # get ext_R ext_R = np.array(ext[6:15], dtype=int) extrude_param = {"value": ext_v, "T": ext_T, "R": ext_R, "op": ext_op} return extrude_param, scale, offset def obj_curve(self, curve, next_curve, cur_str, scale, offset): cur = curve.split(",") next_cur = next_curve.split(",") if cur[0] == "circle": if len(cur) != 9: return False p1, p2, p3, p4 = self.dequantize_circle_points( cur, next_cur, scale, offset) center = np.asarray([0.5 * (p1[0] + p2[0]), 0.5 * (p3[1] + p4[1])]) radius = (np.linalg.norm(p1 - p2) + np.linalg.norm(p3 - p4)) / 4.0 center = center * scale + offset radius = radius * scale center_idx = self.save_vertex(center[0], center[1], "p") radius_idx = self.save_vertex(radius, 0.0, "r") cur_str.append(f"c {center_idx} {radius_idx}") elif cur[0] == "arc": if len(cur) != 5: return False if ( cur[1:3] == cur[3:5] or cur[1:3] == next_cur[1:3] or cur[3:5] == next_cur[3:5] ): # invalid arc return False start_v, mid_v, end_v = self.dequantize_arc_points( cur, next_cur, scale, offset ) try: center, _, _, _ = find_arc_geometry(start_v, mid_v, end_v) except Exception: return False start_v = start_v * scale + offset mid_v = mid_v * scale + offset end_v = end_v * scale + offset center = center * scale + offset center_idx = self.save_vertex(center[0], center[1], "p") start_idx = self.save_vertex(start_v[0], start_v[1], "p") mid_idx = self.save_vertex(mid_v[0], mid_v[1], "p") end_idx = self.save_vertex(end_v[0], end_v[1], "p") cur_str.append(f"a {start_idx} {mid_idx} {center_idx} {end_idx}") elif cur[0] == "line": if len(cur) != 3: return False if cur[1:3] == next_cur[1:3]: return False start_v, end_v = self.dequantize_line_points( cur, next_cur, scale, offset) start_v = start_v * scale + offset end_v = end_v * scale + offset start_idx = self.save_vertex(start_v[0], start_v[1], "p") end_idx = self.save_vertex(end_v[0], end_v[1], "p") cur_str.append(f"l {start_idx} {end_idx}") else: return False return True def get_SE(self, cad_seq): # sketches: 1) between sequence start and sketch_end, sketches_from_start = re.findall(r"^(.+?)(?=)", cad_seq) # sketches: 2) between extrude_end and sketch_end sketches_after_extrude = re.findall( r"(?<=)(.+?)(?=)", cad_seq ) sketches = [x.strip() for x in sketches_from_start] + [ x.strip() for x in sketches_after_extrude ] # extrudes: between sketch_end and extrude_end extrudes = [ x.strip() for x in re.findall(r"(.+?)", cad_seq) ] if len(sketches) != len(extrudes): return None, None return sketches, extrudes def get_faces(self, sketch): faces = sketch.split("") return [x.strip() for x in faces if x.strip() != ""] def get_loops(self, face): loops = face.split("") return [x.strip() for x in loops if x.strip() != ""] def get_curves(self, loop): curves = loop.split("") return [x.strip() for x in curves if x.strip() != ""] def dequantize_circle_points(self, curve, next_curve, scale, offset): p1 = dequantize_verts( np.array(curve[1:3], dtype=int), n_bits=self.bit, min_range=-SKETCH_R, max_range=SKETCH_R, add_noise=False, ) p2 = dequantize_verts( np.array(curve[3:5], dtype=int), n_bits=self.bit, min_range=-SKETCH_R, max_range=SKETCH_R, add_noise=False, ) p3 = dequantize_verts( np.array(curve[5:7], dtype=int), n_bits=self.bit, min_range=-SKETCH_R, max_range=SKETCH_R, add_noise=False, ) p4 = dequantize_verts( np.array(curve[7:9], dtype=int), n_bits=self.bit, min_range=-SKETCH_R, max_range=SKETCH_R, add_noise=False, ) return p1, p2, p3, p4 def dequantize_arc_points(self, curve, next_curve, scale, offset): start_v = dequantize_verts( np.array(curve[1:3], dtype=int), n_bits=self.bit, min_range=-SKETCH_R, max_range=SKETCH_R, add_noise=False, ) mid_v = dequantize_verts( np.array(curve[3:5], dtype=int), n_bits=self.bit, min_range=-SKETCH_R, max_range=SKETCH_R, add_noise=False, ) end_v = dequantize_verts( np.array(next_curve[1:3], dtype=int), n_bits=self.bit, min_range=-SKETCH_R, max_range=SKETCH_R, add_noise=False, ) return start_v, mid_v, end_v def dequantize_line_points(self, curve, next_curve, scale, offset): start_v = dequantize_verts( np.array(curve[1:3], dtype=int), n_bits=self.bit, min_range=-SKETCH_R, max_range=SKETCH_R, add_noise=False, ) end_v = dequantize_verts( np.array(next_curve[1:3], dtype=int), n_bits=self.bit, min_range=-SKETCH_R, max_range=SKETCH_R, add_noise=False, ) return start_v, end_v def dequantize_extrude_params(self, extrude): ext_v = dequantize_verts( np.array(extrude[1:3], dtype=int), n_bits=self.bit, min_range=-EXTRUDE_R, max_range=EXTRUDE_R, add_noise=False, ) ext_T = dequantize_verts( np.array(extrude[3:6], dtype=int), n_bits=self.bit, min_range=-EXTRUDE_R, max_range=EXTRUDE_R, add_noise=False, ) scale = dequantize_verts( np.array(extrude[15], dtype=int), n_bits=self.bit, min_range=0.0, max_range=SCALE_R, add_noise=False, ) offset = dequantize_verts( np.array(extrude[16:18], dtype=int), n_bits=self.bit, min_range=-OFFSET_R, max_range=OFFSET_R, add_noise=False, ) return ext_v, ext_T, scale, offset def save_vertex(self, h_x, h_y, text): unique_key = f"{text}:x{h_x}y{h_y}" index = 0 for key in self.vertex_dict.keys(): # Vertex location already exist in dict if unique_key == key: return index index += 1 # Vertex location does not exist in dict self.vertex_dict[unique_key] = [h_x, h_y] return index def convert_vertices(self): """Convert all the vertices to .obj format""" vertex_strings = "" for pt in self.vertex_dict.values(): # e.g. v 0.123 0.234 0.345 1.0 vertex_string = f"v {pt[0]} {pt[1]}\n" vertex_strings += vertex_string return vertex_strings def find_arc_geometry(a, b, c): A = b[0] - a[0] B = b[1] - a[1] C = c[0] - a[0] D = c[1] - a[1] E = A*(a[0] + b[0]) + B*(a[1] + b[1]) F = C*(a[0] + c[0]) + D*(a[1] + c[1]) G = 2.0*(A*(c[1] - b[1])-B*(c[0] - b[0])) if G == 0: raise Exception("zero G") p_0 = (D*E - B*F) / G p_1 = (A*F - C*E) / G center = np.array([p_0, p_1]) radius = np.linalg.norm(center - a) angles = [] for xx in [a, b, c]: angle = angle_from_vector_to_x(xx - center) angles.append(angle) ab = b-a ac = c-a cp = np.cross(ab, ac) if cp >= 0: start_angle_rads = angles[0] end_angle_rads = angles[2] else: start_angle_rads = angles[2] end_angle_rads = angles[0] return center, radius, start_angle_rads, end_angle_rads def angle_from_vector_to_x(vec): assert vec.size == 2 # We need to find a unit vector angle = 0.0 l = np.linalg.norm(vec) uvec = vec/l # 2 | 1 # ------- # 3 | 4 if uvec[0] >= 0: if uvec[1] >= 0: # Qadrant 1 angle = math.asin(uvec[1]) else: # Qadrant 4 angle = 2.0*math.pi - math.asin(-uvec[1]) else: if vec[1] >= 0: # Qadrant 2 angle = math.pi - math.asin(uvec[1]) else: # Qadrant 3 angle = math.pi + math.asin(-uvec[1]) return angle def dequantize_verts(verts, n_bits=8, min_range=-0.5, max_range=0.5, add_noise=False): """Convert quantized vertices to floats.""" range_quantize = 2**n_bits - 1 verts = verts.astype("float32") verts = verts * (max_range - min_range) / range_quantize + min_range return verts def write_obj_sample(save_folder, data): for idx, write_data in enumerate(data): obj_name = Path(save_folder).stem + "_" + \ str(idx).zfill(3) + "_param.obj" obj_file = Path(save_folder) / obj_name extrude_param = write_data["extrude"] vertex_strings = write_data["vertex"] curve_strings = write_data["curve"] """Write an .obj file with the curves and verts""" if extrude_param["op"] == 1: # 'add' set_op = "NewBodyFeatureOperation" elif extrude_param["op"] == 2: # 'cut' set_op = "CutFeatureOperation" elif extrude_param["op"] == 3: # 'cut' set_op = "IntersectFeatureOperation" with open(obj_file, "w") as fh: # Write Meta info fh.write("# WaveFront *.obj file\n") fh.write("# ExtrudeOperation: " + set_op + "\n") fh.write("\n") # Write vertex and curve fh.write(vertex_strings) fh.write("\n") fh.write(curve_strings) fh.write("\n") # Write extrude value extrude_string = "Extrude " for value in extrude_param["value"]: extrude_string += str(value) + " " fh.write(extrude_string) fh.write("\n") # Write refe plane value p_orig = parse3d_sample(extrude_param["T"]) x_axis = parse3d_sample(extrude_param["R"][0:3]) y_axis = parse3d_sample(extrude_param["R"][3:6]) z_axis = parse3d_sample(extrude_param["R"][6:9]) fh.write("T_origin " + p_orig) fh.write("\n") fh.write("T_xaxis " + x_axis) fh.write("\n") fh.write("T_yaxis " + y_axis) fh.write("\n") fh.write("T_zaxis " + z_axis) def parse3d_sample(point3d): x = point3d[0] y = point3d[1] z = point3d[2] return str(x) + " " + str(y) + " " + str(z) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--in-path", type=str, required=True) parser.add_argument("--out-path", type=str, required=True) args = parser.parse_args() # with open(args.in_path, "r") as f: # data = f.readlines() with open(args.in_path, 'r') as file: data = file.read() data = json.loads(data) num_valid_str = 0 for idx, item in enumerate(data): try: cad_parser = CADparser(bit=6) # print(idx) if type(item) == str: parsed_data = cad_parser.perform(item) elif type(item) == dict: parsed_data = cad_parser.perform(item['output']) else: raise ValueError("Invalid data type") out_path = os.path.join(args.out_path, str(idx).zfill(6)) os.makedirs(out_path, exist_ok=True) if parsed_data is not None: num_valid_str += 1 write_obj_sample(out_path, parsed_data) except Exception as e: print(e) pass print(f"Number of valid CAD strings: {num_valid_str}/{len(data)}")