HartOMat/render-worker/scripts/import_usd.py

"""USD import helper for Blender headless renders.

Runs inside Blender's Python environment (bpy available).
Imports a USD stage and restores seam + sharp edges from
schaeffler:*EdgeVertexPairs primvars (mapped as Blender mesh attributes
by Blender's built-in USD importer).

USD stage convention: mm Y-up, metersPerUnit=0.001.
Blender's USD importer respects metersPerUnit and scales objects to metres.
"""
import sys

import bpy      # type: ignore[import]
import bmesh    # type: ignore[import]
from mathutils import Vector  # type: ignore[import]


def import_usd_file(usd_path: str) -> list | tuple:
    """Import USD stage into current Blender scene.

    Returns a tuple of (parts, material_lookup) where:
    - parts: list of imported mesh objects, centred at world origin
    - material_lookup: dict mapping blender_object_name → canonical_material_name
      (populated from schaeffler:canonicalMaterialName primvars, empty dict if absent)

    USD stage is mm Y-up with metersPerUnit=0.001 — Blender scales to metres.
    """
    bpy.ops.object.select_all(action='DESELECT')
    bpy.ops.wm.usd_import(filepath=usd_path)
    parts = [o for o in bpy.context.selected_objects if o.type == 'MESH']

    if not parts:
        print(f"[import_usd] ERROR: No mesh objects imported from {usd_path}")
        sys.exit(1)

    print(f"[import_usd] imported {len(parts)} part(s) from USD: "
          f"{[p.name for p in parts[:5]]}", flush=True)

    _rename_usd_objects(parts)

    # Restore seam + sharp edges from primvars mapped to Blender mesh attributes.
    # Blender's USD importer converts Int2Array primvars to INT32_2D attributes.
    # Attribute names: "schaeffler:sharpEdgeVertexPairs" / "schaeffler:seamEdgeVertexPairs"
    restored = 0
    for part in parts:
        restored += _restore_seam_sharp(part)
    if restored:
        print(f"[import_usd] restored seam/sharp on {restored} mesh(es)", flush=True)

    # Extract material lookup via pxr direct read of the USD file.
    # Blender's USD importer does NOT expose STRING primvars or customData as
    # Python-accessible properties — but the pxr module (available in render-worker)
    # can read them perfectly from the same file.
    material_lookup: dict[str, str] = {}
    try:
        from pxr import Usd, UsdGeom  # type: ignore[import]
        stage = Usd.Stage.Open(usd_path)
        for prim in stage.Traverse():
            if prim.GetTypeName() != "Mesh":
                continue
            part_key = prim.GetCustomDataByKey("schaeffler:partKey") or ""
            mat_name = prim.GetCustomDataByKey("schaeffler:canonicalMaterialName") or ""
            if not part_key or not mat_name:
                # Also check parent Xform prim (metadata may be on container)
                parent = prim.GetParent()
                if parent:
                    part_key = part_key or (parent.GetCustomDataByKey("schaeffler:partKey") or "")
                    mat_name = mat_name or (parent.GetCustomDataByKey("schaeffler:canonicalMaterialName") or "")
            if part_key and mat_name:
                # Blender object name = mesh prim leaf name (part_key)
                material_lookup[part_key] = mat_name
    except Exception as exc:
        print(f"[import_usd] WARNING: pxr material lookup failed: {exc}", flush=True)

    if material_lookup:
        print(f"[import_usd] pxr material lookup: {len(material_lookup)}/{len(parts)} parts",
              flush=True)
    else:
        print("[import_usd] no schaeffler:canonicalMaterialName metadata found (legacy USD)",
              flush=True)

    # Centre combined bbox at world origin (same as import_glb convention)
    all_corners = []
    for p in parts:
        all_corners.extend(p.matrix_world @ Vector(c) for c in p.bound_box)
    if all_corners:
        mins = Vector((min(v.x for v in all_corners),
                       min(v.y for v in all_corners),
                       min(v.z for v in all_corners)))
        maxs = Vector((max(v.x for v in all_corners),
                       max(v.y for v in all_corners),
                       max(v.z for v in all_corners)))
        center = (mins + maxs) * 0.5
        all_imported = list(bpy.context.selected_objects)
        root_objects = [o for o in all_imported if o.parent is None]
        for obj in root_objects:
            obj.location -= center

    return parts, material_lookup


def _rename_usd_objects(parts: list) -> None:
    """No-op: mesh prims are now named after part_key in export_step_to_usd.py.

    Blender 5.0 collapses single-child Xform+Mesh into just the Mesh object,
    using the mesh prim leaf name as the Blender object name. Since the mesh
    prim is now named after the part_key (not "Mesh"), Blender imports each
    object with the correct part name, making material matching work directly.

    The object["usd:path"] custom property is NOT set by Blender's importer,
    so the previous path-based rename approach did not work.
    """
    print(f"[import_usd] mesh objects named from USD prim paths: {[p.name for p in parts[:3]]!r}",
          flush=True)


def _restore_seam_sharp(obj) -> int:
    """Apply seam+sharp edges from USD primvars mapped as Blender mesh attributes.

    Blender's USD importer maps primvars:schaeffler:sharpEdgeVertexPairs (Int2Array)
    to a mesh attribute named "schaeffler:sharpEdgeVertexPairs" with type INT32_2D.
    Each attribute element has a .value property returning a 2-tuple (v0, v1).

    Returns 1 if any edge data was applied, 0 otherwise.
    """
    mesh = obj.data
    sharp_attr = mesh.attributes.get("schaeffler:sharpEdgeVertexPairs")
    seam_attr  = mesh.attributes.get("schaeffler:seamEdgeVertexPairs")
    if not sharp_attr and not seam_attr:
        return 0

    # Ensure single-user data block before bmesh edit
    if mesh.users > 1:
        obj.data = mesh.copy()
        mesh = obj.data

    bm = bmesh.new()
    bm.from_mesh(mesh)
    bm.verts.ensure_lookup_table()

    n_verts = len(bm.verts)

    def _apply_pairs(attr, mark_fn):
        applied = 0
        for elem in attr.data:
            v = elem.value  # 2-tuple for INT32_2D
            if len(v) >= 2 and 0 <= v[0] < n_verts and 0 <= v[1] < n_verts:
                edge = bm.edges.get([bm.verts[v[0]], bm.verts[v[1]]])
                if edge:
                    mark_fn(edge)
                    applied += 1
        return applied

    n_sharp = n_seam = 0
    if sharp_attr:
        n_sharp = _apply_pairs(sharp_attr, lambda e: setattr(e, 'smooth', False))
    if seam_attr:
        n_seam = _apply_pairs(seam_attr, lambda e: setattr(e, 'seam', True))

    bm.to_mesh(mesh)
    bm.free()

    if n_sharp or n_seam:
        print(f"[import_usd] {obj.name}: {n_sharp} sharp edges, {n_seam} seam edges",
              flush=True)
    return 1