From 7b6e0703c2e43c156bd77d21f7b8b97839494d00 Mon Sep 17 00:00:00 2001 From: Yuki-Kokomi Date: Wed, 7 Aug 2024 10:07:38 +0800 Subject: [PATCH] first commit --- README.md | 26 +++ examples/example01.py | 46 +++++ examples/example02.py | 43 +++++ lib/DataExchange.py | 40 ++++ lib/__init__.py | 0 lib/curves.py | 422 ++++++++++++++++++++++++++++++++++++++++++ lib/extrude.py | 324 ++++++++++++++++++++++++++++++++ lib/file_utils.py | 52 ++++++ lib/macro.py | 40 ++++ lib/math_utils.py | 367 ++++++++++++++++++++++++++++++++++++ lib/sketch.py | 263 ++++++++++++++++++++++++++ lib/timeout.py | 31 ++++ lib/visualize.py | 197 ++++++++++++++++++++ py2step.py | 214 +++++++++++++++++++++ requirements.txt | 4 + 15 files changed, 2069 insertions(+) create mode 100644 README.md create mode 100644 examples/example01.py create mode 100644 examples/example02.py create mode 100644 lib/DataExchange.py create mode 100644 lib/__init__.py create mode 100644 lib/curves.py create mode 100644 lib/extrude.py create mode 100644 lib/file_utils.py create mode 100644 lib/macro.py create mode 100644 lib/math_utils.py create mode 100644 lib/sketch.py create mode 100644 lib/timeout.py create mode 100644 lib/visualize.py create mode 100644 py2step.py create mode 100644 requirements.txt diff --git a/README.md b/README.md new file mode 100644 index 0000000..56f5fb2 --- /dev/null +++ b/README.md @@ -0,0 +1,26 @@ +# OpenECADv2toSTEP + +## Prerequisites + +- Linux +- Python 3.10 + +## Dependencies + +Install python package dependencies through pip: + +```bash +$ pip install -r requirements.txt +``` + +Install [pythonocc](https://github.com/tpaviot/pythonocc-core) (OpenCASCADE) by conda: + +```bash +$ conda install -c conda-forge pythonocc-core=7.5.1 +``` + +## Usage + +```bash +python py2step.py --src examples/ -o outputs/ +``` \ No newline at end of file diff --git a/examples/example01.py b/examples/example01.py new file mode 100644 index 0000000..c03729e --- /dev/null +++ b/examples/example01.py @@ -0,0 +1,46 @@ +SketchPlane0 = add_sketchplane( + origin= [0., 0., 0.], normal= [ 0., -1., 0.], x_axis= [ 1., 0., -0.]) +Loops0 = [] +Curves0_0 = [] +Line0_0_0 = add_line(start= [-50.8, 0. ], end= [0., 0.]) +Line0_0_1 = add_line(start= [0., 0.], end= [ 0. , 50.8]) +Line0_0_2 = add_line(start= [ 0. , 50.8], end= [-50.8, 50.8]) +Line0_0_3 = add_line(start= [-50.8, 50.8], end= [-50.8, 0. ]) +Loop0_0 = add_loop(Curves0_0) +Profile0 = add_profile(Loops0) +Sketch0 = add_sketch(sketch_plane= SketchPlane0, profile= Profile0) +Extrude0 = add_extrude(sketch= Sketch0, + operation= 0, type= 0, extent_one= 50.8, extent_two= 0.) +SketchPlane1 = add_sketchplane_ref( + Extrude0, origin= [25.4, 25.4], type= "line", line= Line0_0_2, angle= np.pi) +Loops1 = [] +Curves1_0 = [] +Circle1_0_0 = add_circle(center= [0., 0.], radius= 12.7) +Loop1_0 = add_loop(Curves1_0) +Profile1 = add_profile(Loops1) +Sketch1 = add_sketch(sketch_plane= SketchPlane1, profile= Profile1) +Extrude1 = add_extrude(sketch= Sketch1, + operation= 2, type= 0, extent_one= -50.8, extent_two= 0.) +SketchPlane2 = add_sketchplane_ref( + Extrude0, origin = [-25.4, 25.4], type = "sameplane", angle = -np.pi, reverse = True) +Loops2 = [] +Curves2_0 = [] +Circle2_0_0 = add_circle(center= [0., 0.], radius= 12.7) +Loop2_0 = add_loop(Curves2_0) +Profile2 = add_profile(Loops2) +Sketch2 = add_sketch(sketch_plane= SketchPlane2, profile= Profile2) +Extrude2 = add_extrude(sketch= Sketch2, + operation= 1, type= 0, extent_one= 12.7, extent_two= 0.) +SketchPlane3 = add_sketchplane_ref( + Extrude0, origin= [25.4, 25.4], type= "line", line= Line0_0_1, angle= -np.pi/2) +Loops3 = [] +Curves3_0 = [] +Line3_0_0 = add_line(start= [-12.7, -12.7], end= [ 12.7, -12.7]) +Line3_0_1 = add_line(start= [ 12.7, -12.7], end= [12.7, 12.7]) +Line3_0_2 = add_line(start= [12.7, 12.7], end= [-12.7, 12.7]) +Line3_0_3 = add_line(start= [-12.7, 12.7], end= [-12.7, -12.7]) +Loop3_0 = add_loop(Curves3_0) +Profile3 = add_profile(Loops3) +Sketch3 = add_sketch(sketch_plane= SketchPlane3, profile= Profile3) +Extrude3 = add_extrude(sketch= Sketch3, + operation= 1, type= 0, extent_one= 12.7, extent_two= 0.) diff --git a/examples/example02.py b/examples/example02.py new file mode 100644 index 0000000..0a077a3 --- /dev/null +++ b/examples/example02.py @@ -0,0 +1,43 @@ +SketchPlane0 = add_sketchplane( + origin= [0., 0., 0.], normal= [ 0., -1., 0.], x_axis= [ 1., 0., -0.]) +Loops0 = [] +Curves0_0 = [] +Circle0_0_0 = add_circle(center= [0., 0.], radius= 41.7551) +Loop0_0 = add_loop(Curves0_0) +Profile0 = add_profile(Loops0) +Sketch0 = add_sketch(sketch_plane= SketchPlane0, profile= Profile0) +Extrude0 = add_extrude(sketch= Sketch0, + operation= 0, type= 0, extent_one= 107.9331, extent_two= 0.) +SketchPlane1 = add_sketchplane_ref( + Extrude0, origin = [-0., -0.], type = "extent_one") +Loops1 = [] +Curves1_0 = [] +Circle1_0_0 = add_circle(center= [0., 0.], radius= 28.914) +Loop1_0 = add_loop(Curves1_0) +Profile1 = add_profile(Loops1) +Sketch1 = add_sketch(sketch_plane= SketchPlane1, profile= Profile1) +Extrude1 = add_extrude(sketch= Sketch1, + operation= 1, type= 0, extent_one= 126.1059, extent_two= 0.) +SketchPlane2 = add_sketchplane_ref( + Extrude1, origin = [-0., -0.], type = "extent_one") +Loops2 = [] +Curves2_0 = [] +Circle2_0_0 = add_circle(center= [0., 0.], radius= 21.4466) +Loop2_0 = add_loop(Curves2_0) +Profile2 = add_profile(Loops2) +Sketch2 = add_sketch(sketch_plane= SketchPlane2, profile= Profile2) +Extrude2 = add_extrude(sketch= Sketch2, + operation= 2, type= 0, extent_one= -265.6521, extent_two= 0.) +SketchPlane3 = add_sketchplane_ref( + Extrude1, origin = [-0., -0.], type = "extent_one") +Loops3 = [] +Curves3_0 = [] +Circle3_0_0 = add_circle(center= [0., 0.], radius= 54.4952) +Loop3_0 = add_loop(Curves3_0) +Curves3_1 = [] +Circle3_1_0 = add_circle(center= [0., 0.], radius= 28.914) +Loop3_1 = add_loop(Curves3_1) +Profile3 = add_profile(Loops3) +Sketch3 = add_sketch(sketch_plane= SketchPlane3, profile= Profile3) +Extrude3 = add_extrude(sketch= Sketch3, + operation= 0, type= 0, extent_one= -30.4578, extent_two= 0.) diff --git a/lib/DataExchange.py b/lib/DataExchange.py new file mode 100644 index 0000000..4f4f82f --- /dev/null +++ b/lib/DataExchange.py @@ -0,0 +1,40 @@ +import os +from OCC.Core.STEPControl import STEPControl_Reader, STEPControl_Writer, STEPControl_AsIs +from OCC.Core.STEPCAFControl import STEPCAFControl_Writer + +from OCC.Core.Interface import Interface_Static_SetCVal +from OCC.Core.IFSelect import IFSelect_RetDone, IFSelect_ItemsByEntity +from OCC.Core.TDocStd import TDocStd_Document +from OCC.Core.TCollection import TCollection_ExtendedString + +try: + import svgwrite + HAVE_SVGWRITE = True +except ImportError: + HAVE_SVGWRITE = False + +def write_step_file(a_shape, filename, application_protocol="AP242DIS"): + """ exports a shape to a STEP file + a_shape: the topods_shape to export (a compound, a solid etc.) + filename: the filename + application protocol: "AP203" or "AP214IS" or "AP242DIS" + """ + # a few checks + #if a_shape.IsNull(): + # raise AssertionError("Shape %s is null." % a_shape) + if application_protocol not in ["AP203", "AP214IS", "AP242DIS"]: + raise AssertionError("application_protocol must be either AP203 or AP214IS. You passed %s." % application_protocol) + if os.path.isfile(filename): + print("Warning: %s file already exists and will be replaced" % filename) + # creates and initialise the step exporter + step_writer = STEPCAFControl_Writer() + Interface_Static_SetCVal("write.step.schema", application_protocol) + + # transfer shapes and write file + step_writer.Transfer(a_shape, STEPControl_AsIs) + status = step_writer.Write(filename) + + if not status == IFSelect_RetDone: + raise IOError("Error while writing shape to STEP file.") + if not os.path.isfile(filename): + raise IOError("File %s was not saved to filesystem." % filename) diff --git a/lib/__init__.py b/lib/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/lib/curves.py b/lib/curves.py new file mode 100644 index 0000000..9c38d21 --- /dev/null +++ b/lib/curves.py @@ -0,0 +1,422 @@ +import numpy as np +import matplotlib.lines as lines +import matplotlib.patches as patches +from .math_utils import rads_to_degs, angle_from_vector_to_x +from .macro import * + + +# FIXME: these two functions can be treated as static method +def construct_curve_from_dict(stat): + if stat['type'] == "Line3D": + return Line.from_dict(stat) + elif stat['type'] == "Circle3D": + return Circle.from_dict(stat) + elif stat['type'] == "Arc3D": + return Arc.from_dict(stat) + else: + raise NotImplementedError("curve type not supported yet: {}".format(stat['type'])) + + +def construct_curve_from_vector(vec, start_point, is_numerical=True): + type = vec[0] + if type == LINE_IDX: + return Line.from_vector(vec, start_point, is_numerical=is_numerical) + elif type == CIRCLE_IDX: + return Circle.from_vector(vec, start_point, is_numerical=is_numerical) + elif type == ARC_IDX: + res = Arc.from_vector(vec, start_point, is_numerical=is_numerical) + if res is None: # for visualization purpose, replace illed arc with line + return Line.from_vector(vec, start_point, is_numerical=is_numerical) + return res + else: + raise NotImplementedError("curve type not supported yet: command idx {}".format(vec[0])) + + +####################### base ####################### +class CurveBase(object): + """Base class for curve. All types of curves shall inherit from this.""" + def __init__(self): + pass + + @staticmethod + def from_dict(stat): + """construct curve from json data""" + raise NotImplementedError + + @staticmethod + def from_vector(vec, start_point, is_numerical=True): + """construct curve from vector representation""" + raise NotImplementedError + + @property + def bbox(self): + """compute bounding box of the curve""" + raise NotImplementedError + + def direction(self, from_start=True): + """return a vector indicating the curve direction""" + raise NotImplementedError + + def transform(self, translate, scale): + """linear transformation""" + raise NotImplementedError + + def flip(self, axis): + """flip the curve about axis""" + raise NotImplementedError + + def reverse(self): + """reverse the curve direction""" + raise NotImplementedError + + def numericalize(self, n=256): + """quantize curve parameters into integers""" + raise NotImplementedError + + def to_vector(self): + """represent curve using a vector. see macro.py""" + raise NotImplementedError + + def draw(self, ax, color): + """draw the curve using matplotlib""" + raise NotImplementedError + + def sample_points(self, n=32): + """uniformly sample points from the curve""" + raise NotImplementedError + + +####################### curves ####################### +class Line(CurveBase): + def __init__(self, start_point, end_point): + super(Line, self).__init__() + self.start_point = start_point + self.end_point = end_point + + def __str__(self): + return "Line: start({}), end({})".format(self.start_point.round(4), self.end_point.round(4)) + + @staticmethod + def from_dict(stat): + assert stat['type'] == "Line3D" + start_point = np.array([stat['start_point']['x'] * 1000, + stat['start_point']['y'] * 1000]) + end_point = np.array([stat['end_point']['x'] * 1000, + stat['end_point']['y'] * 1000]) + return Line(start_point, end_point) + + @staticmethod + def from_vector(vec, start_point, is_numerical=True): + return Line(start_point, vec[1:3]) + + @property + def bbox(self): + points = np.stack([self.start_point, self.end_point], axis=0) + return np.stack([np.min(points, axis=0), np.max(points, axis=0)], axis=0) + + def direction(self, from_start=True): + return self.end_point - self.start_point + + def transform(self, translate, scale): + self.start_point = (self.start_point + translate) * scale + self.end_point = (self.end_point + translate) * scale + + def flip(self, axis): + if axis == 'x': + self.start_point[1], self.end_point[1] = -self.start_point[1], -self.end_point[1] + elif axis == 'y': + self.start_point[0], self.end_point[0] = -self.start_point[0], -self.end_point[0] + elif axis == 'xy': + self.start_point = self.start_point * -1 + self.end_point = self.end_point * -1 + else: + raise ValueError("axis = {}".format(axis)) + + def reverse(self): + self.start_point, self.end_point = self.end_point, self.start_point + + def numericalize(self, n=256): + self.start_point = self.start_point.round().clip(min=0, max=n-1).astype(np.int) + self.end_point = self.end_point.round().clip(min=0, max=n-1).astype(np.int) + + def to_vector(self): + vec = [LINE_IDX, self.end_point[0], self.end_point[1]] + return np.array(vec + [PAD_VAL] * (1 + N_ARGS - len(vec))) + + def draw(self, ax, color): + xdata = [self.start_point[0], self.end_point[0]] + ydata = [self.start_point[1], self.end_point[1]] + l1 = lines.Line2D(xdata, ydata, lw=1, color=color, axes=ax) + ax.add_line(l1) + ax.plot(self.start_point[0], self.start_point[1], 'ok', color=color) + # ax.plot(self.end_point[0], self.end_point[1], 'ok') + + def sample_points(self, n=32): + return np.linspace(self.start_point, self.end_point, num=n) + + +class Arc(CurveBase): + def __init__(self, start_point, end_point, center, radius, + normal=None, start_angle=None, end_angle=None, ref_vec=None, mid_point=None): + super(Arc, self).__init__() + self.start_point = start_point + self.end_point = end_point + self.center = center + self.radius = radius + self.normal = normal + self.start_angle = start_angle + self.end_angle = end_angle + self.ref_vec = ref_vec + if mid_point is None: + self.mid_point = self.get_mid_point() + else: + self.mid_point = mid_point + + def __str__(self): + return "Arc: start({}), end({}), mid({})".format(self.start_point.round(4), self.end_point.round(4), + self.mid_point.round(4)) + + @staticmethod + def from_dict(stat): + assert stat['type'] == "Arc3D" + start_point = np.array([stat['start_point']['x'] * 1000, + stat['start_point']['y'] * 1000]) + end_point = np.array([stat['end_point']['x'] * 1000, + stat['end_point']['y'] * 1000]) + center = np.array([stat['center_point']['x'] * 1000, + stat['center_point']['y'] * 1000]) + radius = stat['radius'] * 1000 + normal = np.array([stat['normal']['x'], + stat['normal']['y'], + stat['normal']['z']]) + start_angle = stat['start_angle'] + end_angle = stat['end_angle'] + ref_vec = np.array([stat['reference_vector']['x'], + stat['reference_vector']['y']]) + return Arc(start_point, end_point, center, radius, normal, start_angle, end_angle, ref_vec) + + @staticmethod + def from_vector(vec, start_point, is_numerical=True): + end_point = vec[1:3] + sweep_angle = vec[3] / 256 * 2 * np.pi if is_numerical else vec[3] + clock_sign = vec[4] + s2e_vec = end_point - start_point + if np.linalg.norm(s2e_vec) == 0: + return None + radius = (np.linalg.norm(s2e_vec) / 2) / np.sin(sweep_angle / 2) + s2e_mid = (start_point + end_point) / 2 + vertical = np.cross(s2e_vec, [0, 0, 1])[:2] + vertical = vertical / np.linalg.norm(vertical) + if clock_sign == 0: + vertical = -vertical + center_point = s2e_mid - vertical * (radius * np.cos(sweep_angle / 2)) + + start_angle = 0 + end_angle = sweep_angle + if clock_sign == 0: + ref_vec = end_point - center_point + else: + ref_vec = start_point - center_point + ref_vec = ref_vec / np.linalg.norm(ref_vec) + + return Arc(start_point, end_point, center_point, radius, + start_angle=start_angle, end_angle=end_angle, ref_vec=ref_vec) + + def get_angles_counterclockwise(self, eps=1e-8): + c2s_vec = (self.start_point - self.center) / (np.linalg.norm(self.start_point - self.center) + eps) + c2m_vec = (self.mid_point - self.center) / (np.linalg.norm(self.mid_point - self.center) + eps) + c2e_vec = (self.end_point - self.center) / (np.linalg.norm(self.end_point - self.center) + eps) + angle_s, angle_m, angle_e = angle_from_vector_to_x(c2s_vec), angle_from_vector_to_x(c2m_vec), \ + angle_from_vector_to_x(c2e_vec) + angle_s, angle_e = min(angle_s, angle_e), max(angle_s, angle_e) + if not angle_s < angle_m < angle_e: + angle_s, angle_e = angle_e - np.pi * 2, angle_s + return angle_s, angle_e + + @property + def bbox(self): + points = [self.start_point, self.end_point] + angle_s, angle_e = self.get_angles_counterclockwise() + if angle_s < 0 < angle_e: + points.append(np.array([self.center[0] + self.radius, self.center[1]])) + if angle_s < np.pi / 2 < angle_e or angle_s < -np.pi / 2 * 3 < angle_e: + points.append(np.array([self.center[0], self.center[1] + self.radius])) + if angle_s < np.pi < angle_e or angle_s < -np.pi < angle_e: + points.append(np.array([self.center[0] - self.radius, self.center[1]])) + if angle_s < np.pi / 2 * 3 < angle_e or angle_s < -np.pi/2 < angle_e: + points.append(np.array([self.center[0], self.center[1] - self.radius])) + points = np.stack(points, axis=0) + return np.stack([np.min(points, axis=0), np.max(points, axis=0)], axis=0) + + def direction(self, from_start=True): + if from_start: + return self.mid_point - self.start_point + else: + return self.end_point - self.mid_point + + @property + def clock_sign(self): + """get a boolean sign indicating whether the arc is on top of s->e """ + s2e = self.end_point - self.start_point + s2m = self.mid_point - self.start_point + sign = np.cross(s2m, s2e) >= 0 # counter-clockwise + return sign + + def get_mid_point(self): + mid_angle = (self.start_angle + self.end_angle) / 2 + rot_mat = np.array([[np.cos(mid_angle), -np.sin(mid_angle)], + [np.sin(mid_angle), np.cos(mid_angle)]]) + mid_vec = rot_mat @ self.ref_vec + return self.center + mid_vec * self.radius + + def transform(self, translate, scale): + self.start_point = (self.start_point + translate) * scale + self.mid_point = (self.mid_point + translate) * scale + self.end_point = (self.end_point + translate) * scale + self.center = (self.center + translate) * scale + if isinstance(scale * 1.0, float): + self.radius = abs(self.radius * scale) + + def flip(self, axis): + if axis == 'x': + self.transform(0, np.array([1, -1])) + new_ref_vec_angle = angle_from_vector_to_x(self.ref_vec) + self.end_angle - self.start_angle + self.ref_vec = np.array([np.cos(new_ref_vec_angle), -np.sin(new_ref_vec_angle)]) + elif axis == 'y': + self.transform(0, np.array([-1, 1])) + new_ref_vec_angle = angle_from_vector_to_x(self.ref_vec) + self.end_angle - self.start_angle + self.ref_vec = np.array([-np.cos(new_ref_vec_angle), np.sin(new_ref_vec_angle)]) + elif axis == 'xy': + self.transform(0, -1) + self.ref_vec = self.ref_vec * -1 + else: + raise ValueError("axis = {}".format(axis)) + + def reverse(self): + self.start_point, self.end_point = self.end_point, self.start_point + + def numericalize(self, n=256): + self.start_point = self.start_point.round().clip(min=0, max=n-1).astype(np.int) + self.mid_point = self.mid_point.round().clip(min=0, max=n-1).astype(np.int) + self.end_point = self.end_point.round().clip(min=0, max=n-1).astype(np.int) + self.center = self.center.round().clip(min=0, max=n-1).astype(np.int) + tmp = np.array([self.start_angle, self.end_angle]) + self.start_angle, self.end_angle = (tmp / (2 * np.pi) * n).round().clip( + min=0, max=n-1).astype(np.int) + + def to_vector(self): + sweep_angle = max(abs(self.start_angle - self.end_angle), 1) + return np.array([ARC_IDX, self.end_point[0], self.end_point[1], sweep_angle, int(self.clock_sign), PAD_VAL, + *[PAD_VAL] * N_ARGS_EXT]) + + def draw(self, ax, color): + ref_vec_angle = rads_to_degs(angle_from_vector_to_x(self.ref_vec)) + start_angle = rads_to_degs(self.start_angle) + end_angle = rads_to_degs(self.end_angle) + diameter = 2.0 * self.radius + ap = patches.Arc( + (self.center[0], self.center[1]), + diameter, + diameter, + angle=ref_vec_angle, + theta1=start_angle, + theta2=end_angle, + lw=1, + color=color + ) + ax.add_patch(ap) + ax.plot(self.start_point[0], self.start_point[1], 'ok', color=color) + # ax.plot(self.center[0], self.center[1], 'ok', color=color) + ax.plot(self.mid_point[0], self.mid_point[1], 'ok', color=color) + # ax.plot(self.end_point[0], self.end_point[1], 'ok') + + def sample_points(self, n=32): + c2s_vec = (self.start_point - self.center) / np.linalg.norm(self.start_point - self.center) + c2m_vec = (self.mid_point - self.center) / np.linalg.norm(self.mid_point - self.center) + c2e_vec = (self.end_point - self.center) / np.linalg.norm(self.end_point - self.center) + angle_s, angle_m, angle_e = angle_from_vector_to_x(c2s_vec), angle_from_vector_to_x(c2m_vec), \ + angle_from_vector_to_x(c2e_vec) + angle_s, angle_e = min(angle_s, angle_e), max(angle_s, angle_e) + if not angle_s < angle_m < angle_e: + angle_s, angle_e = angle_e - np.pi * 2, angle_s + + angles = np.linspace(angle_s, angle_e, num=n) + points = np.stack([np.cos(angles), np.sin(angles)], axis=1) * self.radius + self.center[np.newaxis] + return points + + +class Circle(CurveBase): + def __init__(self, center, radius, normal=None): + super(Circle, self).__init__() + self.center = center + self.radius = radius + self.normal = normal + + def __str__(self): + return "Circle: center({}), radius({})".format(self.center.round(4), round(self.radius, 4)) + + @staticmethod + def from_dict(stat): + assert stat['type'] == "Circle3D" + center = np.array([stat['center_point']['x'] * 1000, + stat['center_point']['y'] * 1000]) + radius = stat['radius'] * 1000 + normal = np.array([stat['normal']['x'], + stat['normal']['y'], + stat['normal']['z']]) + return Circle(center, radius, normal) + + @staticmethod + def from_vector(vec, start_point=None, is_numerical=True): + return Circle(vec[1:3], vec[5]) + + @property + def bbox(self): + return np.stack([self.center - self.radius, self.center + self.radius], axis=0) + + def direction(self, from_start=True): + return self.center - self.start_point + + @property + def start_point(self): + return np.array([self.center[0] - self.radius, self.center[1]]) + + @property + def end_point(self): + return np.array([self.center[0] + self.radius, self.center[1]]) + + def transform(self, translate, scale): + self.center = (self.center + translate) * scale + self.radius = self.radius * scale + + def flip(self, axis): + if axis == 'x': + self.center[1] = -self.center[1] + elif axis == 'y': + self.center[0] = -self.center[0] + elif axis == 'xy': + self.center = self.center * -1 + else: + raise ValueError("axis = {}".format(axis)) + + def reverse(self): + pass + + def numericalize(self, n=256): + self.center = self.center.round().clip(min=0, max=n-1).astype(np.int) + self.radius = np.round(self.radius).clip(min=1, max=n-1).astype(np.int) + + def to_vector(self): + vec = [CIRCLE_IDX, self.center[0], self.center[1], PAD_VAL, PAD_VAL, self.radius] + return np.array(vec + [PAD_VAL] * (1 + N_ARGS - len(vec))) + + def draw(self, ax, color): + ap = patches.Circle((self.center[0], self.center[1]), self.radius, + lw=1, fill=None, color=color) + ax.add_patch(ap) + ax.plot(self.center[0], self.center[1], 'ok') + + def sample_points(self, n=32): + angles = np.linspace(0, np.pi * 2, num=n, endpoint=False) + points = np.stack([np.cos(angles), np.sin(angles)], axis=1) * self.radius + self.center[np.newaxis] + return points diff --git a/lib/extrude.py b/lib/extrude.py new file mode 100644 index 0000000..263bb01 --- /dev/null +++ b/lib/extrude.py @@ -0,0 +1,324 @@ +import numpy as np +import random +from .sketch import Profile +from .macro import * +from .math_utils import cartesian2polar, polar2cartesian, polar_parameterization, polar_parameterization_inverse + + +class CoordSystem(object): + """Local coordinate system for sketch plane.""" + def __init__(self, origin, theta, phi, gamma, y_axis=None, is_numerical=False): + self.origin = origin + self._theta = theta # 0~pi + self._phi = phi # -pi~pi + self._gamma = gamma # -pi~pi + self._y_axis = y_axis # (theta, phi) + self.is_numerical = is_numerical + + @property + def normal(self): + return polar2cartesian([self._theta, self._phi]) + + @property + def x_axis(self): + normal_3d, x_axis_3d = polar_parameterization_inverse(self._theta, self._phi, self._gamma) + return x_axis_3d + + @property + def y_axis(self): + if self._y_axis is None: + return np.cross(self.normal, self.x_axis) + return polar2cartesian(self._y_axis) + + @staticmethod + def from_dict(stat): + origin = np.array([stat["origin"]["x"] * 1000, stat["origin"]["y"] * 1000, stat["origin"]["z"] * 1000]) + normal_3d = np.array([stat["z_axis"]["x"], stat["z_axis"]["y"], stat["z_axis"]["z"]]) + x_axis_3d = np.array([stat["x_axis"]["x"], stat["x_axis"]["y"], stat["x_axis"]["z"]]) + y_axis_3d = np.array([stat["y_axis"]["x"], stat["y_axis"]["y"], stat["y_axis"]["z"]]) + theta, phi, gamma = polar_parameterization(normal_3d, x_axis_3d) + return CoordSystem(origin, theta, phi, gamma, y_axis=cartesian2polar(y_axis_3d)) + + @staticmethod + def from_vector(vec, is_numerical=False, n=256): + origin = vec[:3] + theta, phi, gamma = vec[3:] + system = CoordSystem(origin, theta, phi, gamma) + if is_numerical: + system.denumericalize(n) + return system + + def __str__(self): + return "origin: {}, normal: {}, x_axis: {}, y_axis: {}".format( + self.origin.round(4), self.normal.round(4), self.x_axis.round(4), self.y_axis.round(4)) + + def transform(self, translation, scale): + self.origin = (self.origin + translation) * scale + + def numericalize(self, n=256): + """NOTE: shall only be called after normalization""" + # assert np.max(self.origin) <= 1.0 and np.min(self.origin) >= -1.0 # TODO: origin can be out-of-bound! + self.origin = ((self.origin + 1.0) / 2 * n).round().clip(min=0, max=n-1).astype(np.int) + tmp = np.array([self._theta, self._phi, self._gamma]) + self._theta, self._phi, self._gamma = ((tmp / np.pi + 1.0) / 2 * n).round().clip( + min=0, max=n-1).astype(np.int) + self.is_numerical = True + + def denumericalize(self, n=256): + self.origin = self.origin / n * 2 - 1.0 + tmp = np.array([self._theta, self._phi, self._gamma]) + self._theta, self._phi, self._gamma = (tmp / n * 2 - 1.0) * np.pi + self.is_numerical = False + + def to_vector(self): + return np.array([*self.origin, self._theta, self._phi, self._gamma]) + + +class Extrude(object): + """Single extrude operation with corresponding a sketch profile. + NOTE: only support single sketch profile. Extrusion with multiple profiles is decomposed.""" + def __init__(self, profile: Profile, sketch_plane: CoordSystem, + operation, extent_type, extent_one, extent_two, sketch_pos, sketch_size): + """ + Args: + profile (Profile): normalized sketch profile + sketch_plane (CoordSystem): coordinate system for sketch plane + operation (int): index of EXTRUDE_OPERATIONS, see macro.py + extent_type (int): index of EXTENT_TYPE, see macro.py + extent_one (float): extrude distance in normal direction (NOTE: it's negative in some data) + extent_two (float): extrude distance in opposite direction + sketch_pos (np.array): the global 3D position of sketch starting point + sketch_size (float): size of the sketch + """ + self.profile = profile # normalized sketch + self.sketch_plane = sketch_plane + self.operation = operation + self.extent_type = extent_type + self.extent_one = extent_one + self.extent_two = extent_two + + self.sketch_pos = sketch_pos + self.sketch_size = sketch_size + + @staticmethod + def from_dict(all_stat, extrude_id, sketch_dim=256): + """construct Extrude from json data + + Args: + all_stat (dict): all json data + extrude_id (str): entity ID for this extrude + sketch_dim (int, optional): sketch normalization size. Defaults to 256. + + Returns: + list: one or more Extrude instances + """ + extrude_entity = all_stat["entities"][extrude_id] + assert extrude_entity["start_extent"]["type"] == "ProfilePlaneStartDefinition" + + all_skets = [] + n = len(extrude_entity["profiles"]) + for i in range(len(extrude_entity["profiles"])): + sket_id, profile_id = extrude_entity["profiles"][i]["sketch"], extrude_entity["profiles"][i]["profile"] + sket_entity = all_stat["entities"][sket_id] + sket_profile = Profile.from_dict(sket_entity["profiles"][profile_id]) + sket_plane = CoordSystem.from_dict(sket_entity["transform"]) + # normalize profile + #point = sket_profile.start_point + point = np.array([0.0,0.0,0.0]) + sket_pos = point[0] * sket_plane.x_axis + point[1] * sket_plane.y_axis + sket_plane.origin + sket_size = sket_profile.bbox_size + sket_profile.normalize(sketch_dim) + all_skets.append((sket_profile, sket_plane, sket_pos, sket_size)) + + operation = EXTRUDE_OPERATIONS.index(extrude_entity["operation"]) + extent_type = EXTENT_TYPE.index(extrude_entity["extent_type"]) + extent_one = extrude_entity["extent_one"]["distance"]["value"] * 1000 + extent_two = 0.0 + if extrude_entity["extent_type"] == "TwoSidesFeatureExtentType": + extent_two = extrude_entity["extent_two"]["distance"]["value"] * 1000 + + if operation == EXTRUDE_OPERATIONS.index("NewBodyFeatureOperation"): + all_operations = [operation] + [EXTRUDE_OPERATIONS.index("JoinFeatureOperation")] * (n - 1) + else: + all_operations = [operation] * n + + return [Extrude(all_skets[i][0], all_skets[i][1], all_operations[i], extent_type, extent_one, extent_two, + all_skets[i][2], all_skets[i][3]) for i in range(n)] + + @staticmethod + def from_vector(vec, is_numerical=False, n=256): + """vector representation: commands [SOL, ..., SOL, ..., EXT]""" + assert vec[-1][0] == EXT_IDX and vec[0][0] == SOL_IDX + profile_vec = np.concatenate([vec[:-1], EOS_VEC[np.newaxis]]) + profile = Profile.from_vector(profile_vec, is_numerical=is_numerical) + ext_vec = vec[-1][-N_ARGS_EXT:] + + sket_pos = ext_vec[N_ARGS_PLANE:N_ARGS_PLANE + 3] + sket_size = ext_vec[N_ARGS_PLANE + N_ARGS_TRANS - 1] + sket_plane = CoordSystem.from_vector(np.concatenate([sket_pos, ext_vec[:N_ARGS_PLANE]])) + ext_param = ext_vec[-N_ARGS_EXT_PARAM:] + + res = Extrude(profile, sket_plane, int(ext_param[2]), int(ext_param[3]), ext_param[0], ext_param[1], + sket_pos, sket_size) + if is_numerical: + res.denumericalize(n) + return res + + def __str__(self): + s = "Sketch-Extrude pair:" + s += "\n -" + str(self.sketch_plane) + s += "\n -sketch position: {}, sketch size: {}".format(self.sketch_pos.round(4), self.sketch_size.round(4)) + s += "\n -operation:{}, type:{}, extent_one:{}, extent_two:{}".format( + self.operation, self.extent_type, self.extent_one.round(4), self.extent_two.round(4)) + s += "\n -" + str(self.profile) + return s + + def transform(self, translation, scale): + """linear transformation""" + # self.profile.transform(np.array([0, 0]), scale) + self.sketch_plane.transform(translation, scale) + self.extent_one *= scale + self.extent_two *= scale + self.sketch_pos = (self.sketch_pos + translation) * scale + self.sketch_size *= scale + + def numericalize(self, n=256): + """quantize the representation. + NOTE: shall only be called after CADSequence.normalize (the shape lies in unit cube, -1~1)""" + assert -2.0 <= self.extent_one <= 2.0 and -2.0 <= self.extent_two <= 2.0 + self.profile.numericalize(n) + self.sketch_plane.numericalize(n) + self.extent_one = ((self.extent_one + 1.0) / 2 * n).round().clip(min=0, max=n-1).astype(np.int) + self.extent_two = ((self.extent_two + 1.0) / 2 * n).round().clip(min=0, max=n-1).astype(np.int) + self.operation = int(self.operation) + self.extent_type = int(self.extent_type) + + self.sketch_pos = ((self.sketch_pos + 1.0) / 2 * n).round().clip(min=0, max=n-1).astype(np.int) + self.sketch_size = (self.sketch_size / 2 * n).round().clip(min=0, max=n-1).astype(np.int) + + def denumericalize(self, n=256): + """de-quantize the representation.""" + self.extent_one = self.extent_one / n * 2 - 1.0 + self.extent_two = self.extent_two / n * 2 - 1.0 + self.sketch_plane.denumericalize(n) + self.sketch_pos = self.sketch_pos / n * 2 - 1.0 + self.sketch_size = self.sketch_size / n * 2 + + self.operation = self.operation + self.extent_type = self.extent_type + + def flip_sketch(self, axis): + self.profile.flip(axis) + self.profile.normalize() + + def to_vector(self, max_n_loops=6, max_len_loop=15, pad=True): + """vector representation: commands [SOL, ..., SOL, ..., EXT]""" + profile_vec = self.profile.to_vector(max_n_loops, max_len_loop, pad=False) + if profile_vec is None: + return None + sket_plane_orientation = self.sketch_plane.to_vector()[3:] + ext_param = list(sket_plane_orientation) + list(self.sketch_pos) + [self.sketch_size] + \ + [self.extent_one, self.extent_two, self.operation, self.extent_type] + ext_vec = np.array([EXT_IDX, *[PAD_VAL] * N_ARGS_SKETCH, *ext_param]) + vec = np.concatenate([profile_vec[:-1], ext_vec[np.newaxis], profile_vec[-1:]], axis=0) # NOTE: last one is EOS + if pad: + pad_len = max_n_loops * max_len_loop - vec.shape[0] + vec = np.concatenate([vec, EOS_VEC[np.newaxis].repeat(pad_len, axis=0)], axis=0) + return vec + + +class CADSequence(object): + """A CAD modeling sequence, a series of extrude operations.""" + def __init__(self, extrude_seq, bbox=None): + self.seq = extrude_seq + self.bbox = bbox + + @staticmethod + def from_dict(all_stat): + """construct CADSequence from json data""" + seq = [] + for item in all_stat["sequence"]: + if item["type"] == "ExtrudeFeature": + extrude_ops = Extrude.from_dict(all_stat, item["entity"]) + seq.extend(extrude_ops) + bbox_info = all_stat["properties"]["bounding_box"] + max_point = np.array([bbox_info["max_point"]["x"] * 1000, bbox_info["max_point"]["y"] * 1000, bbox_info["max_point"]["z"] * 1000]) + min_point = np.array([bbox_info["min_point"]["x"] * 1000, bbox_info["min_point"]["y"] * 1000, bbox_info["min_point"]["z"] * 1000]) + bbox = np.stack([max_point, min_point], axis=0) + return CADSequence(seq, bbox) + + @staticmethod + def from_vector(vec, is_numerical=False, n=256): + commands = vec[:, 0] + ext_indices = [-1] + np.where(commands == EXT_IDX)[0].tolist() + ext_seq = [] + for i in range(len(ext_indices) - 1): + start, end = ext_indices[i], ext_indices[i + 1] + ext_seq.append(Extrude.from_vector(vec[start+1:end+1], is_numerical, n)) + cad_seq = CADSequence(ext_seq) + return cad_seq + + def __str__(self): + return "" + "\n".join(["({})".format(i) + str(ext) for i, ext in enumerate(self.seq)]) + + def to_vector(self, max_n_ext=10, max_n_loops=6, max_len_loop=15, max_total_len=60, pad=False): + if len(self.seq) > max_n_ext: + return None + vec_seq = [] + for item in self.seq: + vec = item.to_vector(max_n_loops, max_len_loop, pad=False) + if vec is None: + return None + vec = vec[:-1] # last one is EOS, removed + vec_seq.append(vec) + + vec_seq = np.concatenate(vec_seq, axis=0) + vec_seq = np.concatenate([vec_seq, EOS_VEC[np.newaxis]], axis=0) + + # add EOS padding + if pad and vec_seq.shape[0] < max_total_len: + pad_len = max_total_len - vec_seq.shape[0] + vec_seq = np.concatenate([vec_seq, EOS_VEC[np.newaxis].repeat(pad_len, axis=0)], axis=0) + + return vec_seq + + def transform(self, translation, scale): + """linear transformation""" + for item in self.seq: + item.transform(translation, scale) + + def normalize(self, size=1.0): + """(1)normalize the shape into unit cube (-1~1). """ + scale = size * NORM_FACTOR / np.max(np.abs(self.bbox)) + self.transform(0.0, scale) + + def numericalize(self, n=256): + for item in self.seq: + item.numericalize(n) + + def flip_sketch(self, axis): + for item in self.seq: + item.flip_sketch(axis) + + def random_transform(self): + for item in self.seq: + # random transform sketch + scale = random.uniform(0.8, 1.2) + item.profile.transform(-np.array([128, 128]), scale) + translate = np.array([random.randint(-5, 5), random.randint(-5, 5)], dtype=np.int) + 128 + item.profile.transform(translate, 1) + + # random transform and scale extrusion + t = 0.05 + translate = np.array([random.uniform(-t, t), random.uniform(-t, t), random.uniform(-t, t)]) + scale = random.uniform(0.8, 1.2) + # item.sketch_plane.transform(translate, scale) + item.sketch_pos = (item.sketch_pos + translate) * scale + item.extent_one *= random.uniform(0.8, 1.2) + item.extent_two *= random.uniform(0.8, 1.2) + + def random_flip_sketch(self): + for item in self.seq: + flip_idx = random.randint(0, 3) + if flip_idx > 0: + item.flip_sketch(['x', 'y', 'xy'][flip_idx - 1]) diff --git a/lib/file_utils.py b/lib/file_utils.py new file mode 100644 index 0000000..2e18319 --- /dev/null +++ b/lib/file_utils.py @@ -0,0 +1,52 @@ +import os +import json +import logging +import shutil +import csv + + +def save_args(args, save_dir): + param_path = os.path.join(save_dir, 'params.json') + + with open(param_path, 'w') as fp: + json.dump(args.__dict__, fp, indent=4, sort_keys=True) + + +def ensure_dir(path): + """ + create path by first checking its existence, + :param paths: path + :return: + """ + if not os.path.exists(path): + os.makedirs(path) + + +def ensure_dirs(paths): + """ + create paths by first checking their existence + :param paths: list of path + :return: + """ + if isinstance(paths, list) and not isinstance(paths, str): + for path in paths: + ensure_dir(path) + else: + ensure_dir(paths) + + +def remkdir(path): + """ + if dir exists, remove it and create a new one + :param path: + :return: + """ + if os.path.exists(path): + shutil.rmtree(path) + os.makedirs(path) + + +def cycle(iterable): + while True: + for x in iterable: + yield x diff --git a/lib/macro.py b/lib/macro.py new file mode 100644 index 0000000..c4659a4 --- /dev/null +++ b/lib/macro.py @@ -0,0 +1,40 @@ +import numpy as np + +ALL_COMMANDS = ['Line', 'Arc', 'Circle', 'EOS', 'SOL', 'Ext'] +LINE_IDX = ALL_COMMANDS.index('Line') +ARC_IDX = ALL_COMMANDS.index('Arc') +CIRCLE_IDX = ALL_COMMANDS.index('Circle') +EOS_IDX = ALL_COMMANDS.index('EOS') +SOL_IDX = ALL_COMMANDS.index('SOL') +EXT_IDX = ALL_COMMANDS.index('Ext') + +EXTRUDE_OPERATIONS = ["NewBodyFeatureOperation", "JoinFeatureOperation", + "CutFeatureOperation", "IntersectFeatureOperation"] +EXTENT_TYPE = ["OneSideFeatureExtentType", "SymmetricFeatureExtentType", + "TwoSidesFeatureExtentType"] + +PAD_VAL = -1 +N_ARGS_SKETCH = 5 # sketch parameters: x, y, alpha, f, r +N_ARGS_PLANE = 3 # sketch plane orientation: theta, phi, gamma +N_ARGS_TRANS = 4 # sketch plane origin + sketch bbox size: p_x, p_y, p_z, s +N_ARGS_EXT_PARAM = 4 # extrusion parameters: e1, e2, b, u +N_ARGS_EXT = N_ARGS_PLANE + N_ARGS_TRANS + N_ARGS_EXT_PARAM +N_ARGS = N_ARGS_SKETCH + N_ARGS_EXT + +SOL_VEC = np.array([SOL_IDX, *([PAD_VAL] * N_ARGS)]) +EOS_VEC = np.array([EOS_IDX, *([PAD_VAL] * N_ARGS)]) + +CMD_ARGS_MASK = np.array([[1, 1, 0, 0, 0, *[0]*N_ARGS_EXT], # line + [1, 1, 1, 1, 0, *[0]*N_ARGS_EXT], # arc + [1, 1, 0, 0, 1, *[0]*N_ARGS_EXT], # circle + [0, 0, 0, 0, 0, *[0]*N_ARGS_EXT], # EOS + [0, 0, 0, 0, 0, *[0]*N_ARGS_EXT], # SOL + [*[0]*N_ARGS_SKETCH, *[1]*N_ARGS_EXT]]) # Extrude + +NORM_FACTOR = 0.75 # scale factor for normalization to prevent overflow during augmentation + +MAX_N_EXT = 10 # maximum number of extrusion +MAX_N_LOOPS = 6 # maximum number of loops per sketch +MAX_N_CURVES = 15 # maximum number of curves per loop +MAX_TOTAL_LEN = 60 # maximum cad sequence length +ARGS_DIM = 256 diff --git a/lib/math_utils.py b/lib/math_utils.py new file mode 100644 index 0000000..bad46d0 --- /dev/null +++ b/lib/math_utils.py @@ -0,0 +1,367 @@ +import math +import numpy as np + + +def rads_to_degs(rads): + """Convert an angle from radians to degrees""" + return 180 * rads / math.pi + +def distance_of_point_and_plane(origin, x_axis, y_axis, point): + # 确保输入是 numpy 数组 + origin = np.array(origin) + x_axis = np.array(x_axis) + y_axis = np.array(y_axis) + point = np.array(point) + normal_vector = np.cross(x_axis, y_axis) + distance = np.dot(point - origin, normal_vector) + return np.abs(distance) + +def is_point_on_plane(origin, x_axis, y_axis, point, tolerance=1e-10): + """ + 判断一个3D点是否在给定的平面上 + + 参数: + origin - 平面的原点坐标，形状为 (3,) + x_axis - 平面X轴的方向向量，形状为 (3,) + y_axis - 平面Y轴的方向向量，形状为 (3,) + point - 要检查的3D点，形状为 (3,) + tolerance - 容忍度，默认值为 1e-10 + + 返回: + 如果点在平面上返回 True，否则返回 False + """ + # 确保输入是 numpy 数组 + origin = np.array(origin) + x_axis = np.array(x_axis) + y_axis = np.array(y_axis) + point = np.array(point) + + # 计算法向量 + normal_vector = np.cross(x_axis, y_axis) + + # 计算点到平面的距离 + distance = np.dot(point - origin, normal_vector) + + # 判断距离是否在容忍度范围内 + return np.abs(distance) < tolerance + + +def number_to_pi_string(number): + """ + 将数字转换为带有 np.pi 的字符串表示形式 + + 参数: + number - 输入数字，可以是 np.pi 的倍数 + + 返回: + 带有 np.pi 的字符串表示形式 + """ + # 定义一个容忍度来比较浮点数 + tolerance = 1e-10 + + # 预定义一些常见的 π 的倍数及其对应的字符串表示 + pi_factors = { + np.pi: 'np.pi', + np.pi / 2: 'np.pi/2', + np.pi / 3: 'np.pi/3', + np.pi / 4: 'np.pi/4', + np.pi / 6: 'np.pi/6', + 2 * np.pi: '2*np.pi', + 3 * np.pi / 2: '3*np.pi/2', + 3 * np.pi / 4: '3*np.pi/4', + 5 * np.pi / 6: '5*np.pi/6', + 5 * np.pi / 3: '5*np.pi/3', + 7 * np.pi / 6: '7*np.pi/6', + 4 * np.pi / 3: '4*np.pi/3', + } + + # 检查输入数字是否接近这些常见的 π 倍数 + for key, value in pi_factors.items(): + if np.abs(np.abs(number) - key) < tolerance: + return value if number > 0 else '-' + value + + # 如果数字不在预定义的 π 倍数中，返回原数字 + return str(number.round(6)) + +def are_parallel(v1, v2, tol=1e-10): + # 计算叉积 + cross_product = np.cross(v1, v2) + # 判断叉积是否接近于零向量 + return np.all(np.abs(cross_product) < tol) + +def find_circle_center_and_radius(start_point, end_point, mid_point): + # Calculate midpoints of the chords + mid_point_start_end = (start_point + end_point) / 2 + mid_point_start_mid = (start_point + mid_point) / 2 + + # Calculate direction vectors of the chords + direction_start_end = end_point - start_point + direction_start_mid = mid_point - start_point + + # Calculate perpendicular direction vectors + perp_start_end = np.array([-direction_start_end[1], direction_start_end[0]]) + perp_start_mid = np.array([-direction_start_mid[1], direction_start_mid[0]]) + + # Solve for the intersection of the perpendicular bisectors + A = np.array([perp_start_end, -perp_start_mid]).T + b = mid_point_start_mid - mid_point_start_end + + # Solve the linear system + t, s = np.linalg.solve(A, b) + + # Calculate the center + center = mid_point_start_end + t * perp_start_end + + # Calculate the radius + radius = np.linalg.norm(center - start_point) + + return center, radius + +def rotate_vector(vector, axis, angle): + """ + 将一个3D向量绕指定轴逆时针旋转给定角度 + + 参数: + vector - 要旋转的3D向量，形状为 (3,) + axis - 旋转轴，形状为 (3,) + angle - 旋转角度（弧度） + + 返回: + 旋转后的3D向量，形状为 (3,) + """ + # 确保输入是 numpy 数组 + vector = np.array(vector) + axis = np.array(axis) + + # 计算单位轴向量 + axis = axis / np.linalg.norm(axis) + + # 计算旋转矩阵的各个分量 + cos_theta = np.cos(angle) + sin_theta = np.sin(angle) + cross_product = np.cross(axis, vector) + dot_product = np.dot(axis, vector) + + # 计算旋转后的向量 + rotated_vector = (vector * cos_theta + + cross_product * sin_theta + + axis * dot_product * (1 - cos_theta)) + + return rotated_vector + +def calculate_rotation_angle(v1, v2, axis): + """ + 计算从向量 v1 到向量 v2 绕给定轴的逆时针旋转角度 + + 参数: + v1 - 初始向量，形状为 (3,) + v2 - 旋转后的向量，形状为 (3,) + axis - 旋转轴，形状为 (3,) + + 返回: + 旋转角度（弧度），范围 (-pi, pi] + """ + # 确保输入是 numpy 数组 + v1 = np.array(v1) + v2 = np.array(v2) + axis = np.array(axis) + + # 计算单位轴向量 + axis = axis / np.linalg.norm(axis) + + # 计算点积 + dot_product = np.dot(v1, v2) + + # 计算叉积 + cross_product = np.cross(v1, v2) + + # 计算叉积在旋转轴上的投影长度 + projection_length = np.dot(cross_product, axis) + + # 计算向量的范数 + norm_v1 = np.linalg.norm(v1) + norm_v2 = np.linalg.norm(v2) + + # 计算角度的余弦值和正弦值 + cos_theta = dot_product / (norm_v1 * norm_v2) + sin_theta = projection_length / (norm_v1 * norm_v2) + + # 使用 arctan2 计算角度 + angle = np.arctan2(sin_theta, cos_theta) + + return angle + +def map_2d_to_3d(origin, x_axis, y_axis, point): + u, v = point + return origin + u * x_axis + v * y_axis + +def map_3d_to_2d(origin, x_axis, y_axis, point_3d): + """ + 将三维空间中的点转换为二维平面上的点 + + 参数: + origin - 原点坐标 (Ox, Oy, Oz)，形状为 (3,) + x_axis - X轴向量 (Xx, Xy, Xz)，形状为 (3,) + y_axis - Y轴向量 (Yx, Yy, Yz)，形状为 (3,) + point_3d - 三维空间中的点 (Px, Py, Pz)，形状为 (3,) + + 返回: + 二维平面上的点 (u, v)，形状为 (2,) + """ + # 确保输入是 numpy 数组 + origin = np.array(origin) + x_axis = np.array(x_axis) + y_axis = np.array(y_axis) + point_3d = np.array(point_3d) + + # 构建矩阵 A 和向量 b + A = np.vstack([x_axis, y_axis]).T + b = point_3d - origin + + # 求解线性方程组 Ax = b + uv = np.linalg.lstsq(A, b, rcond=None)[0] + + return uv + + +def unit_vector(vector): + """ + 计算给定向量的单位向量 + + 参数: + vector - 输入向量，形状为 (n,) + + 返回: + 单位向量，形状为 (n,) + """ + # 计算向量的范数 + norm = np.linalg.norm(vector) + if norm == 0: + raise ValueError("零向量没有单位向量") + # 计算单位向量 + unit_vector = vector / norm + return unit_vector + + +def find_n_from_x_and_y(x, y): + """ + Given vectors x and y, find a vector n such that y = n × x. + Assumes that n is orthogonal to x. + + Parameters: + x (numpy array): The vector x. + y (numpy array): The vector y. + + Returns: + numpy array: The vector n. + """ + # Step 1: Compute the cross product of x and y to get n' + n_prime = np.cross(x, y) + + # Step 2: Normalize n' to get the unit vector + n_prime_unit = n_prime / np.linalg.norm(n_prime) + + # Step 3: Determine the correct sign of n_prime_unit + # To ensure y = n × x, we should check if the direction is correct + if np.allclose(np.cross(n_prime_unit, x), y): + n = n_prime_unit + else: + n = -n_prime_unit + + return n + +def angle_from_vector_to_x(vec): + """computer the angle (0~2pi) between a unit vector and positive x axis""" + angle = 0.0 + # 2 | 1 + # ------- + # 3 | 4 + if vec[0] >= 0: + if vec[1] >= 0: + # Qadrant 1 + angle = math.asin(vec[1]) + else: + # Qadrant 4 + angle = 2.0 * math.pi - math.asin(-vec[1]) + else: + if vec[1] >= 0: + # Qadrant 2 + angle = math.pi - math.asin(vec[1]) + else: + # Qadrant 3 + angle = math.pi + math.asin(-vec[1]) + return angle + + +def cartesian2polar(vec, with_radius=False): + """convert a vector in cartesian coordinates to polar(spherical) coordinates""" + vec = vec.round(6) + norm = np.linalg.norm(vec) + theta = np.arccos(vec[2] / norm) # (0, pi) + phi = np.arctan(vec[1] / (vec[0] + 1e-15)) # (-pi, pi) # FIXME: -0.0 cannot be identified here + if not with_radius: + return np.array([theta, phi]) + else: + return np.array([theta, phi, norm]) + + +def polar2cartesian(vec): + """convert a vector in polar(spherical) coordinates to cartesian coordinates""" + r = 1 if len(vec) == 2 else vec[2] + theta, phi = vec[0], vec[1] + x = r * np.sin(theta) * np.cos(phi) + y = r * np.sin(theta) * np.sin(phi) + z = r * np.cos(theta) + return np.array([x, y, z]) + + +def rotate_by_x(vec, theta): + mat = np.array([[1, 0, 0], + [0, np.cos(theta), -np.sin(theta)], + [0, np.sin(theta), np.cos(theta)]]) + return np.dot(mat, vec) + + +def rotate_by_y(vec, theta): + mat = np.array([[np.cos(theta), 0, np.sin(theta)], + [0, 1, 0], + [-np.sin(theta), 0, np.cos(theta)]]) + return np.dot(mat, vec) + + +def rotate_by_z(vec, phi): + mat = np.array([[np.cos(phi), -np.sin(phi), 0], + [np.sin(phi), np.cos(phi), 0], + [0, 0, 1]]) + return np.dot(mat, vec) + + +def polar_parameterization(normal_3d, x_axis_3d): + """represent a coordinate system by its rotation from the standard 3D coordinate system + + Args: + normal_3d (np.array): unit vector for normal direction (z-axis) + x_axis_3d (np.array): unit vector for x-axis + + Returns: + theta, phi, gamma: axis-angle rotation + """ + normal_polar = cartesian2polar(normal_3d) + theta = normal_polar[0] + phi = normal_polar[1] + + ref_x = rotate_by_z(rotate_by_y(np.array([1, 0, 0]), theta), phi) + + gamma = np.arccos(np.dot(x_axis_3d, ref_x).round(6)) + if np.dot(np.cross(ref_x, x_axis_3d), normal_3d) < 0: + gamma = -gamma + return theta, phi, gamma + + +def polar_parameterization_inverse(theta, phi, gamma): + """build a coordinate system by the given rotation from the standard 3D coordinate system""" + normal_3d = polar2cartesian([theta, phi]) + ref_x = rotate_by_z(rotate_by_y(np.array([1, 0, 0]), theta), phi) + ref_y = np.cross(normal_3d, ref_x) + x_axis_3d = ref_x * np.cos(gamma) + ref_y * np.sin(gamma) + return normal_3d, x_axis_3d diff --git a/lib/sketch.py b/lib/sketch.py new file mode 100644 index 0000000..23aed34 --- /dev/null +++ b/lib/sketch.py @@ -0,0 +1,263 @@ +import numpy as np +import matplotlib +matplotlib.use('TkAgg') +import matplotlib.pyplot as plt +from .curves import * +from .macro import * + + +########################## base ########################### +class SketchBase(object): + """Base class for sketch (a collection of curves). """ + def __init__(self, children, reorder=True): + self.children = children + + if reorder: + self.reorder() + + @staticmethod + def from_dict(stat): + """construct sketch from json data + + Args: + stat (dict): dict from json data + """ + raise NotImplementedError + + @staticmethod + def from_vector(vec, start_point, is_numerical=True): + """construct sketch from vector representation + + Args: + vec (np.array): (seq_len, n_args) + start_point (np.array): (2, ). If none, implicitly defined as the last end point. + """ + raise NotImplementedError + + def reorder(self): + """rearrange the curves to follow counter-clockwise direction""" + raise NotImplementedError + + @property + def start_point(self): + return self.children[0].start_point + + @property + def end_point(self): + return self.children[-1].end_point + + @property + def bbox(self): + """compute bounding box (min/max points) of the sketch""" + all_points = np.concatenate([child.bbox for child in self.children], axis=0) + return np.stack([np.min(all_points, axis=0), np.max(all_points, axis=0)], axis=0) + + @property + def bbox_size(self): + """compute bounding box size (max of height and width)""" + bbox_min, bbox_max = self.bbox[0], self.bbox[1] + bbox_size = np.max(np.abs(np.concatenate([bbox_max - self.start_point, bbox_min - self.start_point]))) + return bbox_size + + @property + def global_trans(self): + """start point + sketch size (bbox_size)""" + return np.concatenate([self.start_point, np.array([self.bbox_size])]) + + def transform(self, translate, scale): + """linear transformation""" + for child in self.children: + child.transform(translate, scale) + + def flip(self, axis): + for child in self.children: + child.flip(axis) + self.reorder() + + def numericalize(self, n=256): + """quantize curve parameters into integers""" + for child in self.children: + child.numericalize(n) + + def normalize(self, size=256): + """normalize within the given size, with start_point in the middle center""" + cur_size = self.bbox_size + scale = (size / 2 * NORM_FACTOR - 1) / cur_size # prevent potential overflow if data augmentation applied + #self.transform(-self.start_point, scale) + #self.transform(np.array((size / 2, size / 2)), 1) + + def denormalize(self, bbox_size, size=256): + """inverse procedure of normalize method""" + scale = bbox_size / (size / 2 * NORM_FACTOR - 1) + #self.transform(-np.array((size / 2, size / 2)), scale) + + def to_vector(self): + """convert to vector representation""" + raise NotImplementedError + + def draw(self, ax): + """draw sketch on matplotlib ax""" + raise NotImplementedError + + def to_image(self): + """convert to image""" + fig, ax = plt.subplots() + self.draw(ax) + ax.axis('equal') + fig.canvas.draw() + X = np.array(fig.canvas.renderer.buffer_rgba())[:, :, :3] + plt.close(fig) + return X + + def sample_points(self, n=32): + """uniformly sample points from the sketch""" + raise NotImplementedError + + +####################### loop & profile ####################### +class Loop(SketchBase): + """Sketch loop, a sequence of connected curves.""" + @staticmethod + def from_dict(stat): + all_curves = [construct_curve_from_dict(item) for item in stat['profile_curves']] + this_loop = Loop(all_curves) + this_loop.is_outer = stat['is_outer'] + return this_loop + + def __str__(self): + return "Loop:" + "\n -" + "\n -".join([str(curve) for curve in self.children]) + + @staticmethod + def from_vector(vec, start_point=None, is_numerical=True): + all_curves = [] + if start_point is None: + # FIXME: explicit for loop can be avoided here + for i in range(vec.shape[0]): + if vec[i][0] == EOS_IDX: + start_point = vec[i - 1][1:3] + break + for i in range(vec.shape[0]): + type = vec[i][0] + if type == SOL_IDX: + continue + elif type == EOS_IDX: + break + else: + curve = construct_curve_from_vector(vec[i], start_point, is_numerical=is_numerical) + start_point = vec[i][1:3] # current curve's end_point serves as next curve's start_point + all_curves.append(curve) + return Loop(all_curves) + + def reorder(self): + """reorder by starting left most and counter-clockwise""" + if len(self.children) <= 1: + return + + start_curve_idx = -1 + sx, sy = 10000, 10000 + + # correct start-end point order + if np.allclose(self.children[0].start_point, self.children[1].start_point) or \ + np.allclose(self.children[0].start_point, self.children[1].end_point): + self.children[0].reverse() + + # correct start-end point order and find left-most point + for i, curve in enumerate(self.children): + if i < len(self.children) - 1 and np.allclose(curve.end_point, self.children[i + 1].end_point): + self.children[i + 1].reverse() + if round(curve.start_point[0], 6) < round(sx, 6) or \ + (round(curve.start_point[0], 6) == round(sx, 6) and round(curve.start_point[1], 6) < round(sy, 6)): + start_curve_idx = i + sx, sy = curve.start_point + + self.children = self.children[start_curve_idx:] + self.children[:start_curve_idx] + + # ensure mostly counter-clock wise + if isinstance(self.children[0], Circle) or isinstance(self.children[-1], Circle): # FIXME: hard-coded + return + start_vec = self.children[0].direction() + end_vec = self.children[-1].direction(from_start=False) + if np.cross(end_vec, start_vec) <= 0: + for curve in self.children: + curve.reverse() + self.children.reverse() + + def to_vector(self, max_len=None, add_sol=True, add_eos=True): + loop_vec = np.stack([curve.to_vector() for curve in self.children], axis=0) + if add_sol: + loop_vec = np.concatenate([SOL_VEC[np.newaxis], loop_vec], axis=0) + if add_eos: + loop_vec = np.concatenate([loop_vec, EOS_VEC[np.newaxis]], axis=0) + if max_len is None: + return loop_vec + + if loop_vec.shape[0] > max_len: + return None + elif loop_vec.shape[0] < max_len: + pad_vec = np.tile(EOS_VEC, max_len - loop_vec.shape[0]).reshape((-1, len(EOS_VEC))) + loop_vec = np.concatenate([loop_vec, pad_vec], axis=0) # (max_len, 1 + N_ARGS) + return loop_vec + + def draw(self, ax): + colors = ['red', 'blue', 'green', 'brown', 'pink', 'yellow', 'purple', 'black'] * 10 + for i, curve in enumerate(self.children): + curve.draw(ax, colors[i]) + + def sample_points(self, n=32): + points = np.stack([curve.sample_points(n) for curve in self.children], axis=0) # (n_curves, n, 2) + return points + + +class Profile(SketchBase): + """Sketch profile，a closed region formed by one or more loops. + The outer-most loop is placed at first.""" + @staticmethod + def from_dict(stat): + all_loops = [Loop.from_dict(item) for item in stat['loops']] + return Profile(all_loops) + + def __str__(self): + return "Profile:" + "\n -".join([str(loop) for loop in self.children]) + + @staticmethod + def from_vector(vec, start_point=None, is_numerical=True): + all_loops = [] + command = vec[:, 0] + end_idx = command.tolist().index(EOS_IDX) + indices = np.where(command[:end_idx] == SOL_IDX)[0].tolist() + [end_idx] + for i in range(len(indices) - 1): + loop_vec = vec[indices[i]:indices[i + 1]] + loop_vec = np.concatenate([loop_vec, EOS_VEC[np.newaxis]], axis=0) + if loop_vec[0][0] == SOL_IDX and loop_vec[1][0] not in [SOL_IDX, EOS_IDX]: + all_loops.append(Loop.from_vector(loop_vec, is_numerical=is_numerical)) + return Profile(all_loops) + + def reorder(self): + if len(self.children) <= 1: + return + all_loops_bbox_min = np.stack([loop.bbox[0] for loop in self.children], axis=0).round(6) + ind = np.lexsort(all_loops_bbox_min.transpose()[[1, 0]]) + self.children = [self.children[i] for i in ind] + + def draw(self, ax): + for i, loop in enumerate(self.children): + loop.draw(ax) + ax.text(loop.start_point[0], loop.start_point[1], str(i)) + + def to_vector(self, max_n_loops=None, max_len_loop=None, pad=True): + loop_vecs = [loop.to_vector(None, add_eos=False) for loop in self.children] + if max_n_loops is not None and len(loop_vecs) > max_n_loops: + return None + for vec in loop_vecs: + if max_len_loop is not None and vec.shape[0] > max_len_loop: + return None + profile_vec = np.concatenate(loop_vecs, axis=0) + profile_vec = np.concatenate([profile_vec, EOS_VEC[np.newaxis]], axis=0) + if pad: + pad_len = max_n_loops * max_len_loop - profile_vec.shape[0] + profile_vec = np.concatenate([profile_vec, EOS_VEC[np.newaxis].repeat(pad_len, axis=0)], axis=0) + return profile_vec + + def sample_points(self, n=32): + points = np.concatenate([loop.sample_points(n) for loop in self.children], axis=0) + return points diff --git a/lib/timeout.py b/lib/timeout.py new file mode 100644 index 0000000..b7363d4 --- /dev/null +++ b/lib/timeout.py @@ -0,0 +1,31 @@ + +import signal + +# 定义超时异常 +class TimeoutException(Exception): + pass + +# 处理超时信号 +def handler(signum, frame): + raise TimeoutException() + +# 设置超时时间（秒） +timeout = 30 + +# 使用装饰器设置超时 +def timeout_decorator(func): + def wrapper(*args, **kwargs): + # 设置信号处理器 + signal.signal(signal.SIGALRM, handler) + # 启动闹钟 + signal.alarm(timeout) + try: + result = func(*args, **kwargs) + except TimeoutException: + print("Function timed out!") + result = None + finally: + # 关闭闹钟 + signal.alarm(0) + return result + return wrapper \ No newline at end of file diff --git a/lib/visualize.py b/lib/visualize.py new file mode 100644 index 0000000..f2dc2d6 --- /dev/null +++ b/lib/visualize.py @@ -0,0 +1,197 @@ +from OCC.Core.gp import gp_Pnt, gp_Dir, gp_Circ, gp_Pln, gp_Vec, gp_Ax3, gp_Ax2, gp_Lin +from OCC.Core.BRepBuilderAPI import (BRepBuilderAPI_MakeEdge, BRepBuilderAPI_MakeFace, BRepBuilderAPI_MakeWire) +from OCC.Core.BRepPrimAPI import BRepPrimAPI_MakePrism +from OCC.Core.BRepAlgoAPI import BRepAlgoAPI_Cut, BRepAlgoAPI_Fuse, BRepAlgoAPI_Common +from OCC.Core.GC import GC_MakeArcOfCircle +from OCC.Extend.DataExchange import write_stl_file +from OCC.Core.Bnd import Bnd_Box +from OCC.Core.BRepBndLib import brepbndlib_Add +from copy import copy +from .extrude import * +from .sketch import Loop, Profile +from .curves import * +import os +import trimesh +from trimesh.sample import sample_surface +import random + +from OCC.Core.Quantity import Quantity_Color, Quantity_TOC_RGB +from OCC.Core.TDocStd import TDocStd_Document +from OCC.Core.XCAFDoc import XCAFDoc_DocumentTool, XCAFDoc_ColorGen +from OCC.Core.AIS import AIS_Shape + +# 创建不同颜色 +red = Quantity_Color(1.0, 0.1, 0.1, Quantity_TOC_RGB) +green = Quantity_Color(0.1, 1.0, 0.1, Quantity_TOC_RGB) +blue = Quantity_Color(0.1, 0.1, 1.0, Quantity_TOC_RGB) +white = Quantity_Color(1.0, 1.0, 1.0, Quantity_TOC_RGB) +gray = Quantity_Color(0.1, 0.1, 0.1, Quantity_TOC_RGB) + +import random +def generate_random_color(): + r = random.uniform(0, 1) + g = random.uniform(0, 1) + b = random.uniform(0, 1) + return Quantity_Color(r, g, b, Quantity_TOC_RGB) +import numpy as np +def color_distance(color1, color2): + rgb1 = np.array([color1.Red(), color1.Green(), color1.Blue()]) + rgb2 = np.array([color2.Red(), color2.Green(), color2.Blue()]) + return np.linalg.norm(rgb1 - rgb2) + +def vec2CADsolid(vec, is_numerical=True, n=256): + cad = CADSequence.from_vector(vec, is_numerical=is_numerical, n=256) + cad = create_CAD(cad) + return cad + + +from copy import deepcopy +def create_CAD_index(doc: TDocStd_Document, cad_seq: CADSequence, index = 0, color = red): + """create a 3D CAD model from CADSequence. Only support extrude with boolean operation.""" + if len(cad_seq.seq) != 1: + _ = create_CAD(doc, cad_seq) + extrude_op = cad_seq.seq[index] + profile = copy(extrude_op.profile) # use copy to prevent changing extrude_op internally + profile.denormalize(extrude_op.sketch_size) + + sketch_plane = copy(extrude_op.sketch_plane) + #sketch_plane.origin = extrude_op.sketch_pos + + face = create_profile_face(profile, sketch_plane) + normal = gp_Dir(*extrude_op.sketch_plane.normal) + ext_vec = gp_Vec(normal).Multiplied(extrude_op.extent_one) + body = BRepPrimAPI_MakePrism(face, ext_vec).Shape() + + shape_tool = XCAFDoc_DocumentTool.ShapeTool(doc.Main()) + color_tool = XCAFDoc_DocumentTool.ColorTool(doc.Main()) + label = shape_tool.AddShape(body) + color_tool.SetColor(label, red, XCAFDoc_ColorGen) + + if extrude_op.extent_type == EXTENT_TYPE.index("SymmetricFeatureExtentType"): + body_sym = BRepPrimAPI_MakePrism(face, ext_vec.Reversed()).Shape() + body = BRepAlgoAPI_Fuse(body, body_sym).Shape() + label_sym = shape_tool.AddShape(body_sym) + color_tool.SetColor(label_sym, red, XCAFDoc_ColorGen) + if extrude_op.extent_type == EXTENT_TYPE.index("TwoSidesFeatureExtentType"): + ext_vec = gp_Vec(normal.Reversed()).Multiplied(extrude_op.extent_two) + body_two = BRepPrimAPI_MakePrism(face, ext_vec).Shape() + body = BRepAlgoAPI_Fuse(body, body_two).Shape() + label_two = shape_tool.AddShape(body_two) + color_tool.SetColor(label_two, red, XCAFDoc_ColorGen) + return body + +def create_CAD(doc: TDocStd_Document, cad_seq: CADSequence): + """create a 3D CAD model from CADSequence. Only support extrude with boolean operation.""" + body = create_by_extrude(doc, cad_seq.seq[0]) + + for extrude_op in cad_seq.seq[1:]: + new_body = create_by_extrude(doc, extrude_op) + + if extrude_op.operation == EXTRUDE_OPERATIONS.index("NewBodyFeatureOperation") or \ + extrude_op.operation == EXTRUDE_OPERATIONS.index("JoinFeatureOperation"): + body = BRepAlgoAPI_Fuse(body, new_body).Shape() + elif extrude_op.operation == EXTRUDE_OPERATIONS.index("CutFeatureOperation"): + body = BRepAlgoAPI_Cut(body, new_body).Shape() + elif extrude_op.operation == EXTRUDE_OPERATIONS.index("IntersectFeatureOperation"): + body = BRepAlgoAPI_Common(body, new_body).Shape() + + shape_tool = XCAFDoc_DocumentTool.ShapeTool(doc.Main()) + _ = shape_tool.AddShape(body) + return body + + +def create_by_extrude(doc: TDocStd_Document, extrude_op: Extrude): + """create a solid body from Extrude instance.""" + profile = copy(extrude_op.profile) # use copy to prevent changing extrude_op internally + profile.denormalize(extrude_op.sketch_size) + + sketch_plane = copy(extrude_op.sketch_plane) + #sketch_plane.origin = extrude_op.sketch_pos + + face = create_profile_face(profile, sketch_plane) + normal = gp_Dir(*extrude_op.sketch_plane.normal) + ext_vec = gp_Vec(normal).Multiplied(extrude_op.extent_one) + body = BRepPrimAPI_MakePrism(face, ext_vec).Shape() + if extrude_op.extent_type == EXTENT_TYPE.index("SymmetricFeatureExtentType"): + body_sym = BRepPrimAPI_MakePrism(face, ext_vec.Reversed()).Shape() + body = BRepAlgoAPI_Fuse(body, body_sym).Shape() + if extrude_op.extent_type == EXTENT_TYPE.index("TwoSidesFeatureExtentType"): + ext_vec = gp_Vec(normal.Reversed()).Multiplied(extrude_op.extent_two) + body_two = BRepPrimAPI_MakePrism(face, ext_vec).Shape() + body = BRepAlgoAPI_Fuse(body, body_two).Shape() + + return body + + +def create_profile_face(profile: Profile, sketch_plane: CoordSystem): + """create a face from a sketch profile and the sketch plane""" + origin = gp_Pnt(*sketch_plane.origin) + normal = gp_Dir(*sketch_plane.normal) + x_axis = gp_Dir(*sketch_plane.x_axis) + gp_face = gp_Pln(gp_Ax3(origin, normal, x_axis)) + + all_loops = [create_loop_3d(loop, sketch_plane) for loop in profile.children] + topo_face = BRepBuilderAPI_MakeFace(gp_face, all_loops[0]) + for loop in all_loops[1:]: + topo_face.Add(loop.Reversed()) + return topo_face.Face() + + +def create_loop_3d(loop: Loop, sketch_plane: CoordSystem): + """create a 3D sketch loop""" + topo_wire = BRepBuilderAPI_MakeWire() + for curve in loop.children: + topo_edge = create_edge_3d(curve, sketch_plane) + if topo_edge == -1: # omitted + continue + topo_wire.Add(topo_edge) + return topo_wire.Wire() + + +def create_edge_3d(curve: CurveBase, sketch_plane: CoordSystem): + """create a 3D edge""" + if isinstance(curve, Line): + if np.allclose(curve.start_point, curve.end_point): + return -1 + start_point = point_local2global(curve.start_point, sketch_plane) + end_point = point_local2global(curve.end_point, sketch_plane) + topo_edge = BRepBuilderAPI_MakeEdge(start_point, end_point) + elif isinstance(curve, Circle): + center = point_local2global(curve.center, sketch_plane) + axis = gp_Dir(*sketch_plane.normal) + gp_circle = gp_Circ(gp_Ax2(center, axis), abs(float(curve.radius))) + topo_edge = BRepBuilderAPI_MakeEdge(gp_circle) + elif isinstance(curve, Arc): + # print(curve.start_point, curve.mid_point, curve.end_point) + start_point = point_local2global(curve.start_point, sketch_plane) + mid_point = point_local2global(curve.mid_point, sketch_plane) + end_point = point_local2global(curve.end_point, sketch_plane) + arc = GC_MakeArcOfCircle(start_point, mid_point, end_point).Value() + topo_edge = BRepBuilderAPI_MakeEdge(arc) + else: + raise NotImplementedError(type(curve)) + return topo_edge.Edge() + + +def point_local2global(point, sketch_plane: CoordSystem, to_gp_Pnt=True): + """convert point in sketch plane local coordinates to global coordinates""" + g_point = point[0] * sketch_plane.x_axis + point[1] * sketch_plane.y_axis + sketch_plane.origin + if to_gp_Pnt: + return gp_Pnt(*g_point) + return g_point + + +def CADsolid2pc(shape, n_points, name=None): + """convert opencascade solid to point clouds""" + bbox = Bnd_Box() + brepbndlib_Add(shape, bbox) + if bbox.IsVoid(): + raise ValueError("box check failed") + + if name is None: + name = random.randint(100000, 999999) + write_stl_file(shape, "tmp_out_{}.stl".format(name)) + out_mesh = trimesh.load("tmp_out_{}.stl".format(name)) + os.system("rm tmp_out_{}.stl".format(name)) + out_pc, _ = sample_surface(out_mesh, n_points) + return out_pc diff --git a/py2step.py b/py2step.py new file mode 100644 index 0000000..af067ca --- /dev/null +++ b/py2step.py @@ -0,0 +1,214 @@ +import os +import glob +import numpy as np +import argparse +import sys +sys.path.append(".") + +from lib.DataExchange import write_step_file +from lib.visualize import create_CAD +from lib.curves import Line, Arc, Circle +from lib.sketch import Loop, Profile +from lib.math_utils import * +from lib.extrude import CoordSystem, Extrude, CADSequence + +# Curves +def add_line(start, end): + start = np.array(start) + end = np.array(end) + return Line(start, end) + +def add_arc(start, end, mid): + # get radius and center point + start = np.array(start) + end = np.array(end) + mid = np.array(mid) + center, radius = find_circle_center_and_radius(start, end, mid) + def get_angles_counterclockwise(eps=1e-8): + c2s_vec = (start - center) / (np.linalg.norm(start - center) + eps) + c2m_vec = (mid - center) / (np.linalg.norm(mid - center) + eps) + c2e_vec = (end - center) / (np.linalg.norm(end - center) + eps) + angle_s, angle_m, angle_e = angle_from_vector_to_x(c2s_vec), angle_from_vector_to_x(c2m_vec), \ + angle_from_vector_to_x(c2e_vec) + angle_s, angle_e = min(angle_s, angle_e), max(angle_s, angle_e) + if not angle_s < angle_m < angle_e: + angle_s, angle_e = angle_e - np.pi * 2, angle_s + return angle_s, angle_e + angle_s, angle_e = get_angles_counterclockwise() + return Arc(start, end, center, radius, start_angle=angle_s, end_angle=angle_e, mid_point=mid) + +def add_circle(center, radius): + center = np.array(center) + return Circle(center, radius) + +# Loops +def add_loop(curves): + res = Loop(curves) + res.reorder() + def autofix(loop): + if len(loop.children) <= 1: + return + if isinstance(loop.children[0], Circle): + return + for i in range(0, len(loop.children) - 1): + if not np.allclose(loop.children[i].end_point, loop.children[i+1].start_point): + loop.children[i+1].start_point = loop.children[i].end_point + print("warning: fixing loop") + if not np.allclose(loop.children[len(loop.children) - 1].end_point, loop.children[0].start_point): + loop.children[len(loop.children) - 1].start_point = loop.children[0].end_point + print("warning: fixing loop") + + autofix(res) + return res + +# Sketch-Profile +def add_profile(loops): + return Profile(loops) + +def add_sketchplane(origin, normal, x_axis):#, y_axis): + #print(origin, normal, x_axis) + origin = np.array(origin) + normal = np.array(normal) + x_axis = np.array(x_axis) + y_axis = find_n_from_x_and_y(normal, x_axis) + # get theta and phi + theta, phi, gamma = polar_parameterization(normal, x_axis) + #print(normal_axis, x_axis) + #print(theta, phi, gamma) + return CoordSystem(origin, theta, phi, gamma, y_axis=cartesian2polar(y_axis)) + +def add_sketchplane_ref(extrude: Extrude, origin, type: str, line: Line = None, reverse = False, angle=0): + origin = np.array(origin) + types_dict = ["sameplane", "extent_one", "extent_two", "line"] + """ + sameplane: 参考Extrude的SketchPlane，angle是以normal为轴，向下看逆时针角度，原点是在SketchPlane内的2D相对坐标。reverse只在最后反转normal，其他参考并不反转。 + line：默认normal为y轴，line start to end 为x轴，它们的叉积为方向向量normal'，原点是在该默认平面的2D相对坐标。angle是以normal'为轴，向下看逆时针角度。reverse只在最后反转normal'，其他参考并不反转。 + """ + if type not in types_dict: + raise ValueError + ref_plane = extrude.sketch_plane + ref_x_axis = unit_vector(ref_plane.x_axis) + ref_y_axis = unit_vector(ref_plane.y_axis) + ref_n_axis = unit_vector(ref_plane.normal) + ref_origin = ref_plane.origin + if type == "sameplane": + real_origin = map_2d_to_3d(ref_origin, ref_x_axis, ref_y_axis, origin) + elif type == "extent_one": + ref_origin_ = ref_origin + extrude.extent_one * ref_n_axis + real_origin = map_2d_to_3d(ref_origin_, ref_x_axis, ref_y_axis, origin) + elif type == "extent_two": + ref_origin_ = ref_origin - extrude.extent_two * ref_n_axis + real_origin = map_2d_to_3d(ref_origin_, ref_x_axis, ref_y_axis, origin) + if type in types_dict[:3]: + real_x_axis = rotate_vector(ref_x_axis, ref_n_axis, angle) + return add_sketchplane(real_origin, ref_n_axis if not reverse else -ref_n_axis, real_x_axis) + if type == "line": + if line is None: + raise TypeError + start_point = map_2d_to_3d(ref_origin, ref_x_axis, ref_y_axis, line.start_point) + end_point = map_2d_to_3d(ref_origin, ref_x_axis, ref_y_axis, line.end_point) + default_x_axis = unit_vector(end_point - start_point) + real_origin = map_2d_to_3d(start_point, default_x_axis, ref_n_axis, origin) # ref_n_axis is y axis of default plane + default_normal = find_n_from_x_and_y(default_x_axis, ref_n_axis) + real_x_axis = rotate_vector(default_x_axis, default_normal, angle) + return add_sketchplane(real_origin, default_normal if not reverse else -default_normal, real_x_axis) + raise ValueError + +class Sketch(object): + def __init__(self, sketch_plane, profile, sketch_position, sketch_size): + self.sketch_plane = sketch_plane + self.profile = profile + self.sketch_position = sketch_position + self.sketch_size = sketch_size + +def add_sketch(sketch_plane, profile, sketch_position=[0.0,0.0,0.0], sketch_size=0): + return Sketch(sketch_plane, profile, np.array(sketch_position), sketch_size) + +cad_seq = [] +def add_extrude(sketch: Sketch, operation, type, extent_one, extent_two): + res = Extrude( + sketch.profile, + sketch.sketch_plane, + np.intc(operation), np.intc(type), np.double(extent_one), np.double(extent_two), + np.double(sketch.sketch_position), + np.double(sketch.sketch_size) + ) + cad_seq.append(res) + return res + +from OCC.Core.TDocStd import TDocStd_Document +from OCC.Core.TCollection import TCollection_ExtendedString + +def _process(path, path_o): + global cad_seq + cad_seq.clear() + with open(path, 'r') as file: + codes = file.read() + # 执行读取到的代码 + + codes = codes.split("\n") + i = 0 + last_curves_list, last_loops_list = None, None + while i < len(codes): + code = codes[i] + last_curve_name, last_loop_name = None, None + if codes[i].startswith('Curves'): + last_curves_list = codes[i].split('=', 1)[0].split()[0] + elif codes[i].startswith('Loops'): + last_loops_list = codes[i].split('=', 1)[0].split()[0] + elif codes[i].startswith('Arc') or codes[i].startswith('Line') or codes[i].startswith('Circle'): + last_curve_name = codes[i].split('=', 1)[0].split()[0] + elif codes[i].startswith('Loop'): + last_loop_name = codes[i].split('=', 1)[0].split()[0] + for j in range(i + 1, len(codes)): + if codes[j].startswith('\t') or codes[j].startswith(' '): + code += '\n' + codes[j] + i += 1 + else: + break + exec(code) + if last_curve_name is not None: + exec(f"{last_curves_list}.append({last_curve_name})") + elif last_loop_name is not None: + exec(f"{last_loops_list}.append({last_loop_name})") + + i += 1 + doc_name = TCollection_ExtendedString("pythonocc-doc") + doc = TDocStd_Document(doc_name) + cad = CADSequence(cad_seq) + #print(cad) + out_shape = create_CAD(doc, cad) + + write_step_file(doc, path_o) + +error_list = [] + +from lib.file_utils import ensure_dir + +parser = argparse.ArgumentParser() +parser.add_argument('--src', type=str, required=True, help="source folder") +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("py")))) +save_dir = args.src + "_step" if args.outputs is None else args.outputs +ensure_dir(save_dir) + +import traceback + +for path in out_paths: + name = path.split("/")[-1].split(".")[0] + try: + save_path = os.path.join(save_dir, name + ".step") + _process(path, save_path) + + except Exception as e: + print("load and create failed.") + traceback.print_exc() + print(name) + error_list.append(name) + +for error in error_list: + print(error) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..b9725f9 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +matplotlib==3.9.1 +numpy==1.26.4 +Pillow==8.3.2 +trimesh==3.23.5 \ No newline at end of file