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)")
    return parser.parse_args(rest)


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

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

    # 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)
    print(f"Imported geometry GLB: {args.glb_path} "
          f"({len([o for o in bpy.data.objects if o.type == 'MESH'])} mesh objects)")

    # Apply smooth shading using the configured angle threshold
    smooth_rad = _math.radians(args.smooth_angle)
    print(f"Applying smooth shading at {args.smooth_angle}° ({smooth_rad:.3f} rad)")
    for obj in bpy.data.objects:
        if obj.type == "MESH":
            bpy.context.view_layer.objects.active = obj
            obj.select_set(True)
            try:
                bpy.ops.object.shade_smooth_by_angle(angle=smooth_rad)
            except Exception:
                # Fallback for older Blender API
                bpy.ops.object.shade_smooth()
                if obj.data.use_auto_smooth is not None:
                    obj.data.use_auto_smooth = True
                    obj.data.auto_smooth_angle = smooth_rad

    # 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.
    #
    # IMPORTANT: OCC-exported GLBs name materials generically (mat_0, mat_1, …)
    # but preserve STEP part names as mesh OBJECT names. We therefore match by
    # obj.name, not by slot.material.name (which is how blender_render.py works).
    if args.asset_library_blend and material_map:
        import re as _re
        mat_map_lower = {k.lower().strip(): v for k, v in material_map.items()}
        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
            mesh_objects = [o for o in bpy.data.objects if o.type == "MESH"]
            for obj in mesh_objects:
                # Strip Blender's .001/.002 deduplication suffix
                base_name = _re.sub(r'\.\d{3}$', '', obj.name)
                # Strip OCC assembly-instance suffix (_AF0, _AF1, … added by
                # RWGltf_CafWriter when the same part appears multiple times).
                # Apply repeatedly in case of nested suffixes (_AF0_AF1, etc.)
                prev = None
                while prev != base_name:
                    prev = base_name
                    base_name = _re.sub(r'_AF\d+$', '', base_name, flags=_re.IGNORECASE)

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

                # Prefix fallback: some sub-assembly nodes have names that
                # extend a known key (e.g. key="Ring" matches "Ring_inner_AF0").
                # Sort by key length descending so the most-specific key wins.
                if not mat_name:
                    lower_base = base_name.lower().strip()
                    for key, val in sorted(mat_map_lower.items(), key=lambda x: len(x[0]), reverse=True):
                        if len(key) >= 5 and len(lower_base) >= 5 and (
                            lower_base.startswith(key) or key.startswith(lower_base)
                        ):
                            mat_name = val
                            break

                if mat_name and mat_name in appended:
                    obj.data.materials.clear()
                    obj.data.materials.append(appended[mat_name])
                    assigned += 1
            print(f"Material substitution: {assigned}/{len(mesh_objects)} mesh objects assigned")

    # Export production GLB with full PBR material data
    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",
        )
    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)