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 bmesh  # type: ignore[import]
    import mathutils  # type: ignore[import]

    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))

    marked_total = 0
    for obj in parts:
        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()

        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 = 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

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

    print(f"[blender_render] OCC sharp edges applied: {marked_total} edges marked across {len(parts)} parts", 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")