| import numpy as np
|
| from collections import OrderedDict
|
| import re
|
| from pathlib import Path
|
| import argparse
|
| import os
|
| import json
|
| import math
|
|
|
|
|
| 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):
|
|
|
| sketches, extrudes = self.get_SE(cad_seq)
|
| if sketches is None or extrudes is None:
|
| return None
|
|
|
| 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):
|
| face_str = "face\n"
|
| loops = self.get_loops(face)
|
| if len(loops) == 0:
|
| return None, None
|
| for loop_idx, loop in enumerate(loops):
|
| 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):
|
| 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
|
|
|
|
|
| ext_op = {"add": 1, "cut": 2, "intersect": 3}.get(ext[0], None)
|
| if ext_op is None:
|
| return None, None, None
|
|
|
| ext_v, ext_T, scale, offset = self.dequantize_extrude_params(ext)
|
|
|
| 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]
|
| ):
|
| 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_from_start = re.findall(r"^(.+?)(?=<sketch_end>)", cad_seq)
|
|
|
| sketches_after_extrude = re.findall(
|
| r"(?<=<extrude_end>)(.+?)(?=<sketch_end>)", cad_seq
|
| )
|
| sketches = [x.strip() for x in sketches_from_start] + [
|
| x.strip() for x in sketches_after_extrude
|
| ]
|
|
|
| extrudes = [
|
| x.strip() for x in re.findall(r"<sketch_end>(.+?)<extrude_end>", cad_seq)
|
| ]
|
| if len(sketches) != len(extrudes):
|
| return None, None
|
| return sketches, extrudes
|
|
|
| def get_faces(self, sketch):
|
| faces = sketch.split("<face_end>")
|
| return [x.strip() for x in faces if x.strip() != ""]
|
|
|
| def get_loops(self, face):
|
| loops = face.split("<loop_end>")
|
| return [x.strip() for x in loops if x.strip() != ""]
|
|
|
| def get_curves(self, loop):
|
| curves = loop.split("<curve_end>")
|
| 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():
|
|
|
| if unique_key == key:
|
| return index
|
| index += 1
|
|
|
| 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():
|
|
|
| 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
|
|
|
| angle = 0.0
|
|
|
| l = np.linalg.norm(vec)
|
| uvec = vec/l
|
|
|
|
|
|
|
|
|
| if uvec[0] >= 0:
|
| if uvec[1] >= 0:
|
|
|
| angle = math.asin(uvec[1])
|
| else:
|
|
|
| angle = 2.0*math.pi - math.asin(-uvec[1])
|
| else:
|
| if vec[1] >= 0:
|
|
|
| angle = math.pi - math.asin(uvec[1])
|
| else:
|
|
|
| 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:
|
| set_op = "NewBodyFeatureOperation"
|
| elif extrude_param["op"] == 2:
|
| set_op = "CutFeatureOperation"
|
| elif extrude_param["op"] == 3:
|
| set_op = "IntersectFeatureOperation"
|
|
|
| with open(obj_file, "w") as fh:
|
|
|
| fh.write("# WaveFront *.obj file\n")
|
| fh.write("# ExtrudeOperation: " + set_op + "\n")
|
| fh.write("\n")
|
|
|
|
|
| fh.write(vertex_strings)
|
| fh.write("\n")
|
| fh.write(curve_strings)
|
| fh.write("\n")
|
|
|
|
|
| extrude_string = "Extrude "
|
| for value in extrude_param["value"]:
|
| extrude_string += str(value) + " "
|
| fh.write(extrude_string)
|
| fh.write("\n")
|
|
|
|
|
| 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 file:
|
| data = file.read()
|
|
|
| data = json.loads(data)
|
|
|
| num_valid_str = 0
|
| for idx, item in enumerate(data):
|
| try:
|
| cad_parser = CADparser(bit=6)
|
|
|
| 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)}")
|
|
|
