"""STEP → USD exporter for HartOMat. Reads a STEP file via OCP/XCAF (preserving part names + embedded colors), tessellates with BRepMesh, builds a USD stage mirroring the full XCAF assembly hierarchy (intermediate Xform prims with local transforms, leaf Mesh prims with definition-space geometry), and writes a .usd file. Coordinate system: OCC is mm Z-up. USD stage is Z-up with a coordinate swap Xform on /Root/Assembly: (X_occ, Y_occ, Z_occ) → (X, -Z, Y). metersPerUnit=0.001 is set so Blender handles mm→m on import. Usage: python3 export_step_to_usd.py \\ --step_path /path/to/file.stp \\ --output_path /path/to/output.usd \\ [--linear_deflection 0.03] \\ [--angular_deflection 0.05] \\ [--color_map '{"Ring": "#4C9BE8"}'] \\ [--sharp_threshold 20.0] \\ [--cad_file_id uuid] \\ [--material_map '{"part_name": "HARTOMAT_010101_Steel-Bare", ...}'] Exit 0 on success, exit 1 on failure. Prints MANIFEST_JSON: {...} to stdout before exit. """ from __future__ import annotations import argparse import hashlib import json import math import re import sys import traceback from pathlib import Path # ── CLI ─────────────────────────────────────────────────────────────────────── def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser() p.add_argument("--step_path", required=True) p.add_argument("--output_path", required=True) p.add_argument("--linear_deflection", type=float, default=0.03) p.add_argument("--angular_deflection", type=float, default=0.05) p.add_argument("--color_map", default="{}") p.add_argument("--sharp_threshold", type=float, default=20.0) p.add_argument("--cad_file_id", default="") p.add_argument("--material_map", default="{}") p.add_argument("--tessellation_engine", choices=["occ", "gmsh"], default="occ", help="Tessellation backend: 'occ' = BRepMesh (default, fast), " "'gmsh' = Frontal-Delaunay (uniform mesh, no fan artifacts but slower)") return p.parse_args() # ── Part key generation ─────────────────────────────────────────────────────── _AF_RE = re.compile(r'_AF\d+$', re.IGNORECASE) def _generate_part_key(xcaf_path: str, source_name: str, existing_keys: set) -> str: """Deterministic slug, max 64 chars, unique within assembly.""" base = _AF_RE.sub('', source_name) if source_name else '' base = re.sub(r'([a-z])([A-Z])', r'\1_\2', base) slug = re.sub(r'[^a-z0-9]+', '_', base.lower()).strip('_') if not slug: slug = f"part_{hashlib.sha256(xcaf_path.encode()).hexdigest()[:8]}" slug = slug[:50] # USD prim names cannot start with a digit if slug and slug[0].isdigit(): slug = f"p_{slug}" key = slug n = 2 while key in existing_keys: key = f"{slug}_{n}" n += 1 existing_keys.add(key) return key # ── Color helpers ───────────────────────────────────────────────────────────── PALETTE_HEX = [ "#4C9BE8", "#E85B4C", "#4CBE72", "#E8A84C", "#A04CE8", "#4CD4E8", "#E84CA8", "#7EC850", "#E86B30", "#5088C8", ] def _occ_color_to_hex(occ_color) -> str: r = int(occ_color.Red() * 255) g = int(occ_color.Green() * 255) b = int(occ_color.Blue() * 255) return f"#{r:02X}{g:02X}{b:02X}" def _hex_to_occ_color(hex_color: str): from OCP.Quantity import Quantity_Color, Quantity_TOC_RGB h = hex_color.lstrip("#") if len(h) < 6: return Quantity_Color(0.7, 0.7, 0.7, Quantity_TOC_RGB) return Quantity_Color( int(h[0:2], 16) / 255.0, int(h[2:4], 16) / 255.0, int(h[4:6], 16) / 255.0, Quantity_TOC_RGB, ) def _hex_to_rgb01(hex_color: str) -> tuple: h = hex_color.lstrip('#') if len(h) < 6: return (0.7, 0.7, 0.7) return (int(h[0:2], 16) / 255.0, int(h[2:4], 16) / 255.0, int(h[4:6], 16) / 255.0) def _get_shape_color(color_tool, shape) -> str | None: """Return hex color for an OCC shape (surface color preferred).""" from OCP.Quantity import Quantity_Color try: from OCP.XCAFDoc import XCAFDoc_ColorSurf as _SURF from OCP.XCAFDoc import XCAFDoc_ColorGen as _GEN except ImportError: _SURF = 1 _GEN = 0 occ_color = Quantity_Color() if color_tool.GetColor(shape, _SURF, occ_color): return _occ_color_to_hex(occ_color) if color_tool.GetColor(shape, _GEN, occ_color): return _occ_color_to_hex(occ_color) return None # ── XCAF color application ──────────────────────────────────────────────────── def _apply_color_map(shape_tool, color_tool, free_labels, color_map: dict) -> None: from OCP.TDF import TDF_LabelSequence from OCP.TDataStd import TDataStd_Name from OCP.XCAFDoc import XCAFDoc_ShapeTool try: from OCP.XCAFDoc import XCAFDoc_ColorSurf as _SURF except ImportError: _SURF = 1 def _visit(label) -> None: name_attr = TDataStd_Name() name = "" if label.FindAttribute(TDataStd_Name.GetID_s(), name_attr): name = name_attr.Get().ToExtString() if name: for part_name, hex_color in color_map.items(): if part_name.lower() in name.lower() or name.lower() in part_name.lower(): color_tool.SetColor(label, _hex_to_occ_color(hex_color), _SURF) break components = TDF_LabelSequence() XCAFDoc_ShapeTool.GetComponents_s(label, components) for i in range(1, components.Length() + 1): _visit(components.Value(i)) for i in range(1, free_labels.Length() + 1): _visit(free_labels.Value(i)) def _apply_palette_colors(shape_tool, color_tool, free_labels) -> None: from OCP.TDF import TDF_LabelSequence from OCP.XCAFDoc import XCAFDoc_ShapeTool try: from OCP.XCAFDoc import XCAFDoc_ColorSurf as _SURF except ImportError: _SURF = 1 leaves: list = [] def _collect(label) -> None: components = TDF_LabelSequence() XCAFDoc_ShapeTool.GetComponents_s(label, components) if components.Length() == 0: leaves.append(label) else: for i in range(1, components.Length() + 1): _collect(components.Value(i)) for i in range(1, free_labels.Length() + 1): _collect(free_labels.Value(i)) for idx, label in enumerate(leaves): color_tool.SetColor(label, _hex_to_occ_color(PALETTE_HEX[idx % len(PALETTE_HEX)]), _SURF) # ── Sharp edge extraction (inlined from export_step_to_gltf.py) ────────────── def _extract_sharp_edge_pairs(shape, sharp_threshold_deg: float = 20.0) -> list: """Extract sharp B-rep edges as dense curve-sample segment pairs (mm, Z-up). Ported from export_step_to_gltf.py to avoid importing that module (its top-level code runs main() on import). """ from OCP.TopTools import TopTools_IndexedDataMapOfShapeListOfShape from OCP.TopExp import TopExp as _TopExp from OCP.TopAbs import TopAbs_EDGE, TopAbs_FACE, TopAbs_FORWARD from OCP.TopoDS import TopoDS as _TopoDS from OCP.BRepAdaptor import BRepAdaptor_Surface, BRepAdaptor_Curve2d, BRepAdaptor_Curve from OCP.BRepLProp import BRepLProp_SLProps from OCP.GCPnts import GCPnts_UniformAbscissa edge_face_map = TopTools_IndexedDataMapOfShapeListOfShape() _TopExp.MapShapesAndAncestors_s(shape, TopAbs_EDGE, TopAbs_FACE, edge_face_map) sharp_pairs: list = [] n_checked = 0 n_sharp = 0 SAMPLE_STEP_MM = 0.3 for i in range(1, edge_face_map.Extent() + 1): edge_shape = edge_face_map.FindKey(i) faces = edge_face_map.FindFromIndex(i) n_checked += 1 if faces.Size() < 2: continue face_shapes = list(faces) if len(face_shapes) < 2: continue try: edge = _TopoDS.Edge_s(edge_shape) face1 = _TopoDS.Face_s(face_shapes[0]) face2 = _TopoDS.Face_s(face_shapes[1]) c2d_1 = BRepAdaptor_Curve2d(edge, face1) uv1 = c2d_1.Value((c2d_1.FirstParameter() + c2d_1.LastParameter()) / 2.0) surf1 = BRepAdaptor_Surface(face1) props1 = BRepLProp_SLProps(surf1, uv1.X(), uv1.Y(), 1, 1e-6) if not props1.IsNormalDefined(): continue n1 = props1.Normal() if face1.Orientation() != TopAbs_FORWARD: n1.Reverse() c2d_2 = BRepAdaptor_Curve2d(edge, face2) uv2 = c2d_2.Value((c2d_2.FirstParameter() + c2d_2.LastParameter()) / 2.0) surf2 = BRepAdaptor_Surface(face2) props2 = BRepLProp_SLProps(surf2, uv2.X(), uv2.Y(), 1, 1e-6) if not props2.IsNormalDefined(): continue n2 = props2.Normal() if face2.Orientation() != TopAbs_FORWARD: n2.Reverse() cos_angle = max(-1.0, min(1.0, n1.Dot(n2))) angle_deg = math.degrees(math.acos(cos_angle)) if angle_deg > 90.0: angle_deg = 180.0 - angle_deg if angle_deg <= sharp_threshold_deg: continue n_sharp += 1 pts: list = [] try: curve3d = BRepAdaptor_Curve(edge) f_param = curve3d.FirstParameter() l_param = curve3d.LastParameter() if math.isfinite(f_param) and math.isfinite(l_param): sampler = GCPnts_UniformAbscissa() sampler.Initialize(curve3d, SAMPLE_STEP_MM, 1e-6) if sampler.IsDone() and sampler.NbPoints() >= 2: for j in range(1, sampler.NbPoints() + 1): p = curve3d.Value(sampler.Parameter(j)) pts.append([round(p.X(), 4), round(p.Y(), 4), round(p.Z(), 4)]) except Exception: pts = [] if len(pts) < 2: continue for k in range(len(pts) - 1): sharp_pairs.append([pts[k], pts[k + 1]]) except Exception: continue print( f"Sharp edge extraction: {n_checked} edges checked, " f"{n_sharp} sharp (>{sharp_threshold_deg:.0f}°), " f"{len(sharp_pairs)} segment pairs total" ) return sharp_pairs def _extract_seam_edge_pairs(shape) -> list: """Extract seam edges (periodic-surface boundary edges) as segment pairs (mm, Z-up). Seam edges are detected via BRep_Tool.IsClosed_s(edge) — edges that are topologically closed (start == end vertex). This includes the UV seams of periodic surfaces (cylinders, cones, spheres) but also full circles on flat faces and bore rims. TODO: Use ShapeAnalysis_Edge().IsSeam(edge, face) to restrict to true UV seams when UV-unwrapped texture mapping is needed (future phase). """ from OCP.BRep import BRep_Tool from OCP.TopExp import TopExp_Explorer from OCP.TopAbs import TopAbs_EDGE from OCP.BRepAdaptor import BRepAdaptor_Curve from OCP.GCPnts import GCPnts_UniformAbscissa seam_pairs: list = [] n_seam = 0 exp = TopExp_Explorer(shape, TopAbs_EDGE) while exp.More(): edge = exp.Current() exp.Next() if not BRep_Tool.IsClosed_s(edge): continue try: curve = BRepAdaptor_Curve(edge) # Use arc-length step (0.3 mm) matching the sharp edge sampler, # so segments are short enough for _world_to_index_pairs (tol=0.5 mm). sampler = GCPnts_UniformAbscissa() sampler.Initialize(curve, 0.3, 1e-6) if not sampler.IsDone() or sampler.NbPoints() < 2: continue pts = [] for i in range(1, sampler.NbPoints() + 1): p = curve.Value(sampler.Parameter(i)) pts.append([p.X(), p.Y(), p.Z()]) for k in range(len(pts) - 1): seam_pairs.append([pts[k], pts[k + 1]]) n_seam += 1 except Exception: continue print(f"Seam edge extraction: {n_seam} seam edges, {len(seam_pairs)} segment pairs total") return seam_pairs # ── XCAF traversal + hierarchical USD authoring ────────────────────────────── def _get_label_name(label) -> str: """Extract the source name string from a TDF_Label.""" from OCP.TDataStd import TDataStd_Name name_attr = TDataStd_Name() if label.FindAttribute(TDataStd_Name.GetID_s(), name_attr): return name_attr.Get().ToExtString() return "" def _occ_trsf_to_usd_matrix(trsf): """Convert an OCC gp_Trsf (column-vector) to a USD Gf.Matrix4d (row-vector). OCC uses p' = R·p + t (column-vector convention). USD uses p' = p·M (row-vector convention). So M = [R^T | 0; t^T | 1]. """ from pxr import Gf return Gf.Matrix4d( trsf.Value(1, 1), trsf.Value(2, 1), trsf.Value(3, 1), 0, trsf.Value(1, 2), trsf.Value(2, 2), trsf.Value(3, 2), 0, trsf.Value(1, 3), trsf.Value(2, 3), trsf.Value(3, 3), 0, trsf.Value(1, 4), trsf.Value(2, 4), trsf.Value(3, 4), 1, ) def _author_xcaf_to_usd( stage, shape_tool, color_tool, label, usd_parent_path, xcaf_path_prefix, existing_keys, mat_map_lower, color_map, args, manifest_parts, counters, prim_names_at_level, depth=0, ): """Recursively author USD prims mirroring the XCAF hierarchy. Assembly nodes → UsdGeom.Xform with local transform from component Location. Leaf shapes → Xform + Mesh with definition-space geometry. Sharp/seam edges are extracted per-part from the definition shape. The local transform for each node comes from GetShape_s(label).Location(), which returns ONLY this label's own placement (not composed with parents). USD scene graph composition handles the full parent-to-leaf chain. """ from OCP.TDF import TDF_LabelSequence, TDF_Label from OCP.XCAFDoc import XCAFDoc_ShapeTool from OCP.TopLoc import TopLoc_Location from pxr import UsdGeom, UsdShade, Sdf, Vt, Gf source_name = _get_label_name(label) xcaf_path = (f"{xcaf_path_prefix}/{source_name}" if source_name else f"{xcaf_path_prefix}/unnamed_{depth}") # Get local transform from this label's shape Location. # GetShape_s(label) returns the shape with ONLY this label's own Location # (not composed with parent locations). label_shape = shape_tool.GetShape_s(label) if label_shape.IsNull(): return local_loc = label_shape.Location() has_local_trsf = not local_loc.IsIdentity() # Follow reference to definition label actual_label = label if XCAFDoc_ShapeTool.IsReference_s(label): ref_label = TDF_Label() if XCAFDoc_ShapeTool.GetReferredShape_s(label, ref_label): actual_label = ref_label # Check for sub-components on the definition components = TDF_LabelSequence() XCAFDoc_ShapeTool.GetComponents_s(actual_label, components) if components.Length() > 0: # ── ASSEMBLY NODE ────────────────────────────────────────────── raw_name = _prim_name(source_name or f"asm_{depth}") unique_name = raw_name n = 2 while unique_name in prim_names_at_level: unique_name = f"{raw_name}_{n}" n += 1 prim_names_at_level.add(unique_name) xform_path = f"{usd_parent_path}/{unique_name}" xform = UsdGeom.Xform.Define(stage, xform_path) if has_local_trsf: xform.AddTransformOp().Set( _occ_trsf_to_usd_matrix(local_loc.Transformation())) prim = xform.GetPrim() prim.SetCustomDataByKey("hartomat:sourceName", source_name) prim.SetCustomDataByKey("hartomat:sourceAssemblyPath", xcaf_path) print(f" {' ' * depth}[asm] {source_name} → {xform_path}" f"{' (transform)' if has_local_trsf else ''}") child_names: set = set() for i in range(1, components.Length() + 1): _author_xcaf_to_usd( stage, shape_tool, color_tool, components.Value(i), xform_path, xcaf_path, existing_keys, mat_map_lower, color_map, args, manifest_parts, counters, child_names, depth + 1, ) else: # ── LEAF SHAPE ───────────────────────────────────────────────── # Get definition shape without instance location def_shape = shape_tool.GetShape_s(actual_label) if def_shape.IsNull(): return # Strip any residual location so _extract_mesh yields definition-space # vertices (face_loc only, no instance placement). bare_def = def_shape.Located(TopLoc_Location()) part_key = _generate_part_key(xcaf_path, source_name, existing_keys) hex_color = _get_shape_color(color_tool, label_shape) if not hex_color: hex_color = _get_shape_color(color_tool, def_shape) # color_map override (substring match) if source_name: for map_name, map_hex in color_map.items(): if (map_name.lower() in source_name.lower() or source_name.lower() in map_name.lower()): hex_color = map_hex break if not hex_color: hex_color = PALETTE_HEX[counters['n_parts'] % len(PALETTE_HEX)] # Extract mesh from definition shape (face_loc only, no instance placement) vertices, triangles, per_vertex_normals = _extract_mesh(bare_def) if not vertices or not triangles: counters['n_empty'] += 1 return # Ensure unique prim name at this level raw_name = _prim_name(part_key) unique_name = raw_name n = 2 while unique_name in prim_names_at_level: unique_name = f"{raw_name}_{n}" n += 1 prim_names_at_level.add(unique_name) part_path = f"{usd_parent_path}/{unique_name}" # Name the Mesh prim after part_key so Blender imports it with the # part name directly (Blender collapses single-child Xform+Mesh). mesh_path = f"{part_path}/{part_key}" # ── Xform prim with local transform ──────────────────────── xform = UsdGeom.Xform.Define(stage, part_path) if has_local_trsf: xform.AddTransformOp().Set( _occ_trsf_to_usd_matrix(local_loc.Transformation())) prim = xform.GetPrim() prim.SetCustomDataByKey("hartomat:partKey", part_key) prim.SetCustomDataByKey("hartomat:sourceName", source_name) prim.SetCustomDataByKey("hartomat:sourceAssemblyPath", xcaf_path) prim.SetCustomDataByKey("hartomat:sourceColor", hex_color) prim.SetCustomDataByKey("hartomat:tessellation:linearDeflectionMm", args.linear_deflection) prim.SetCustomDataByKey("hartomat:tessellation:angularDeflectionRad", args.angular_deflection) if args.cad_file_id: prim.SetCustomDataByKey("hartomat:cadFileId", args.cad_file_id) # ── UsdGeomMesh ──────────────────────────────────────────── mesh = UsdGeom.Mesh.Define(stage, mesh_path) mesh.CreateSubdivisionSchemeAttr(UsdGeom.Tokens.none) # Vertices in OCC definition space (mm, Z-up). # The /Root/Assembly Xform carries the OCC→USD coordinate swap # so no per-vertex (X,-Z,Y) transformation is needed here. mesh.CreatePointsAttr(Vt.Vec3fArray([ Gf.Vec3f(x, y, z) for (x, y, z) in vertices ])) mesh.CreateFaceVertexCountsAttr(Vt.IntArray([3] * len(triangles))) mesh.CreateFaceVertexIndicesAttr( Vt.IntArray([idx for tri in triangles for idx in tri]) ) r, g, b = _hex_to_rgb01(hex_color) mesh.CreateDisplayColorAttr(Vt.Vec3fArray([Gf.Vec3f(r, g, b)])) # ── Per-vertex normals ───────────────────────────────────── # Blender's USD importer reads normals and creates custom split # normals. Using 'vertex' interpolation (1 normal per vertex) is # ~6× smaller than 'faceVarying' (1 per triangle corner) and # produces identical smooth shading for our tessellated meshes. if per_vertex_normals and len(per_vertex_normals) == len(vertices): mesh.CreateNormalsAttr(Vt.Vec3fArray([ Gf.Vec3f(nx, ny, nz) for (nx, ny, nz) in per_vertex_normals ])) mesh.SetNormalsInterpolation(UsdGeom.Tokens.vertex) # ── Material metadata on mesh prim (customData) ─────────── mesh_prim = mesh.GetPrim() mesh_prim.SetCustomDataByKey("hartomat:partKey", part_key) mesh_prim.SetCustomDataByKey("hartomat:sourceName", source_name) canonical_mat = _lookup_material(source_name, part_key, mat_map_lower) if canonical_mat: mesh_prim.SetCustomDataByKey( "hartomat:canonicalMaterialName", canonical_mat) primvars_api = UsdGeom.PrimvarsAPI(mesh) # ── Per-part sharp + seam edge primvars (definition space) ─ try: sharp_pairs = _extract_sharp_edge_pairs(bare_def, args.sharp_threshold) if sharp_pairs: idx_pairs = _world_to_index_pairs(vertices, sharp_pairs) if idx_pairs: pv = primvars_api.CreatePrimvar( "hartomat:sharpEdgeVertexPairs", Sdf.ValueTypeNames.Int2Array, UsdGeom.Tokens.constant, ) pv.Set(Vt.Vec2iArray([Gf.Vec2i(a, b) for a, b in idx_pairs])) except Exception as exc: print(f"WARNING: sharp edge extraction for {part_key}: {exc}") try: seam_pairs = _extract_seam_edge_pairs(bare_def) if seam_pairs: seam_idx_pairs = _world_to_index_pairs(vertices, seam_pairs) if seam_idx_pairs: pv_seam = primvars_api.CreatePrimvar( "hartomat:seamEdgeVertexPairs", Sdf.ValueTypeNames.Int2Array, UsdGeom.Tokens.constant, ) pv_seam.Set(Vt.Vec2iArray([ Gf.Vec2i(a, b) for a, b in seam_idx_pairs])) except Exception as exc: print(f"WARNING: seam edge extraction for {part_key}: {exc}") # ── Material binding ────────────────────────────────────── if canonical_mat: mat_prim_name = _prim_name(canonical_mat) else: mat_prim_name = (_prim_name(source_name) if source_name else f"mat_{part_key}") mat_usd_path = f"/Root/Looks/{mat_prim_name}" if not stage.GetPrimAtPath(mat_usd_path): UsdShade.Material.Define(stage, mat_usd_path) UsdShade.MaterialBindingAPI(mesh.GetPrim()).Bind( UsdShade.Material(stage.GetPrimAtPath(mat_usd_path)) ) manifest_parts.append({ "part_key": part_key, "source_name": source_name, "prim_path": part_path, "canonical_material": canonical_mat, }) counters['n_parts'] += 1 # ── Mesh geometry extraction ────────────────────────────────────────────────── def _extract_mesh(shape) -> tuple[list, list, list]: """Return (vertices, triangles, normals) from a tessellated OCC shape. Vertices are in OCC space (mm, Z-up). Triangles are 0-based index triples. Normals are per-vertex (same count/order as vertices). Normal sources (in priority order): 1. B-Rep surface normals — evaluated from the mathematical surface at each vertex UV parameter via BRepLProp_SLProps. Gives perfectly smooth normals regardless of triangle topology (same approach as Stepper/CAD-Exchanger). 2. Angle-weighted vertex normals from triangle cross products — fallback when UV nodes are unavailable (e.g. after GMSH tessellation). Transform strategy: strip the shape's own Location before exploring faces so that face_loc from BRep_Tool.Triangulation_s is always relative to the shape's DEFINITION space (not contaminated by instance placement). Then uniformly apply the shape's Location to every vertex. This avoids both double-transform (when face_loc already includes placement) and missing- transform (when face_loc is identity but shape has placement). """ from OCP.TopExp import TopExp_Explorer from OCP.TopAbs import TopAbs_FACE, TopAbs_REVERSED from OCP.TopoDS import TopoDS from OCP.BRep import BRep_Tool from OCP.TopLoc import TopLoc_Location from OCP.BRepAdaptor import BRepAdaptor_Surface from OCP.BRepLProp import BRepLProp_SLProps vertices: list = [] normals: list = [] triangles: list = [] v_offset = 0 faces_without_normals = False shape_trsf = shape.Location().Transformation() shape_has_loc = not shape.Location().IsIdentity() # Strip instance placement so face exploration yields definition-space locs bare = shape.Located(TopLoc_Location()) exp = TopExp_Explorer(bare, TopAbs_FACE) while exp.More(): face = TopoDS.Face_s(exp.Current()) face_loc = TopLoc_Location() tri = BRep_Tool.Triangulation_s(face, face_loc) if tri is not None and tri.NbNodes() > 0: reversed_face = (face.Orientation() == TopAbs_REVERSED) face_has_loc = not face_loc.IsIdentity() # Evaluate surface normals from B-Rep geometry (analytically exact). # This gives smooth shading across face boundaries because the normals # come from the mathematical surface, not from triangle approximation. has_uv = tri.HasUVNodes() surf = None if has_uv: try: surf = BRepAdaptor_Surface(face) except Exception: surf = None for i in range(1, tri.NbNodes() + 1): node = tri.Node(i) if face_has_loc: node = node.Transformed(face_loc.Transformation()) if shape_has_loc: node = node.Transformed(shape_trsf) vertices.append((node.X(), node.Y(), node.Z())) nrm_set = False if surf is not None: try: uv = tri.UVNode(i) props = BRepLProp_SLProps(surf, uv.X(), uv.Y(), 1, 1e-6) if props.IsNormalDefined(): nrm = props.Normal() if face_has_loc: nrm = nrm.Transformed(face_loc.Transformation()) if shape_has_loc: nrm = nrm.Transformed(shape_trsf) if reversed_face: normals.append((-nrm.X(), -nrm.Y(), -nrm.Z())) else: normals.append((nrm.X(), nrm.Y(), nrm.Z())) nrm_set = True except Exception: pass if not nrm_set: faces_without_normals = True normals.append(None) # type: ignore[arg-type] for i in range(1, tri.NbTriangles() + 1): n1, n2, n3 = tri.Triangle(i).Get() v0 = n1 - 1 + v_offset v1 = n2 - 1 + v_offset v2 = n3 - 1 + v_offset triangles.append((v0, v2, v1) if reversed_face else (v0, v1, v2)) v_offset += tri.NbNodes() exp.Next() # Fallback: compute normals from triangle geometry for faces without UV data if faces_without_normals and vertices and triangles: _compute_vertex_normals(vertices, triangles, normals) return vertices, triangles, normals def _compute_vertex_normals(vertices: list, triangles: list, normals: list) -> None: """Fill in None entries in normals with angle-weighted vertex normals. For each triangle, computes the face normal (cross product) and adds it to each vertex's accumulator weighted by the interior angle at that vertex. This produces smooth normals identical to what Blender's "Shade Smooth" does. """ # Accumulate weighted normals for vertices that need them accum = {} # vertex_index → [nx, ny, nz] for tri in triangles: v0, v1, v2 = tri # Only process if any vertex in this triangle needs normals needs = any(normals[vi] is None for vi in (v0, v1, v2)) if not needs: continue p0 = vertices[v0] p1 = vertices[v1] p2 = vertices[v2] # Edge vectors e01 = (p1[0]-p0[0], p1[1]-p0[1], p1[2]-p0[2]) e02 = (p2[0]-p0[0], p2[1]-p0[1], p2[2]-p0[2]) e12 = (p2[0]-p1[0], p2[1]-p1[1], p2[2]-p1[2]) # Face normal (unnormalized cross product) fn = ( e01[1]*e02[2] - e01[2]*e02[1], e01[2]*e02[0] - e01[0]*e02[2], e01[0]*e02[1] - e01[1]*e02[0], ) fn_len = (fn[0]**2 + fn[1]**2 + fn[2]**2) ** 0.5 if fn_len < 1e-12: continue # Angle at each vertex (for weighting) def _dot(a, b): return a[0]*b[0] + a[1]*b[1] + a[2]*b[2] def _len(a): return (a[0]**2 + a[1]**2 + a[2]**2) ** 0.5 def _neg(a): return (-a[0], -a[1], -a[2]) edges_at_vert = [ (v0, e01, e02), (v1, _neg(e01), e12), (v2, _neg(e02), _neg(e12)), ] for vi, ea, eb in edges_at_vert: if normals[vi] is not None: continue la, lb = _len(ea), _len(eb) if la < 1e-12 or lb < 1e-12: continue cos_a = max(-1.0, min(1.0, _dot(ea, eb) / (la * lb))) angle = math.acos(cos_a) if vi not in accum: accum[vi] = [0.0, 0.0, 0.0] accum[vi][0] += fn[0] * angle accum[vi][1] += fn[1] * angle accum[vi][2] += fn[2] * angle # Normalize and write back for vi, acc in accum.items(): length = (acc[0]**2 + acc[1]**2 + acc[2]**2) ** 0.5 if length > 1e-12: normals[vi] = (acc[0]/length, acc[1]/length, acc[2]/length) else: normals[vi] = (0.0, 0.0, 1.0) # Fill any remaining None entries (shouldn't happen, but safety) for i in range(len(normals)): if normals[i] is None: normals[i] = (0.0, 0.0, 1.0) # ── GMSH Frontal-Delaunay tessellation ─────────────────────────────────────── def _tessellate_with_gmsh(shape, linear_deflection: float, angular_deflection: float) -> None: """Tessellate an OCC TopoDS_Shape using GMSH Frontal-Delaunay mesher. Writes the resulting Poly_Triangulation back to each TopoDS_Face via BRep_Builder.UpdateFace(), so _extract_mesh() picks up the uniform topology from BRep_Tool.Triangulation_s. GMSH surface tags correspond 1:1 to faces in TopExp_Explorer(FACE) order after importShapes() from a .brep file. Falls back to BRepMesh for any face that GMSH cannot mesh. """ import tempfile from OCP.BRep import BRep_Builder from OCP.BRepTools import BRepTools from OCP.BRepMesh import BRepMesh_IncrementalMesh from OCP.Poly import Poly_Triangulation, Poly_Array1OfTriangle, Poly_Triangle from OCP.TColgp import TColgp_Array1OfPnt from OCP.TopExp import TopExp_Explorer from OCP.TopAbs import TopAbs_FACE from OCP.TopoDS import TopoDS as _TopoDS from OCP.gp import gp_Pnt import gmsh with tempfile.NamedTemporaryFile(suffix=".brep", delete=False) as tmp: brep_path = tmp.name n_faces_gmsh = 0 n_faces_fallback = 0 n_triangles_total = 0 try: BRepTools.Write_s(shape, brep_path) import os as _os n_threads = min(_os.cpu_count() or 1, 16) gmsh.initialize() gmsh.option.setNumber("General.Terminal", 0) gmsh.option.setNumber("General.NumThreads", n_threads) gmsh.option.setNumber("Mesh.MaxNumThreads1D", n_threads) gmsh.option.setNumber("Mesh.MaxNumThreads2D", n_threads) gmsh.option.setNumber("Mesh.Algorithm", 6) # Frontal-Delaunay 2D gmsh.option.setNumber("Mesh.RecombineAll", 0) gmsh.option.setNumber("Mesh.CharacteristicLengthMin", linear_deflection * 0.5) gmsh.option.setNumber("Mesh.CharacteristicLengthMax", linear_deflection * 50.0) min_circle_pts = min(12, max(6, int(math.ceil( 2.0 * math.pi / max(angular_deflection, 0.01))))) gmsh.option.setNumber("Mesh.MinimumCirclePoints", min_circle_pts) gmsh.option.setNumber("Mesh.MinimumCurvePoints", 3) gmsh.option.setNumber("General.Verbosity", 1) gmsh.model.add("shape") gmsh.model.occ.importShapes(brep_path) gmsh.model.occ.synchronize() gmsh.model.mesh.generate(2) surface_tags = [tag for (_, tag) in gmsh.model.getEntities(2)] surface_mesh: dict[int, tuple] = {} for stag in surface_tags: try: node_tags, coords, _ = gmsh.model.mesh.getNodes( dim=2, tag=stag, includeBoundary=True) if len(node_tags) == 0: continue node_map = {} pts_xyz = [] for i, ntag in enumerate(node_tags): x, y, z = coords[3*i], coords[3*i+1], coords[3*i+2] node_map[ntag] = i + 1 pts_xyz.append((x, y, z)) elem_types, elem_tags, elem_node_tags = gmsh.model.mesh.getElements( dim=2, tag=stag) tris = [] for etype, etags, entags in zip(elem_types, elem_tags, elem_node_tags): if etype != 2: continue n_elems = len(etags) for k in range(n_elems): a = node_map.get(entags[3*k]) b = node_map.get(entags[3*k+1]) c = node_map.get(entags[3*k+2]) if a and b and c: tris.append((a, b, c)) if pts_xyz and tris: surface_mesh[stag] = (pts_xyz, tris) except Exception: continue except Exception as _gmsh_err: print(f"WARNING: GMSH failed ({_gmsh_err}), falling back to BRepMesh", file=sys.stderr) BRepMesh_IncrementalMesh(shape, linear_deflection, False, angular_deflection, True) return finally: try: gmsh.finalize() except Exception: pass try: Path(brep_path).unlink(missing_ok=True) except Exception: pass builder = BRep_Builder() explorer = TopExp_Explorer(shape, TopAbs_FACE) face_index = 0 while explorer.More(): face = _TopoDS.Face_s(explorer.Current()) if face_index < len(surface_tags): stag = surface_tags[face_index] mesh_data = surface_mesh.get(stag) if mesh_data: pts_xyz, tris = mesh_data n_nodes = len(pts_xyz) n_tris = len(tris) try: arr_pts = TColgp_Array1OfPnt(1, n_nodes) for idx, (x, y, z) in enumerate(pts_xyz, 1): arr_pts.SetValue(idx, gp_Pnt(x, y, z)) arr_tris = Poly_Array1OfTriangle(1, n_tris) for idx, (a, b, c) in enumerate(tris, 1): arr_tris.SetValue(idx, Poly_Triangle(a, b, c)) tri = Poly_Triangulation(arr_pts, arr_tris) tri.ComputeNormals() # smooth normals from triangle averaging builder.UpdateFace(face, tri) n_faces_gmsh += 1 n_triangles_total += n_tris except Exception: BRepMesh_IncrementalMesh( face, linear_deflection, False, angular_deflection, False) n_faces_fallback += 1 else: BRepMesh_IncrementalMesh( face, linear_deflection, False, angular_deflection, False) n_faces_fallback += 1 else: BRepMesh_IncrementalMesh( face, linear_deflection, False, angular_deflection, False) n_faces_fallback += 1 face_index += 1 explorer.Next() print( f"GMSH tessellation: {n_faces_gmsh} faces meshed, " f"{n_faces_fallback} BRepMesh fallback, " f"{n_triangles_total} triangles total" f" (threads={n_threads})" ) # ── Index-space sharp edge mapping ──────────────────────────────────────────── def _world_to_index_pairs(vertices: list, world_pairs: list, tol: float = 0.5) -> list: """Map world-space (mm, Z-up) segment pairs → local vertex index pairs.""" def _k(x, y, z): return (round(x / tol) * tol, round(y / tol) * tol, round(z / tol) * tol) coord_map: dict = {} for idx, (x, y, z) in enumerate(vertices): k = _k(x, y, z) if k not in coord_map: coord_map[k] = idx result = [] for p0, p1 in world_pairs: i0 = coord_map.get(_k(p0[0], p0[1], p0[2])) i1 = coord_map.get(_k(p1[0], p1[1], p1[2])) if i0 is not None and i1 is not None and i0 != i1: result.append((i0, i1)) return result # ── USD prim name sanitizer ─────────────────────────────────────────────────── def _prim_name(name: str) -> str: safe = re.sub(r'[^A-Za-z0-9_]', '_', name) if safe and safe[0].isdigit(): safe = f"_{safe}" return safe or "unnamed" # ── Material map lookup (mirrors _blender_materials.build_mat_map_lower) ───── def _build_mat_map_lower(material_map: dict) -> dict: """Build a lowercased material_map with AF-stripped and slug variants. Same normalization as _blender_materials.build_mat_map_lower() so that source_name → canonical material name lookup works consistently. """ mat_map_lower: dict = {} for k, v in material_map.items(): kl = k.lower().strip() mat_map_lower[kl] = v # Slug variant: replace non-alphanumeric with '_' (same as _generate_part_key) slug_key = re.sub(r'[^a-z0-9]+', '_', kl).strip('_') if slug_key and slug_key != kl: mat_map_lower.setdefault(slug_key, v) # Strip OCC assembly-frame suffixes: _AF0, _AF0_1, _AF0_1_AF0, etc. stripped = re.sub(r'(_af\d+(_\d+)?)+$', '', kl) if stripped != kl: mat_map_lower.setdefault(stripped, v) slug_stripped = re.sub(r'[^a-z0-9]+', '_', stripped).strip('_') if slug_stripped and slug_stripped != stripped: mat_map_lower.setdefault(slug_stripped, v) return mat_map_lower def _lookup_material(source_name: str, part_key: str, mat_map_lower: dict) -> str | None: """Look up canonical material name for a part, trying multiple key variants.""" if not mat_map_lower: return None # Try source_name (lowered) sn = source_name.lower().strip() if sn in mat_map_lower: return mat_map_lower[sn] # Try AF-stripped source_name stripped = re.sub(r'(_af\d+(_\d+)?)+$', '', sn, flags=re.IGNORECASE) if stripped != sn and stripped in mat_map_lower: return mat_map_lower[stripped] # Try slug of source_name (matches part_key generation logic) slug = re.sub(r'[^a-z0-9]+', '_', sn).strip('_') if slug and slug in mat_map_lower: return mat_map_lower[slug] # Try part_key directly pk = part_key.lower().strip() if pk in mat_map_lower: return mat_map_lower[pk] # Prefix fallback: longest key that starts with or is started by part_key for key in sorted(mat_map_lower.keys(), key=len, reverse=True): if len(key) >= 5 and len(pk) >= 5 and (pk.startswith(key) or key.startswith(pk)): return mat_map_lower[key] return None # ── Main ────────────────────────────────────────────────────────────────────── def main() -> None: args = parse_args() color_map: dict = json.loads(args.color_map) raw_material_map: dict = json.loads(args.material_map) mat_map_lower = _build_mat_map_lower(raw_material_map) if raw_material_map else {} if mat_map_lower: print(f"Material map: {len(raw_material_map)} entries ({len(mat_map_lower)} with variants)") step_path = Path(args.step_path) output_path = Path(args.output_path) if not step_path.exists(): print(f"ERROR: STEP file not found: {step_path}", file=sys.stderr) sys.exit(1) output_path.parent.mkdir(parents=True, exist_ok=True) # ── OCC / XCAF imports ──────────────────────────────────────────────────── from OCP.STEPCAFControl import STEPCAFControl_Reader from OCP.TDocStd import TDocStd_Document from OCP.XCAFApp import XCAFApp_Application from OCP.XCAFDoc import XCAFDoc_DocumentTool from OCP.TCollection import TCollection_ExtendedString from OCP.TDF import TDF_LabelSequence from OCP.BRepMesh import BRepMesh_IncrementalMesh from OCP.IFSelect import IFSelect_RetDone # ── pxr imports ─────────────────────────────────────────────────────────── from pxr import Usd, UsdGeom, UsdShade, Sdf, Vt, Gf # ── Read STEP ───────────────────────────────────────────────────────────── app = XCAFApp_Application.GetApplication_s() doc = TDocStd_Document(TCollection_ExtendedString("MDTV-CAF")) app.InitDocument(doc) reader = STEPCAFControl_Reader() reader.SetNameMode(True) reader.SetColorMode(True) reader.SetLayerMode(True) status = reader.ReadFile(str(step_path)) if status != IFSelect_RetDone: print(f"ERROR: STEPCAFControl_Reader failed (status={status})", file=sys.stderr) sys.exit(1) reader.Transfer(doc) shape_tool = XCAFDoc_DocumentTool.ShapeTool_s(doc.Main()) color_tool = XCAFDoc_DocumentTool.ColorTool_s(doc.Main()) free_labels = TDF_LabelSequence() shape_tool.GetFreeShapes(free_labels) print( f"Found {free_labels.Length()} root shape(s), tessellating " f"(linear={args.linear_deflection}mm, angular={args.angular_deflection}rad) …" ) # ── Tessellate ──────────────────────────────────────────────────────────── engine = getattr(args, "tessellation_engine", "gmsh") if engine == "gmsh": # GMSH: tessellate per-solid (same strategy as export_step_to_gltf.py). # 1. BRepMesh baseline on full root shape — catches free faces/shells. # 2. GMSH override per unique SOLID — uniform Frontal-Delaunay topology. # Skips REVERSED (mirrored) solids to avoid inverted-Jacobian issues. # Deduplication via IsSame() (TShape pointer), not id() (unreliable in OCP). from OCP.TopExp import TopExp_Explorer as _Explorer from OCP.TopAbs import (TopAbs_SOLID as _SOLID, TopAbs_SHELL as _SHELL, TopAbs_REVERSED as _REVERSED) from OCP.TopLoc import TopLoc_Location as _TopLoc_Location for i in range(1, free_labels.Length() + 1): root_shape = shape_tool.GetShape_s(free_labels.Value(i)) if root_shape.IsNull(): continue # Step 1: BRepMesh baseline BRepMesh_IncrementalMesh( root_shape, args.linear_deflection, False, args.angular_deflection, True) # Step 2: GMSH override for unique SOLID shapes _seen_shapes: list = [] solids = [] exp = _Explorer(root_shape, _SOLID) while exp.More(): solids.append(exp.Current()) exp.Next() if not solids: exp = _Explorer(root_shape, _SHELL) while exp.More(): solids.append(exp.Current()) exp.Next() from OCP.TopoDS import TopoDS_Compound as _Compound from OCP.BRep import BRep_Builder as _BBuilder eligible = [] for solid in solids: if solid.Orientation() == _REVERSED: continue if any(solid.IsSame(s) for s in _seen_shapes): continue eligible.append(solid.Located(_TopLoc_Location())) _seen_shapes.append(solid) if eligible: try: if len(eligible) == 1: _tessellate_with_gmsh( eligible[0], args.linear_deflection, args.angular_deflection) else: compound = _Compound() bb = _BBuilder() bb.MakeCompound(compound) for s in eligible: bb.Add(compound, s) _tessellate_with_gmsh( compound, args.linear_deflection, args.angular_deflection) except Exception as exc: print(f"WARNING: GMSH fallback to BRepMesh: {exc}", file=sys.stderr) print(f"GMSH: {len(eligible)} eligible solids out of {len(solids)} total") else: for i in range(1, free_labels.Length() + 1): shape = shape_tool.GetShape_s(free_labels.Value(i)) if not shape.IsNull(): BRepMesh_IncrementalMesh( shape, args.linear_deflection, False, args.angular_deflection, True) print("Tessellation complete.") # ── Apply colors ────────────────────────────────────────────────────── if color_map: try: _apply_color_map(shape_tool, color_tool, free_labels, color_map) print(f"Applied color_map ({len(color_map)} entries)") except Exception as exc: print(f"WARNING: color_map application failed (non-fatal): {exc}", file=sys.stderr) else: try: _apply_palette_colors(shape_tool, color_tool, free_labels) print("Applied palette colors") except Exception as exc: print(f"WARNING: palette colors failed (non-fatal): {exc}", file=sys.stderr) # ── Create USD stage ────────────────────────────────────────────────── stage = Usd.Stage.CreateNew(str(output_path)) UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.z) UsdGeom.SetStageMetersPerUnit(stage, 0.001) # mm; Blender handles m conversion on import root_prim = UsdGeom.Xform.Define(stage, "/Root") stage.SetDefaultPrim(root_prim.GetPrim()) # /Root/Assembly carries the OCC→Blender coordinate swap. # OCC is mm Z-up Y-forward; Blender/GLB convention is Z-up Y-backward. # Transform: (X_occ, Y_occ, Z_occ) → (X, -Z, Y) (= Rx(-90°)). # Authored as a USD row-vector matrix on the Assembly Xform so that all # child XCAF transforms (authored in OCC space) are correctly composed. assembly_xform = UsdGeom.Xform.Define(stage, "/Root/Assembly") assembly_xform.AddTransformOp().Set(Gf.Matrix4d( 1, 0, 0, 0, 0, 0, 1, 0, 0, -1, 0, 0, 0, 0, 0, 1, )) stage.DefinePrim("/Root/Looks", "Scope") # ── Walk XCAF tree → author USD prims (hierarchical) ────────────────── # Sharp/seam edges are extracted per-part inside _author_xcaf_to_usd # (in definition space, matching definition-space mesh vertices). existing_keys: set = set() manifest_parts: list = [] counters = {"n_parts": 0, "n_empty": 0} for root_idx in range(1, free_labels.Length() + 1): root_label = free_labels.Value(root_idx) root_names: set = set() _author_xcaf_to_usd( stage, shape_tool, color_tool, root_label, "/Root/Assembly", "", existing_keys, mat_map_lower, color_map, args, manifest_parts, counters, root_names, ) n_parts = counters["n_parts"] n_empty = counters["n_empty"] stage.Save() sz = output_path.stat().st_size // 1024 if output_path.exists() else 0 n_mat_assigned = sum(1 for p in manifest_parts if p.get("canonical_material")) print(f"USD exported: {output_path.name} ({sz} KB), " f"{n_parts} parts, {n_empty} empty shapes skipped, " f"{n_mat_assigned}/{n_parts} material primvars written") # ── Stdout manifest (one line, parsed by Celery task) ───────────────────── print(f"MANIFEST_JSON: {json.dumps({'parts': manifest_parts})}") try: main() except SystemExit: raise except Exception: traceback.print_exc() sys.exit(1)