HartOMat/render-worker/scripts/export_gltf.py

"""Blender headless script: export a STEP-derived scene as a production GLB.

Usage:
    blender --background --python export_gltf.py -- \\
        --stl_path /path/to/file.stl \\
        --output_path /path/to/output.glb \\
        [--asset_library_blend /path/to/library.blend] \\
        [--material_map '{"SrcMat": "LibMat"}']

The script:
1. Imports the STL file (with mm→m scale).
2. Optionally applies asset library materials from a .blend.
3. Exports as GLB (Draco-compressed if available, otherwise standard).
"""
from __future__ import annotations

import argparse
import json
import sys
import traceback


def parse_args() -> argparse.Namespace:
    argv = sys.argv
    if "--" not in argv:
        print("No arguments after --", file=sys.stderr)
        sys.exit(1)
    rest = argv[argv.index("--") + 1:]
    parser = argparse.ArgumentParser()
    parser.add_argument("--glb_path", required=True,
                        help="Geometry GLB from export_step_to_gltf.py (already in metres)")
    parser.add_argument("--output_path", required=True)
    parser.add_argument("--asset_library_blend", default=None)
    parser.add_argument("--material_map", default="{}")
    parser.add_argument("--smooth_angle", type=float, default=30.0,
                        help="Auto-smooth angle in degrees (default 30)")
    parser.add_argument("--mesh_attributes", default="{}",
                        help="JSON dict from cad_file.mesh_attributes (sharp_edge_pairs etc.)")
    return parser.parse_args(rest)


def _apply_sharp_edges_from_occ(mesh_objects: list, sharp_edge_pairs: list) -> None:
    """Mark edges sharp using OCC vertex-pair data (same approach as blender_render.py).

    sharp_edge_pairs: [[x0,y0,z0],[x1,y1,z1]] in mm.
    Blender mesh coords are in metres (×0.001 scale already applied by OCC export).
    """
    if not sharp_edge_pairs:
        return

    import bmesh
    import mathutils

    SCALE = 0.001   # mm → m
    TOL = 0.0005    # 0.5 mm tolerance in metres

    # OCC STEP space (Z-up, mm) → Blender (Z-up, m):
    # RWGltf applies Z→Y-up, Blender import applies Y→Z-up.
    # Net: Blender(X, Y, Z) = OCC(X*0.001, -Z*0.001, Y*0.001)
    occ_pairs = []
    for pair in sharp_edge_pairs:
        v0 = mathutils.Vector((pair[0][0] * SCALE, -pair[0][2] * SCALE, pair[0][1] * SCALE))
        v1 = mathutils.Vector((pair[1][0] * SCALE, -pair[1][2] * SCALE, pair[1][1] * SCALE))
        occ_pairs.append((v0, v1))

    marked_total = 0
    for obj in mesh_objects:
        bm = bmesh.new()
        bm.from_mesh(obj.data)
        bm.verts.ensure_lookup_table()
        bm.edges.ensure_lookup_table()

        # Build KD-tree in WORLD space — OCC pairs are world coords, but mesh
        # vertices are in local space (assembly node transform in GLB hierarchy).
        world_mat = obj.matrix_world
        kd = mathutils.kdtree.KDTree(len(bm.verts))
        for v in bm.verts:
            kd.insert(world_mat @ v.co, v.index)
        kd.balance()

        marked = 0
        for v0_occ, v1_occ in occ_pairs:
            _co0, idx0, dist0 = kd.find(v0_occ)
            _co1, idx1, dist1 = kd.find(v1_occ)
            if dist0 > TOL or dist1 > TOL:
                continue
            if idx0 == idx1:
                continue  # degenerate — both endpoints map to same vertex
            bv0, bv1 = bm.verts[idx0], bm.verts[idx1]
            edge = bm.edges.get((bv0, bv1)) or bm.edges.get((bv1, bv0))
            if edge is not None:
                # Mark sharp (for normal splitting) AND seam (for UV unwrap).
                # Both are needed: sharp controls glTF vertex splits / shading;
                # seam defines UV island boundaries for correct UV unwrapping.
                edge.smooth = False
                edge.seam = True
                marked += 1

        bm.to_mesh(obj.data)
        bm.free()
        marked_total += marked

    print(f"OCC sharp edges applied: {marked_total} edges marked across {len(mesh_objects)} objects")


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

    import bpy  # type: ignore[import]
    import math as _math
    import re as _re

    # Clean scene
    bpy.ops.wm.read_factory_settings(use_empty=True)

    # Import geometry GLB from export_step_to_gltf.py (already in metres, Y-up)
    bpy.ops.import_scene.gltf(filepath=args.glb_path)
    mesh_objects = [o for o in bpy.data.objects if o.type == "MESH"]
    print(f"Imported geometry GLB: {args.glb_path} ({len(mesh_objects)} mesh objects)")

    # Read OCC sharp edge pairs embedded by export_step_to_gltf.py into GLB extras.
    # Blender 5.0 maps glTF scenes[0].extras as scene custom properties on import.
    # These take priority over the mesh_attributes CLI argument (which only has 2
    # endpoints per edge — see V02 refactor for why this matters).
    glb_sharp_pairs = bpy.context.scene.get("schaeffler_sharp_edge_pairs") or []
    if glb_sharp_pairs:
        print(f"Loaded {len(glb_sharp_pairs)} OCC sharp edge pairs from GLB extras")

    # Remove OCC-baked custom normals from the geometry GLB.
    # RWGltf_CafWriter embeds per-corner normals from OCC tessellation as a
    # 'custom_normal' attribute (CORNER, INT16_2D). If left in place, Blender's
    # glTF exporter re-exports these pre-baked normals unchanged, ignoring our
    # shade_smooth_by_angle processing and sharp edge marks entirely.
    # Removing this attribute forces Blender to recompute normals from scratch.
    cleared_normals = 0
    for obj in mesh_objects:
        if "custom_normal" in obj.data.attributes:
            obj.data.attributes.remove(obj.data.attributes["custom_normal"])
            cleared_normals += 1
    if cleared_normals:
        print(f"Cleared OCC custom_normal attribute from {cleared_normals} mesh objects")

    # Mark sharp edges and seams using the configured angle threshold.
    # We use Blender's edit-mode operators (mark_sharp + mark_seam) rather than
    # shade_smooth_by_angle alone, because:
    # 1. mark_sharp() sets the sharp_edge boolean attribute on edges — the glTF
    #    exporter creates vertex splits (duplicate vertices with different normals)
    #    at sharp edges, which is how glTF encodes hard edges.
    # 2. mark_seam() ensures UV splits at the same edges (stepper-addon behaviour).
    # Note: calc_normals_split() was removed in Blender 5.0 — not needed here
    #       because export_apply=True triggers vertex splitting automatically.
    smooth_rad = _math.radians(args.smooth_angle)
    print(f"Marking sharp edges + seams at {args.smooth_angle}° ({smooth_rad:.3f} rad)")

    bpy.ops.object.select_all(action='DESELECT')
    total_sharp = 0
    for obj in mesh_objects:
        bpy.context.view_layer.objects.active = obj
        obj.select_set(True)

        # Set all faces smooth
        bpy.ops.object.mode_set(mode='OBJECT')
        for poly in obj.data.polygons:
            poly.use_smooth = True

        # Enter edit mode, deselect, select sharp edges by angle, mark sharp+seam
        bpy.ops.object.mode_set(mode='EDIT')
        bpy.ops.mesh.select_all(action='DESELECT')
        bpy.ops.mesh.edges_select_sharp(sharpness=smooth_rad)
        bpy.ops.mesh.mark_sharp()
        bpy.ops.mesh.mark_seam()
        bpy.ops.object.mode_set(mode='OBJECT')

        # Count how many edges were marked
        n_sharp = sum(1 for e in obj.data.edges if e.use_edge_sharp)
        total_sharp += n_sharp
        obj.select_set(False)

    print(f"Marked {total_sharp} sharp/seam edges across {len(mesh_objects)} objects")

    # Apply OCC sharp edges from GLB extras (V02: dense tessellation segment pairs).
    # Prefer GLB-embedded pairs over mesh_attributes CLI argument — the GLB extras
    # contain the full tessellated polyline for each sharp B-rep edge (all intermediate
    # points), while mesh_attributes only has 2 endpoints per edge (too sparse for
    # reliable KD-tree matching). Fall back to mesh_attributes if GLB extras absent.
    occ_pairs = list(glb_sharp_pairs) or (mesh_attributes.get("sharp_edge_pairs") or [])
    if occ_pairs:
        _apply_sharp_edges_from_occ(mesh_objects, occ_pairs)

    # Apply asset library materials if provided.
    # link=False (append) is required: the GLTF exporter can only traverse
    # local (appended) Principled BSDF node trees to extract PBR values.
    #
    # Matching strategy (mirrors blender_render.py):
    # Build mat_map_lower with BOTH the original key AND the _AF-stripped key,
    # so keys like "RingOuter_AF0" match object names "RingOuter" and vice-versa.
    # Object names from RWGltf_CafWriter preserve the original STEP part name
    # (including any _AF suffixes), so we strip from both sides.
    if args.asset_library_blend and material_map:
        mat_map_lower: dict = {}
        for k, v in material_map.items():
            kl = k.lower().strip()
            mat_map_lower[kl] = v
            # Also add the _AF-stripped version so either form matches
            stripped = kl
            prev = None
            while prev != stripped:
                prev = stripped
                stripped = _re.sub(r'_af\d+$', '', stripped)
            if stripped != kl:
                mat_map_lower.setdefault(stripped, v)

        needed = set(mat_map_lower.values())

        # Append materials from library (link=False so glTF exporter can read nodes)
        appended: dict = {}
        for mat_name in needed:
            try:
                bpy.ops.wm.append(
                    filepath=f"{args.asset_library_blend}/Material/{mat_name}",
                    directory=f"{args.asset_library_blend}/Material/",
                    filename=mat_name,
                    link=False,
                )
                if mat_name in bpy.data.materials:
                    appended[mat_name] = bpy.data.materials[mat_name]
                    print(f"Appended material: {mat_name}")
                else:
                    print(f"WARNING: material '{mat_name}' not found in library after append",
                          file=sys.stderr)
            except Exception as exc:
                print(f"WARNING: failed to append material '{mat_name}': {exc}", file=sys.stderr)

        if appended:
            assigned = 0
            assigned_names: set = set()
            for obj in mesh_objects:
                # Strip Blender's .001/.002 deduplication suffix
                base_name = _re.sub(r'\.\d{3}$', '', obj.name)
                # Also strip _AF suffix from object name so both directions match
                prev = None
                while prev != base_name:
                    prev = base_name
                    base_name = _re.sub(r'_AF\d+$', '', base_name, flags=_re.IGNORECASE)

                lower_base = base_name.lower().strip()
                mat_name = mat_map_lower.get(lower_base)

                # Prefix fallback for sub-assembly nodes
                if not mat_name:
                    for key, val in sorted(mat_map_lower.items(), key=lambda x: len(x[0]), reverse=True):
                        if len(key) >= 3 and len(lower_base) >= 3 and (
                            lower_base.startswith(key) or key.startswith(lower_base)
                        ):
                            mat_name = val
                            break

                if mat_name and mat_name in appended:
                    # Make mesh data single-user before modifying material slots;
                    # otherwise clearing materials on a shared data block removes
                    # slots from ALL objects that share it.
                    if obj.data.users > 1:
                        obj.data = obj.data.copy()
                    obj.data.materials.clear()
                    obj.data.materials.append(appended[mat_name])
                    assigned += 1
                    assigned_names.add(obj.name)
                else:
                    pass  # name-matching miss — may be covered by single-material fallback below
            print(f"Material substitution: {assigned}/{len(mesh_objects)} mesh objects assigned")

            # Single-material fallback: if only one library material was loaded,
            # apply it to every object that name-matching missed.
            # (mat_map_lower may contain unresolvable pass-through values like
            # "Stahl; Durotect CMT", so checking appended is more reliable.)
            if len(appended) == 1:
                default_mat_name, default_mat = next(iter(appended.items()))
                if default_mat:
                    fallback = 0
                    for obj in mesh_objects:
                        if obj.name not in assigned_names:
                            if obj.data.users > 1:
                                obj.data = obj.data.copy()
                            obj.data.materials.clear()
                            obj.data.materials.append(default_mat)
                            fallback += 1
                    if fallback:
                        print(f"Single-material fallback: applied '{default_mat_name}' to {fallback} unmatched objects")

    # Purge orphan data-blocks (palette materials mat_0/mat_1/... from the geometry
    # GLB that now have users=0 after library material substitution).
    # This prevents stale materials from appearing as duplicates in the export.
    try:
        bpy.ops.outliner.orphans_purge(do_recursive=True)
    except Exception:
        pass  # non-critical; export proceeds regardless

    # Store the sharp angle in the scene so it is embedded in the GLB extras.
    # After importing the production GLB in Blender, running restore_sharp_marks.py
    # reads this value and re-applies mark_sharp()+mark_seam() on all mesh objects.
    bpy.context.scene["schaeffler_sharp_angle_deg"] = args.smooth_angle

    # Export production GLB with full PBR material data.
    # export_extras=True embeds scene custom properties (incl. schaeffler_sharp_angle_deg)
    # in the glTF scenes[0].extras JSON field, surviving the round-trip intact.
    try:
        bpy.ops.export_scene.gltf(
            filepath=args.output_path,
            export_format="GLB",
            export_apply=True,
            use_selection=False,
            export_materials="EXPORT",
            export_image_format="AUTO",
            export_extras=True,
        )
    except Exception as exc:
        print(f"GLB export failed: {exc}", file=sys.stderr)
        sys.exit(1)

    print(f"Production GLB exported to {args.output_path}")


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