Andrea Pozzetti
2859c0a3e5
Split SECAD-Net nodes and utilities from the CADabra monorepo. Voxel-based sketch-extrude CAD reconstruction. Original source: https://github.com/PozzettiAndrea/ComfyUI-CADabra
600 lines
19 KiB
Python
600 lines
19 KiB
Python
# SPDX-License-Identifier: GPL-3.0-or-later
|
|
# Originally from ComfyUI-CADabra: https://github.com/PozzettiAndrea/ComfyUI-CADabra
|
|
# Copyright (C) 2025 ComfyUI-CADabra Contributors
|
|
|
|
"""
|
|
OpenCASCADE (OCC) geometry conversion utilities.
|
|
|
|
Converts fitted surface parameters to OCC geometry objects for CAD export.
|
|
"""
|
|
import numpy as np
|
|
|
|
try:
|
|
from OCC.Core.gp import gp_Pnt, gp_Dir, gp_Vec, gp_Ax1, gp_Ax2, gp_Ax3
|
|
from OCC.Core.Geom import (
|
|
Geom_Plane, Geom_SphericalSurface, Geom_CylindricalSurface,
|
|
Geom_ConicalSurface, Geom_BSplineSurface
|
|
)
|
|
from OCC.Core.GeomAPI import (
|
|
GeomAPI_PointsToBSplineSurface, GeomAPI_ProjectPointOnSurf
|
|
)
|
|
from OCC.Core.TColgp import TColgp_Array2OfPnt, TColgp_Array1OfPnt
|
|
from OCC.Core.TColStd import TColStd_Array1OfReal, TColStd_Array1OfInteger
|
|
from OCC.Core.BRepBuilderAPI import BRepBuilderAPI_MakeFace, BRepBuilderAPI_MakeEdge, BRepBuilderAPI_MakeWire
|
|
from OCC.Core.TopoDS import TopoDS_Compound, TopoDS_Face
|
|
from OCC.Core.BRep import BRep_Builder
|
|
from OCC.Core.STEPControl import STEPControl_Writer, STEPControl_AsIs
|
|
from OCC.Core.IFSelect import IFSelect_RetDone
|
|
HAS_OCC = True
|
|
except ImportError as e:
|
|
HAS_OCC = False
|
|
print(f"[occ_conversion] WARNING: OCC import failed: {e}")
|
|
|
|
# Note: We use polyline edges for boundary wires, no need for GeomAPI_PointsToBSplineCurve
|
|
|
|
|
|
def check_occ():
|
|
"""Check if OCC is available."""
|
|
if not HAS_OCC:
|
|
raise ImportError(
|
|
"OpenCASCADE not available. Install pythonocc-core or use cadquery."
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# Primitive to OCC conversion
|
|
# =============================================================================
|
|
|
|
def plane_to_occ(axis, distance, center=None):
|
|
"""
|
|
Convert plane parameters to OCC Geom_Plane.
|
|
|
|
Args:
|
|
axis: Normal vector [nx, ny, nz]
|
|
distance: Distance from origin along normal
|
|
center: Optional center point [x, y, z]. If None, computed from axis*distance
|
|
|
|
Returns:
|
|
Geom_Plane object
|
|
"""
|
|
check_occ()
|
|
|
|
axis = np.array(axis, dtype=np.float64).flatten() # Ensure 1D float array
|
|
axis = axis / np.linalg.norm(axis) # Normalize
|
|
|
|
if center is None:
|
|
center = axis * distance
|
|
else:
|
|
center = np.array(center, dtype=np.float64).flatten() # Also flatten center
|
|
|
|
normal = gp_Dir(float(axis[0]), float(axis[1]), float(axis[2]))
|
|
point = gp_Pnt(float(center[0]), float(center[1]), float(center[2]))
|
|
|
|
return Geom_Plane(point, normal)
|
|
|
|
|
|
def sphere_to_occ(center, radius):
|
|
"""
|
|
Convert sphere parameters to OCC Geom_SphericalSurface.
|
|
|
|
Args:
|
|
center: Center point [x, y, z]
|
|
radius: Sphere radius
|
|
|
|
Returns:
|
|
Geom_SphericalSurface object
|
|
"""
|
|
check_occ()
|
|
|
|
center = np.array(center)
|
|
ax = gp_Ax3(
|
|
gp_Pnt(float(center[0]), float(center[1]), float(center[2])),
|
|
gp_Dir(0, 0, 1) # Z-up orientation
|
|
)
|
|
|
|
return Geom_SphericalSurface(ax, float(radius))
|
|
|
|
|
|
def cylinder_to_occ(axis, center, radius):
|
|
"""
|
|
Convert cylinder parameters to OCC Geom_CylindricalSurface.
|
|
|
|
Args:
|
|
axis: Cylinder axis direction [ax, ay, az]
|
|
center: Point on cylinder axis [x, y, z]
|
|
radius: Cylinder radius
|
|
|
|
Returns:
|
|
Geom_CylindricalSurface object
|
|
"""
|
|
check_occ()
|
|
|
|
axis = np.array(axis)
|
|
axis = axis / np.linalg.norm(axis)
|
|
center = np.array(center)
|
|
|
|
ax = gp_Ax3(
|
|
gp_Pnt(float(center[0]), float(center[1]), float(center[2])),
|
|
gp_Dir(float(axis[0]), float(axis[1]), float(axis[2]))
|
|
)
|
|
|
|
return Geom_CylindricalSurface(ax, float(radius))
|
|
|
|
|
|
def cone_to_occ(apex, axis, half_angle):
|
|
"""
|
|
Convert cone parameters to OCC Geom_ConicalSurface.
|
|
|
|
Args:
|
|
apex: Cone apex point [x, y, z]
|
|
axis: Cone axis direction [ax, ay, az]
|
|
half_angle: Cone half-angle in radians
|
|
|
|
Returns:
|
|
Geom_ConicalSurface object
|
|
"""
|
|
check_occ()
|
|
|
|
apex = np.array(apex)
|
|
axis = np.array(axis)
|
|
axis = axis / np.linalg.norm(axis)
|
|
|
|
ax = gp_Ax3(
|
|
gp_Pnt(float(apex[0]), float(apex[1]), float(apex[2])),
|
|
gp_Dir(float(axis[0]), float(axis[1]), float(axis[2]))
|
|
)
|
|
|
|
# OCC uses reference radius at distance 1 from apex
|
|
# For half_angle θ, ref_radius = tan(θ)
|
|
ref_radius = np.tan(half_angle)
|
|
|
|
return Geom_ConicalSurface(ax, float(half_angle), float(ref_radius))
|
|
|
|
|
|
# =============================================================================
|
|
# Bounds computation for primitive surfaces
|
|
# =============================================================================
|
|
|
|
def compute_surface_bounds_from_points(surface, points, margin=0.2):
|
|
"""
|
|
Compute UV bounds for a surface based on point cloud extent.
|
|
|
|
For planes and other primitive surfaces, uses a simple approach based on
|
|
the spatial extent of the point cloud to avoid huge UV ranges.
|
|
|
|
Args:
|
|
surface: OCC Geom_Surface object (Geom_Plane, etc.)
|
|
points: Nx3 numpy array of points
|
|
margin: Fractional margin to add around bounds (default 0.2 = 20%)
|
|
|
|
Returns:
|
|
(umin, umax, vmin, vmax) tuple, or None if computation fails
|
|
"""
|
|
check_occ()
|
|
|
|
points = np.array(points, dtype=np.float64)
|
|
if len(points) == 0:
|
|
return None
|
|
|
|
# Compute bounding box extent of the point cloud
|
|
# Use this as a reasonable UV range (centered at origin)
|
|
bbox_min = points.min(axis=0)
|
|
bbox_max = points.max(axis=0)
|
|
extent = bbox_max - bbox_min
|
|
max_extent = float(np.max(extent))
|
|
|
|
if max_extent < 1e-6:
|
|
return None
|
|
|
|
# Use half-extent with margin as UV bounds (centered at 0)
|
|
half_extent = max_extent * (1.0 + margin) / 2.0
|
|
|
|
return (-half_extent, half_extent, -half_extent, half_extent)
|
|
|
|
|
|
# =============================================================================
|
|
# Mesh boundary extraction for trimmed faces
|
|
# =============================================================================
|
|
|
|
def mesh_boundary_to_wire(mesh):
|
|
"""
|
|
Extract boundary edges from trimesh and create OCC wire using polyline edges.
|
|
|
|
Args:
|
|
mesh: trimesh.Trimesh object (e.g., grid mesh from surface fitting)
|
|
|
|
Returns:
|
|
TopoDS_Wire or None if no boundary found
|
|
"""
|
|
check_occ()
|
|
|
|
try:
|
|
from trimesh.grouping import group_rows
|
|
except ImportError:
|
|
print("[occ_conversion] Warning: trimesh.grouping not available")
|
|
return None
|
|
|
|
# Find boundary edges (edges that appear only once - not shared between faces)
|
|
edges_sorted = mesh.edges_sorted
|
|
boundary_indices = group_rows(edges_sorted, require_count=1)
|
|
if len(boundary_indices) == 0:
|
|
print("[occ_conversion] Warning: No boundary edges found")
|
|
return None
|
|
|
|
boundary_edges = edges_sorted[boundary_indices]
|
|
|
|
# Connect edges into ordered loop
|
|
loop_vertices = _connect_edges_to_loop(boundary_edges)
|
|
if loop_vertices is None or len(loop_vertices) < 3:
|
|
print("[occ_conversion] Warning: Could not form closed loop from boundary edges")
|
|
return None
|
|
|
|
# Get 3D coordinates
|
|
loop_points = mesh.vertices[loop_vertices]
|
|
|
|
# Create wire from polyline edges (line segments between consecutive points)
|
|
wire_maker = BRepBuilderAPI_MakeWire()
|
|
|
|
for i in range(len(loop_points)):
|
|
p1 = loop_points[i]
|
|
p2 = loop_points[(i + 1) % len(loop_points)] # Wrap to close loop
|
|
|
|
pt1 = gp_Pnt(float(p1[0]), float(p1[1]), float(p1[2]))
|
|
pt2 = gp_Pnt(float(p2[0]), float(p2[1]), float(p2[2]))
|
|
|
|
edge_maker = BRepBuilderAPI_MakeEdge(pt1, pt2)
|
|
if edge_maker.IsDone():
|
|
wire_maker.Add(edge_maker.Edge())
|
|
|
|
if not wire_maker.IsDone():
|
|
print("[occ_conversion] Warning: Wire creation failed")
|
|
return None
|
|
|
|
print(f"[occ_conversion] Created boundary wire with {len(loop_points)} vertices")
|
|
return wire_maker.Wire()
|
|
|
|
|
|
def _connect_edges_to_loop(boundary_edges):
|
|
"""
|
|
Connect disconnected edges into a single ordered vertex loop.
|
|
|
|
Args:
|
|
boundary_edges: Nx2 array of vertex index pairs
|
|
|
|
Returns:
|
|
List of vertex indices forming a connected loop, or None if failed
|
|
"""
|
|
if len(boundary_edges) == 0:
|
|
return None
|
|
|
|
# Build adjacency: vertex -> list of (edge_idx, other_vertex)
|
|
vertex_to_edges = {}
|
|
for i, (v1, v2) in enumerate(boundary_edges):
|
|
vertex_to_edges.setdefault(v1, []).append((i, v2))
|
|
vertex_to_edges.setdefault(v2, []).append((i, v1))
|
|
|
|
# Walk the boundary starting from first edge
|
|
used = set()
|
|
start_vertex = boundary_edges[0][0]
|
|
loop = [start_vertex]
|
|
current = boundary_edges[0][1]
|
|
used.add(0)
|
|
|
|
# Walk until we return to start or run out of edges
|
|
while current != start_vertex and len(used) < len(boundary_edges):
|
|
loop.append(current)
|
|
found = False
|
|
for edge_idx, next_v in vertex_to_edges.get(current, []):
|
|
if edge_idx not in used:
|
|
used.add(edge_idx)
|
|
current = next_v
|
|
found = True
|
|
break
|
|
if not found:
|
|
break
|
|
|
|
return loop if len(loop) >= 3 else None
|
|
|
|
|
|
# =============================================================================
|
|
# B-spline to OCC conversion
|
|
# =============================================================================
|
|
|
|
def geomdl_to_occ_bspline(surf):
|
|
"""
|
|
Convert geomdl BSpline.Surface to OCC Geom_BSplineSurface.
|
|
|
|
Args:
|
|
surf: geomdl BSpline.Surface object
|
|
|
|
Returns:
|
|
Geom_BSplineSurface object
|
|
"""
|
|
check_occ()
|
|
|
|
# Get control points
|
|
ctrl_pts = np.array(surf.ctrlpts2d)
|
|
n_u, n_v = ctrl_pts.shape[:2]
|
|
|
|
# Create OCC control point array (1-indexed)
|
|
poles = TColgp_Array2OfPnt(1, n_u, 1, n_v)
|
|
for i in range(n_u):
|
|
for j in range(n_v):
|
|
pt = ctrl_pts[i, j]
|
|
poles.SetValue(i + 1, j + 1, gp_Pnt(float(pt[0]), float(pt[1]), float(pt[2])))
|
|
|
|
# Get knot vectors
|
|
kv_u = np.array(surf.knotvector_u)
|
|
kv_v = np.array(surf.knotvector_v)
|
|
|
|
# Convert to unique knots and multiplicities
|
|
def knots_to_occ(kv, degree):
|
|
"""Convert flat knot vector to unique knots and multiplicities."""
|
|
unique_knots = []
|
|
mults = []
|
|
prev = None
|
|
for k in kv:
|
|
if prev is None or k != prev:
|
|
unique_knots.append(k)
|
|
mults.append(1)
|
|
else:
|
|
mults[-1] += 1
|
|
prev = k
|
|
return unique_knots, mults
|
|
|
|
u_knots, u_mults = knots_to_occ(kv_u, surf.degree_u)
|
|
v_knots, v_mults = knots_to_occ(kv_v, surf.degree_v)
|
|
|
|
# Create OCC arrays
|
|
u_knots_arr = TColStd_Array1OfReal(1, len(u_knots))
|
|
u_mults_arr = TColStd_Array1OfInteger(1, len(u_mults))
|
|
for i, (k, m) in enumerate(zip(u_knots, u_mults)):
|
|
u_knots_arr.SetValue(i + 1, float(k))
|
|
u_mults_arr.SetValue(i + 1, int(m))
|
|
|
|
v_knots_arr = TColStd_Array1OfReal(1, len(v_knots))
|
|
v_mults_arr = TColStd_Array1OfInteger(1, len(v_mults))
|
|
for i, (k, m) in enumerate(zip(v_knots, v_mults)):
|
|
v_knots_arr.SetValue(i + 1, float(k))
|
|
v_mults_arr.SetValue(i + 1, int(m))
|
|
|
|
# Create B-spline surface
|
|
return Geom_BSplineSurface(
|
|
poles,
|
|
u_knots_arr, v_knots_arr,
|
|
u_mults_arr, v_mults_arr,
|
|
surf.degree_u, surf.degree_v
|
|
)
|
|
|
|
|
|
def nurbsdiff_to_occ_bspline(nurbs_result):
|
|
"""
|
|
Convert NURBSDiff result to OCC Geom_BSplineSurface.
|
|
|
|
Args:
|
|
nurbs_result: Dict from fit_nurbs_nurbsdiff containing control_points, degree
|
|
|
|
Returns:
|
|
Geom_BSplineSurface object
|
|
"""
|
|
check_occ()
|
|
|
|
ctrl_pts = nurbs_result["control_points"][0] # (n_u, n_v, 4) with weights
|
|
n_u, n_v = ctrl_pts.shape[:2]
|
|
degree_u = nurbs_result["degree_u"]
|
|
degree_v = nurbs_result["degree_v"]
|
|
|
|
# Create OCC control point array
|
|
poles = TColgp_Array2OfPnt(1, n_u, 1, n_v)
|
|
for i in range(n_u):
|
|
for j in range(n_v):
|
|
pt = ctrl_pts[i, j, :3]
|
|
poles.SetValue(i + 1, j + 1,
|
|
gp_Pnt(float(pt[0]), float(pt[1]), float(pt[2])))
|
|
|
|
# Create uniform knot vectors
|
|
def make_uniform_knots(n_ctrl, degree):
|
|
"""Create uniform clamped knot vector."""
|
|
n_knots = n_ctrl + degree + 1
|
|
knots = []
|
|
for i in range(n_knots):
|
|
if i < degree + 1:
|
|
knots.append(0.0)
|
|
elif i >= n_knots - degree - 1:
|
|
knots.append(1.0)
|
|
else:
|
|
knots.append((i - degree) / (n_knots - 2 * degree - 1))
|
|
return knots
|
|
|
|
kv_u = make_uniform_knots(n_u, degree_u)
|
|
kv_v = make_uniform_knots(n_v, degree_v)
|
|
|
|
# Convert to unique knots and multiplicities
|
|
def knots_to_occ(kv):
|
|
unique_knots = []
|
|
mults = []
|
|
prev = None
|
|
for k in kv:
|
|
if prev is None or abs(k - prev) > 1e-10:
|
|
unique_knots.append(k)
|
|
mults.append(1)
|
|
else:
|
|
mults[-1] += 1
|
|
prev = k
|
|
return unique_knots, mults
|
|
|
|
u_knots, u_mults = knots_to_occ(kv_u)
|
|
v_knots, v_mults = knots_to_occ(kv_v)
|
|
|
|
# Create OCC arrays
|
|
u_knots_arr = TColStd_Array1OfReal(1, len(u_knots))
|
|
u_mults_arr = TColStd_Array1OfInteger(1, len(u_mults))
|
|
for i, (k, m) in enumerate(zip(u_knots, u_mults)):
|
|
u_knots_arr.SetValue(i + 1, float(k))
|
|
u_mults_arr.SetValue(i + 1, int(m))
|
|
|
|
v_knots_arr = TColStd_Array1OfReal(1, len(v_knots))
|
|
v_mults_arr = TColStd_Array1OfInteger(1, len(v_mults))
|
|
for i, (k, m) in enumerate(zip(v_knots, v_mults)):
|
|
v_knots_arr.SetValue(i + 1, float(k))
|
|
v_mults_arr.SetValue(i + 1, int(m))
|
|
|
|
return Geom_BSplineSurface(
|
|
poles,
|
|
u_knots_arr, v_knots_arr,
|
|
u_mults_arr, v_mults_arr,
|
|
degree_u, degree_v
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# Surface params to OCC
|
|
# =============================================================================
|
|
|
|
def surface_params_to_occ(surface_result):
|
|
"""
|
|
Convert surface fitting result to OCC geometry.
|
|
|
|
Args:
|
|
surface_result: Dict with 'type' and 'params' from surface fitting
|
|
|
|
Returns:
|
|
OCC Geom_Surface object
|
|
"""
|
|
surface_type = surface_result["type"]
|
|
params = surface_result.get("params")
|
|
|
|
if surface_type == "plane":
|
|
axis, distance = params
|
|
# Get center from mesh if available
|
|
if "mesh" in surface_result:
|
|
center = surface_result["mesh"].centroid
|
|
else:
|
|
center = np.array(axis) * distance
|
|
return plane_to_occ(axis, distance, center)
|
|
|
|
elif surface_type == "sphere":
|
|
center, radius = params
|
|
return sphere_to_occ(center, radius)
|
|
|
|
elif surface_type == "cylinder":
|
|
axis, center, radius = params
|
|
return cylinder_to_occ(axis, center, radius)
|
|
|
|
elif surface_type == "cone":
|
|
apex, axis, half_angle = params
|
|
return cone_to_occ(apex, axis, half_angle)
|
|
|
|
elif surface_type == "open_spline":
|
|
# B-spline needs to be fit separately
|
|
raise ValueError("open_spline requires B-spline fitting first")
|
|
|
|
else:
|
|
raise ValueError(f"Unknown surface type: {surface_type}")
|
|
|
|
|
|
# =============================================================================
|
|
# Compound building and export
|
|
# =============================================================================
|
|
|
|
def build_compound(surfaces, bounds_list=None, boundary_wires=None):
|
|
"""
|
|
Build OCC TopoDS_Compound from list of surfaces.
|
|
|
|
Args:
|
|
surfaces: List of OCC Geom_Surface objects
|
|
bounds_list: Optional list of (umin, umax, vmin, vmax) tuples per surface.
|
|
If None or entry is None, uses default bounds based on surface type.
|
|
boundary_wires: Optional list of TopoDS_Wire objects for trimmed faces.
|
|
If provided, wire is used to create a trimmed face instead of UV bounds.
|
|
|
|
Returns:
|
|
TopoDS_Compound
|
|
"""
|
|
check_occ()
|
|
|
|
if bounds_list is None:
|
|
bounds_list = [None] * len(surfaces)
|
|
|
|
if boundary_wires is None:
|
|
boundary_wires = [None] * len(surfaces)
|
|
|
|
builder = BRep_Builder()
|
|
compound = TopoDS_Compound()
|
|
builder.MakeCompound(compound)
|
|
|
|
for i, surf in enumerate(surfaces):
|
|
bounds = bounds_list[i] if i < len(bounds_list) else None
|
|
wire = boundary_wires[i] if i < len(boundary_wires) else None
|
|
|
|
try:
|
|
face_maker = None
|
|
face = None
|
|
|
|
# Priority 1: Use boundary wire for trimmed face on surface
|
|
if wire is not None:
|
|
try:
|
|
# Create face from surface bounded by wire
|
|
face_maker = BRepBuilderAPI_MakeFace(surf, wire, True)
|
|
if face_maker.IsDone():
|
|
face = face_maker.Face()
|
|
print(f" [build_compound] Surface {i}: created trimmed face from surface + wire")
|
|
except Exception as e:
|
|
print(f" [build_compound] Surface {i}: wire-based face failed ({e}), falling back")
|
|
face = None
|
|
|
|
# Priority 2: Use UV bounds
|
|
if face is None and bounds is not None:
|
|
umin, umax, vmin, vmax = bounds
|
|
face_maker = BRepBuilderAPI_MakeFace(surf, umin, umax, vmin, vmax, 1e-6)
|
|
if face_maker.IsDone():
|
|
face = face_maker.Face()
|
|
|
|
# Priority 3: Default based on surface type
|
|
if face is None:
|
|
try:
|
|
if surf.IsKind(Geom_Plane.get_type_descriptor()):
|
|
# Use reasonable default bounds for planes (10 units)
|
|
face_maker = BRepBuilderAPI_MakeFace(surf, -5, 5, -5, 5, 1e-6)
|
|
else:
|
|
# B-splines and other bounded surfaces work with just tolerance
|
|
face_maker = BRepBuilderAPI_MakeFace(surf, 1e-6)
|
|
|
|
if face_maker.IsDone():
|
|
face = face_maker.Face()
|
|
except Exception:
|
|
# Fallback if type check fails
|
|
face_maker = BRepBuilderAPI_MakeFace(surf, 1e-6)
|
|
if face_maker.IsDone():
|
|
face = face_maker.Face()
|
|
|
|
if face is not None:
|
|
builder.Add(compound, face)
|
|
else:
|
|
print(f"Warning: Face maker failed for surface {i}")
|
|
|
|
except Exception as e:
|
|
print(f"Warning: Could not create face from surface {i}: {e}")
|
|
|
|
return compound
|
|
|
|
|
|
def export_step(compound, filename):
|
|
"""
|
|
Export OCC compound to STEP file.
|
|
|
|
Args:
|
|
compound: TopoDS_Compound or TopoDS_Shape
|
|
filename: Output STEP file path
|
|
|
|
Returns:
|
|
True if successful
|
|
"""
|
|
check_occ()
|
|
|
|
writer = STEPControl_Writer()
|
|
writer.Transfer(compound, STEPControl_AsIs)
|
|
status = writer.Write(filename)
|
|
|
|
return status == IFSelect_RetDone
|