HartOMat/render-worker/scripts/export_step_to_gltf.py

"""OCC-native STEP → GLB export script.

Reads a STEP file via OCP/XCAF (preserving part names and embedded colors),
tessellates with BRepMesh, optionally applies per-part hex colors, and writes
a binary GLB in meters (Y-up, glTF convention).

No Blender required. Uses the same OCP bindings that cadquery ships with.

Usage:
    python3 export_step_to_gltf.py \
        --step_path /path/to/file.stp \
        --output_path /path/to/output.glb \
        [--linear_deflection 0.1] \
        [--angular_deflection 0.5] \
        [--color_map '{"RingInner": "#4C9BE8", "RingOuter": "#E85B4C"}']

Exit 0 on success, exit 1 on failure.
"""
from __future__ import annotations

import argparse
import json
import sys
import traceback
from pathlib import Path

PALETTE_HEX = [
    "#4C9BE8", "#E85B4C", "#4CBE72", "#E8A84C", "#A04CE8",
    "#4CD4E8", "#E84CA8", "#7EC850", "#E86B30", "#5088C8",
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--step_path", required=True)
    parser.add_argument("--output_path", required=True)
    parser.add_argument(
        "--linear_deflection", type=float, default=0.1,
        help="OCC linear deflection for tessellation (mm). Smaller = finer mesh. Default 0.1",
    )
    parser.add_argument(
        "--angular_deflection", type=float, default=0.5,
        help="OCC angular deflection (radians). Default 0.5",
    )
    parser.add_argument(
        "--color_map", default="{}",
        help='JSON dict mapping part name → hex color, e.g. \'{"Ring": "#4C9BE8"}\'',
    )
    parser.add_argument(
        "--sharp_threshold", type=float, default=20.0,
        help="Dihedral angle threshold (degrees) for sharp B-rep edge detection. Default 20.0",
    )
    return parser.parse_args()


def _hex_to_occ_color(hex_color: str):
    """Convert '#RRGGBB' → Quantity_Color (linear float)."""
    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)
    r = int(h[0:2], 16) / 255.0
    g = int(h[2:4], 16) / 255.0
    b = int(h[4:6], 16) / 255.0
    return Quantity_Color(r, g, b, Quantity_TOC_RGB)


def _apply_color_map(shape_tool, color_tool, free_labels, color_map: dict) -> None:
    """Apply hex colors from color_map to matching shapes by name (case-insensitive substring)."""
    from OCP.TDF import TDF_LabelSequence
    from OCP.TDataStd import TDataStd_Name
    from OCP.XCAFDoc import XCAFDoc_ShapeTool

    # XCAFDoc_ColorType: XCAFDoc_ColorGen=0, XCAFDoc_ColorSurf=1, XCAFDoc_ColorCurv=2
    try:
        from OCP.XCAFDoc import XCAFDoc_ColorSurf as COLOR_SURF
    except ImportError:
        COLOR_SURF = 1  # integer fallback

    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), 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:
    """Assign palette colors to leaf shapes when no color_map is provided."""
    from OCP.TDF import TDF_LabelSequence
    from OCP.XCAFDoc import XCAFDoc_ShapeTool

    try:
        from OCP.XCAFDoc import XCAFDoc_ColorSurf as COLOR_SURF
    except ImportError:
        COLOR_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):
        occ_color = _hex_to_occ_color(PALETTE_HEX[idx % len(PALETTE_HEX)])
        color_tool.SetColor(label, occ_color, COLOR_SURF)


def _extract_sharp_edge_pairs(shape, sharp_threshold_deg: float = 20.0) -> list:
    """Extract geometrically sharp B-rep edges as dense curve sample segment pairs.

    For each edge shared by exactly 2 faces, evaluates the dihedral angle using
    PCurve-based surface normal evaluation. When the angle exceeds the threshold,
    samples the analytical 3D curve uniformly at 0.3mm intervals via
    GCPnts_UniformAbscissa. Consecutive sample pairs give the KD-tree in
    export_gltf.py enough density to find and mark the correct Blender mesh edges.

    Note: BRep_Tool.Polygon3D_s() and PolygonOnTriangulation_s() return None in
    XCAF compound context — the tessellation data is stored on component instances,
    not on the compound edges. Curve sampling bypasses this entirely.

    Args:
        shape: OCC TopoDS_Shape (tessellated with BRepMesh_IncrementalMesh)
        sharp_threshold_deg: dihedral angle threshold in degrees (default 20°)

    Returns:
        List of [[x0,y0,z0],[x1,y1,z1]] consecutive segment pairs in mm (OCC
        coordinate space, Z-up). Must be called BEFORE mm→m scaling.
    """
    import math as _math

    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 0.3mm — well below the KD-tree TOL=0.5mm in export_gltf.py.
    # Tessellation vertex spacing for default deflection is ~0.78-1.55mm, so at
    # least one consecutive sample pair will straddle each tessellation edge.
    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)
        if faces.Size() < 2:
            n_checked += 1
            continue

        face_shapes = list(faces)
        if len(face_shapes) < 2:
            n_checked += 1
            continue

        n_checked += 1
        try:
            edge = _TopoDS.Edge_s(edge_shape)
            face1 = _TopoDS.Face_s(face_shapes[0])
            face2 = _TopoDS.Face_s(face_shapes[1])

            # PCurve-based normal evaluation at edge midpoint
            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))
            # Use exterior (supplement) angle so convex and concave edges both work
            if angle_deg > 90.0:
                angle_deg = 180.0 - angle_deg

            if angle_deg <= sharp_threshold_deg:
                continue  # smooth transition — skip

            n_sharp += 1

            # Sample the analytical 3D curve at fixed arc-length intervals.
            # GCPnts_UniformAbscissa works on the exact B-rep curve regardless of
            # whether tessellation polygon data is stored on the edge or not.
            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):
                            t = sampler.Parameter(j)
                            p = curve3d.Value(t)
                            pts.append([round(p.X(), 4), round(p.Y(), 4), round(p.Z(), 4)])
            except Exception:
                pts = []

            if len(pts) < 2:
                continue

            # Consecutive segment pairs — KD-tree in export_gltf.py maps each
            # endpoint to its nearest Blender vertex; if they differ and share a
            # mesh edge, that edge is marked sharp+seam.
            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 _inject_glb_extras(glb_path: Path, extras: dict) -> None:
    """Patch a GLB binary to add/update scenes[0].extras JSON field.

    The GLB format stores a JSON chunk immediately after the 12-byte header.
    We re-serialize it with the new extras and update chunk + total lengths.
    No external dependencies — pure stdlib struct/json.
    """
    import struct as _struct

    data = glb_path.read_bytes()
    # GLB header: magic(4) + version(4) + total_length(4) = 12 bytes
    # JSON chunk: chunk_length(4) + chunk_type(4) + chunk_data(chunk_length bytes)
    json_len = _struct.unpack_from("<I", data, 12)[0]
    json_type = _struct.unpack_from("<I", data, 16)[0]
    if json_type != 0x4E4F534A:  # "JSON"
        print("WARNING: _inject_glb_extras: unexpected chunk type, skipping extras injection",
              file=sys.stderr)
        return

    j = json.loads(data[20: 20 + json_len])
    if "scenes" in j and j["scenes"]:
        existing = j["scenes"][0].get("extras") or {}
        existing.update(extras)
        j["scenes"][0]["extras"] = existing
    else:
        j.setdefault("extras", {}).update(extras)

    new_json = json.dumps(j, separators=(",", ":"))
    # Pad to 4-byte boundary with spaces (required by GLB spec)
    pad = (4 - len(new_json) % 4) % 4
    new_json_bytes = new_json.encode() + b" " * pad

    rest = data[20 + json_len:]  # BIN chunk and anything after
    new_chunk = _struct.pack("<II", len(new_json_bytes), 0x4E4F534A) + new_json_bytes
    new_total = 12 + len(new_chunk) + len(rest)
    new_header = _struct.pack("<III", 0x46546C67, 2, new_total)
    glb_path.write_bytes(new_header + new_chunk + rest)


def main() -> None:
    args = parse_args()
    color_map: dict = json.loads(args.color_map)

    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, TCollection_AsciiString
    from OCP.TDF import TDF_LabelSequence
    from OCP.BRepMesh import BRepMesh_IncrementalMesh
    from OCP.IFSelect import IFSelect_RetDone
    from OCP.Message import Message_ProgressRange

    # --- Init XDE document ---
    app = XCAFApp_Application.GetApplication_s()
    doc = TDocStd_Document(TCollection_ExtendedString("MDTV-CAF"))
    app.InitDocument(doc)

    # --- Read STEP into XDE (preserves part names + embedded colors) ---
    reader = STEPCAFControl_Reader()
    reader.SetNameMode(True)
    reader.SetColorMode(True)
    reader.SetLayerMode(True)
    status = reader.ReadFile(args.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())

    # --- Tessellate all free shapes ---
    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) …")

    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,   # isRelative
                args.angular_deflection,
                True,    # isInParallel
            )

    # --- Extract sharp B-rep edge pairs (before mm→m scaling so coords are in mm) ---
    # Collect all free shapes into one list for the extraction function.
    # The extraction uses the freshly tessellated XCAF shapes.
    sharp_pairs: list = []
    try:
        for i in range(1, free_labels.Length() + 1):
            root_shape = shape_tool.GetShape_s(free_labels.Value(i))
            if not root_shape.IsNull():
                pairs = _extract_sharp_edge_pairs(root_shape, args.sharp_threshold)
                sharp_pairs.extend(pairs)
        print(f"Total OCC sharp segment pairs: {len(sharp_pairs)}")
    except Exception as _exc:
        print(f"WARNING: sharp edge extraction failed (non-fatal): {_exc}", file=sys.stderr)
        sharp_pairs = []

    # --- Apply colors ---
    if color_map:
        _apply_color_map(shape_tool, color_tool, free_labels, color_map)
        print(f"Applied color_map ({len(color_map)} entries)")
    else:
        _apply_palette_colors(shape_tool, color_tool, free_labels)
        print("Applied palette colors (no color_map provided)")

    # --- Scale shapes mm → m before GLB export ---
    # RWMesh_CoordinateSystemConverter is not wrapped in OCP Python bindings.
    # Pre-scale each free shape by 0.001 (mm → m) using BRepBuilderAPI_Transform.
    from OCP.gp import gp_Trsf
    from OCP.BRepBuilderAPI import BRepBuilderAPI_Transform

    trsf = gp_Trsf()
    trsf.SetScaleFactor(0.001)

    for i in range(1, free_labels.Length() + 1):
        label = free_labels.Value(i)
        orig_shape = shape_tool.GetShape_s(label)
        if not orig_shape.IsNull():
            scaled = BRepBuilderAPI_Transform(orig_shape, trsf, True).Shape()
            shape_tool.SetShape(label, scaled)

    print("Shapes scaled mm → m")

    # --- Export GLB via RWGltf_CafWriter ---
    from OCP.RWGltf import RWGltf_CafWriter

    writer = RWGltf_CafWriter(TCollection_AsciiString(args.output_path), True)  # True = binary GLB
    # Z-up → Y-up rotation is applied by RWGltf_CafWriter by default (OCC 7.6+).

    # Perform export
    try:
        from OCP.TColStd import TColStd_IndexedDataMapOfStringString
        metadata = TColStd_IndexedDataMapOfStringString()
        ok = writer.Perform(doc, metadata, Message_ProgressRange())
    except TypeError:
        # Older API without metadata dict
        ok = writer.Perform(doc, Message_ProgressRange())

    out = Path(args.output_path)
    if not ok or not out.exists() or out.stat().st_size == 0:
        print(f"ERROR: RWGltf_CafWriter.Perform returned ok={ok}, file exists={out.exists()}",
              file=sys.stderr)
        sys.exit(1)

    print(f"GLB exported: {out.name} ({out.stat().st_size // 1024} KB)")

    # --- Inject sharp edge pairs into GLB extras ---
    # Blender 5.0 reads scenes[0].extras as scene custom properties on import,
    # making the data available to export_gltf.py as bpy.context.scene["key"].
    if sharp_pairs:
        try:
            _inject_glb_extras(out, {
                "schaeffler_sharp_edge_pairs": sharp_pairs,
                "schaeffler_sharp_threshold_deg": args.sharp_threshold,
            })
            print(f"Injected {len(sharp_pairs)} sharp edge segment pairs into GLB extras")
        except Exception as _exc:
            print(f"WARNING: GLB extras injection failed (non-fatal): {_exc}", file=sys.stderr)


try:
    main()
except SystemExit:
    raise
except Exception:
    traceback.print_exc()
    sys.exit(1)