HartOMat/render-worker/scripts/_blender_scene.py

"""Scene-level helpers for Blender headless renders."""
from __future__ import annotations

import math


def ensure_collection(name: str):
    """Return a collection by name, creating it if needed."""
    import bpy  # type: ignore[import]

    if name in bpy.data.collections:
        return bpy.data.collections[name]
    col = bpy.data.collections.new(name)
    bpy.context.scene.collection.children.link(col)
    return col


def apply_smooth_batch(parts: list, angle_deg: float) -> None:
    """Apply smooth shading to ALL parts in a single operator call.

    bpy.ops.object.shade_smooth_by_angle() operates on all selected objects
    at once (one C-level call), so batching reduces O(n) operator overhead to O(1).
    Per-part calls cost ~90ms each × 175 parts = 16s; batch call costs ~0.2s total.
    """
    import bpy  # type: ignore[import]

    bpy.ops.object.select_all(action='DESELECT')
    mesh_parts = [p for p in parts if p.type == 'MESH']
    for part in mesh_parts:
        part.select_set(True)
    if not mesh_parts:
        return
    bpy.context.view_layer.objects.active = mesh_parts[0]
    if angle_deg > 0:
        try:
            bpy.ops.object.shade_smooth_by_angle(angle=math.radians(angle_deg))
        except AttributeError:
            bpy.ops.object.shade_smooth()
            for part in mesh_parts:
                if hasattr(part.data, 'use_auto_smooth'):
                    part.data.use_auto_smooth = True
                    part.data.auto_smooth_angle = math.radians(angle_deg)
    else:
        bpy.ops.object.shade_flat()
    bpy.ops.object.select_all(action='DESELECT')


def apply_sharp_edges_from_occ(parts: list, sharp_edge_pairs: list) -> None:
    """Mark edges sharp using OCC-derived vertex-pair data.

    `sharp_edge_pairs` is a list of [[x0,y0,z0],[x1,y1,z1]] in mm.
    Blender mesh coordinates are in metres (STEP mm * 0.001 scale applied).
    We match each OCC vertex pair against bmesh vertex positions with a 0.5 mm
    tolerance (0.0005 m) and mark the matched edge as sharp.
    """
    if not sharp_edge_pairs:
        return

    import time as _time
    import bmesh  # type: ignore[import]
    import mathutils  # type: ignore[import]

    _t0 = _time.monotonic()

    SCALE = 0.001   # mm → m
    TOL   = 0.0005  # 0.5 mm 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))

    _t_convert = _time.monotonic()
    print(f"[blender_render] TIMING sharp_edges_convert={_t_convert - _t0:.3f}s ({len(occ_pairs)} pairs)", flush=True)

    # ── Spatial pre-filter: build a KD-tree over OCC pair midpoints ────────
    # For each part, query the midpoint KD-tree with the part's bbox radius
    # to get only nearby pairs instead of testing all N pairs × M parts.
    _t_spatial = _time.monotonic()
    pair_midpoints = []
    pair_radii = []  # half-length of each pair (max distance from midpoint to endpoint)
    for v0, v1 in occ_pairs:
        mid = (v0 + v1) * 0.5
        pair_midpoints.append(mid)
        pair_radii.append((v0 - mid).length)

    pair_kd = mathutils.kdtree.KDTree(len(pair_midpoints))
    for i, mid in enumerate(pair_midpoints):
        pair_kd.insert(mid, i)
    pair_kd.balance()
    _t_spatial_done = _time.monotonic()
    print(f"[blender_render] TIMING sharp_edges_spatial_index={_t_spatial_done - _t_spatial:.3f}s", flush=True)

    marked_total = 0
    kd_build_time = 0.0
    match_time = 0.0
    pairs_tested_total = 0
    for obj in parts:
        _t_kd = _time.monotonic()
        bm = bmesh.new()
        bm.from_mesh(obj.data)
        bm.verts.ensure_lookup_table()
        bm.edges.ensure_lookup_table()

        # Build KD-tree on vertices in WORLD space — OCC pairs are world coords,
        # but mesh vertices are in local space (assembly node transform in GLB).
        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()

        # Compute part's world-space bounding box center and search radius
        from mathutils import Vector  # type: ignore[import]
        corners = [world_mat @ Vector(c) for c in obj.bound_box]
        bbox_min = Vector((min(c.x for c in corners), min(c.y for c in corners), min(c.z for c in corners)))
        bbox_max = Vector((max(c.x for c in corners), max(c.y for c in corners), max(c.z for c in corners)))
        bbox_center = (bbox_min + bbox_max) * 0.5
        bbox_half_diag = (bbox_max - bbox_min).length * 0.5

        kd_build_time += _time.monotonic() - _t_kd

        _t_match = _time.monotonic()
        marked = 0

        # Query pair midpoints within bbox_half_diag + max_pair_radius + tolerance
        # This guarantees we don't miss any pair whose endpoints could be inside the bbox
        max_pair_radius = max(pair_radii) if pair_radii else 0.0
        search_radius = bbox_half_diag + max_pair_radius + TOL
        nearby = pair_kd.find_range(bbox_center, search_radius)
        pairs_tested_total += len(nearby)

        for _co, pair_idx, _dist in nearby:
            v0_occ, v1_occ = occ_pairs[pair_idx]
            _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 = bm.verts[idx0]
            bv1 = bm.verts[idx1]
            edge = bm.edges.get((bv0, bv1))
            if edge is None:
                edge = bm.edges.get((bv1, bv0))
            if edge is not None and edge.smooth:
                edge.smooth = False
                marked += 1
        match_time += _time.monotonic() - _t_match

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

    _total = _time.monotonic() - _t0
    pairs_skipped = len(occ_pairs) * len(parts) - pairs_tested_total
    print(f"[blender_render] TIMING sharp_edges={_total:.2f}s "
          f"(kd_build={kd_build_time:.2f}s, matching={match_time:.2f}s, "
          f"pairs={len(occ_pairs)}, parts={len(parts)}, marked={marked_total}, "
          f"tested={pairs_tested_total}, skipped={pairs_skipped})", flush=True)


def setup_shadow_catcher(parts: list) -> None:
    """Enable the Shadowcatcher collection in the template and position its plane.

    The template must contain a 'Shadowcatcher' collection with a 'Shadowcatcher'
    mesh object. The plane is moved to the lowest Z of the product bounding box.
    """
    import bpy  # type: ignore[import]
    from mathutils import Vector  # type: ignore[import]

    sc_col_name = "Shadowcatcher"
    sc_obj_name = "Shadowcatcher"

    # Enable the Shadowcatcher collection in all view layers
    for vl in bpy.context.scene.view_layers:
        def _enable_col_recursive(layer_col):
            if layer_col.collection.name == sc_col_name:
                layer_col.exclude = False
                layer_col.collection.hide_render = False
                layer_col.collection.hide_viewport = False
                return True
            for child in layer_col.children:
                if _enable_col_recursive(child):
                    return True
            return False
        _enable_col_recursive(vl.layer_collection)

    sc_obj = bpy.data.objects.get(sc_obj_name)
    if sc_obj:
        all_world_z = []
        for part in parts:
            for corner in part.bound_box:
                all_world_z.append((part.matrix_world @ Vector(corner)).z)
        if all_world_z:
            sc_obj.location.z = min(all_world_z)
        print(f"[blender_render] shadow catcher enabled, plane Z={sc_obj.location.z:.4f}")
    else:
        print(f"[blender_render] WARNING: shadow catcher object '{sc_obj_name}' not found in template")