first commit

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)

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

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])

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

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

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

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

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

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