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("--stl_path", required=True)
    parser.add_argument("--output_path", required=True)
    parser.add_argument("--asset_library_blend", default=None)
    parser.add_argument("--material_map", default="{}")
    parser.add_argument("--sharp_edges_json", default="[]",
                        help="JSON array of [x, y, z] midpoints (mm) to mark as sharp edges")
    return parser.parse_args(rest)


def mark_sharp_edges_by_proximity(midpoints_mm: list, threshold_mm: float = 1.0) -> None:
    """Mark Blender mesh edges as sharp based on proximity to OCC-derived midpoints.

    midpoints_mm: list of [x, y, z] in mm (from OCC coordinate space).
    After STL import + scale-apply (mm→m), Blender vertices are in meters, so we
    convert the edge midpoint back to mm before comparing.
    threshold_mm: snap distance in mm (default 1.0 mm).
    """
    if not midpoints_mm:
        return

    import bpy  # type: ignore[import]
    import math

    for obj in bpy.data.objects:
        if obj.type != "MESH":
            continue
        mesh = obj.data
        # Blender 4.1+ removed use_auto_smooth — use shade_smooth_by_angle instead
        bpy.context.view_layer.objects.active = obj
        obj.select_set(True)
        try:
            bpy.ops.object.shade_smooth_by_angle(angle=math.radians(30))
        except Exception:
            pass  # fallback: stay flat-shaded
        mw = obj.matrix_world
        for edge in mesh.edges:
            v1 = mw @ mesh.vertices[edge.vertices[0]].co
            v2 = mw @ mesh.vertices[edge.vertices[1]].co
            # Convert Blender meters → mm for comparison
            mid_mm = [
                (v1.x + v2.x) / 2 * 1000,
                (v1.y + v2.y) / 2 * 1000,
                (v1.z + v2.z) / 2 * 1000,
            ]
            for hint in midpoints_mm:
                dist_sq = sum((a - b) ** 2 for a, b in zip(mid_mm, hint))
                if dist_sq < threshold_mm ** 2:
                    edge.use_edge_sharp = True
                    break


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

    import bpy  # type: ignore[import]

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

    # Import STL (bpy.ops.wm.stl_import is the Blender 4.0+ API)
    bpy.ops.wm.stl_import(filepath=args.stl_path)

    # Scale mm → m
    for obj in bpy.context.selected_objects:
        obj.scale = (0.001, 0.001, 0.001)
        bpy.context.view_layer.objects.active = obj
        bpy.ops.object.transform_apply(scale=True)

    # Apply smooth shading with 30° angle threshold (Blender 4.1+ API)
    import math as _math
    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=_math.radians(30))
            except Exception:
                pass

    # Mark sharp edges for better UV seams
    if sharp_edge_midpoints:
        mark_sharp_edges_by_proximity(sharp_edge_midpoints)
        print(f"Marked sharp edges from {len(sharp_edge_midpoints)} hint points")

    # Apply asset library materials if provided.
    # link=False (append) is required for GLB export: the GLTF exporter can only
    # traverse local (appended) Principled BSDF node trees to extract PBR values.
    # Linked materials are external references whose node data is not accessible.
    if args.asset_library_blend and material_map:
        import os
        sys.path.insert(0, os.path.dirname(__file__))
        from asset_library import apply_asset_library_materials
        apply_asset_library_materials(args.asset_library_blend, material_map, link=False)

    # Export GLB with full PBR material data
    # Note: export_colors was removed in Blender 4.x — do not pass it.
    try:
        bpy.ops.export_scene.gltf(
            filepath=args.output_path,
            export_format="GLB",
            export_apply=True,
            use_selection=False,
            export_materials="EXPORT",       # export all materials (Principled BSDF → glTF PBR)
            export_image_format="AUTO",      # embed textures (base color, normal, roughness maps)
        )
    except Exception as exc:
        print(f"GLB export failed: {exc}", file=sys.stderr)
        sys.exit(1)

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


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