1207 lines
49 KiB
Python
1207 lines
49 KiB
Python
"""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 os
|
||
import re
|
||
import sys
|
||
import traceback
|
||
from pathlib import Path
|
||
|
||
from _blender_materials import build_mat_map_lower, lookup_material_name
|
||
|
||
|
||
# ── 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)
|
||
binding_api = UsdShade.MaterialBindingAPI.Apply(mesh.GetPrim())
|
||
binding_api.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"
|
||
|
||
|
||
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."""
|
||
return lookup_material_name(source_name, mat_map_lower, part_key)
|
||
|
||
|
||
def _atomic_output_path(output_path: Path) -> Path:
|
||
return output_path.with_name(
|
||
f".{output_path.stem}.{os.getpid()}.tmp{output_path.suffix}"
|
||
)
|
||
|
||
|
||
# ── 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 ──────────────────────────────────────────────────
|
||
temp_output_path = _atomic_output_path(output_path)
|
||
try:
|
||
if temp_output_path.exists():
|
||
temp_output_path.unlink()
|
||
except OSError:
|
||
pass
|
||
|
||
stage = Usd.Stage.CreateNew(str(temp_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"]
|
||
try:
|
||
stage.Save()
|
||
if temp_output_path.exists():
|
||
os.chmod(temp_output_path, 0o664)
|
||
os.replace(temp_output_path, output_path)
|
||
os.chmod(output_path, 0o664)
|
||
except Exception:
|
||
try:
|
||
if temp_output_path.exists():
|
||
temp_output_path.unlink()
|
||
except OSError:
|
||
pass
|
||
raise
|
||
|
||
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)
|