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
hartomat:*EdgeVertexPairs primvars.  Blender's built-in USD importer does
NOT map arbitrary custom primvars (constant Int2Array) to mesh attributes,
so we read them directly via the pxr module and apply via bmesh.

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 hartomat:canonicalMaterialName customData, 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)

    # Read primvars + customData directly via pxr (Blender's USD importer does
    # NOT expose custom primvars or customData as Python-accessible properties).
    material_lookup: dict[str, str] = {}
    edge_data: dict[str, dict] = {}  # part_key → {sharp: [...], seam: [...]}
    try:
        from pxr import Usd, UsdGeom, Vt  # type: ignore[import]
        stage = Usd.Stage.Open(usd_path)
        for prim in stage.Traverse():
            if prim.GetTypeName() != "Mesh":
                continue
            part_key = prim.GetCustomDataByKey("hartomat:partKey") or ""
            mat_name = prim.GetCustomDataByKey("hartomat:canonicalMaterialName") or ""
            if not part_key or not mat_name:
                parent = prim.GetParent()
                if parent:
                    part_key = part_key or (parent.GetCustomDataByKey("hartomat:partKey") or "")
                    mat_name = mat_name or (parent.GetCustomDataByKey("hartomat:canonicalMaterialName") or "")
            if part_key and mat_name:
                material_lookup[part_key] = mat_name

            # Read seam/sharp primvars from USD mesh prim
            pvs_api = UsdGeom.PrimvarsAPI(prim)
            sharp_pv = pvs_api.GetPrimvar("hartomat:sharpEdgeVertexPairs")
            seam_pv = pvs_api.GetPrimvar("hartomat:seamEdgeVertexPairs")
            sharp_list = []
            seam_list = []
            if sharp_pv and sharp_pv.HasValue():
                raw = sharp_pv.Get()
                if raw is not None:
                    sharp_list = [(int(v[0]), int(v[1])) for v in raw]
            if seam_pv and seam_pv.HasValue():
                raw = seam_pv.Get()
                if raw is not None:
                    seam_list = [(int(v[0]), int(v[1])) for v in raw]
            if sharp_list or seam_list:
                # Use part_key as lookup key (matches Blender object name)
                edge_data[part_key] = {"sharp": sharp_list, "seam": seam_list}
    except Exception as exc:
        print(f"[import_usd] WARNING: pxr read 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 hartomat:canonicalMaterialName metadata found (legacy USD)",
              flush=True)

    if edge_data:
        print(f"[import_usd] pxr edge data: {len(edge_data)} parts with seam/sharp primvars",
              flush=True)

    # Apply seam + sharp edges to Blender meshes using pxr-read data
    restored = 0
    for part in parts:
        # Match Blender object name to part_key. Blender may add .001 suffixes
        # for duplicate names, so try exact match first, then strip suffix.
        obj_name = part.name
        data = edge_data.get(obj_name)
        if data is None:
            # Try stripping Blender's .NNN duplicate suffix
            base = obj_name.rsplit('.', 1)[0] if '.' in obj_name else obj_name
            data = edge_data.get(base)
        if data:
            restored += _restore_seam_sharp(part, data["sharp"], data["seam"])
    if restored:
        print(f"[import_usd] restored seam/sharp on {restored} mesh(es)", flush=True)
    else:
        print("[import_usd] no seam/sharp edges restored (no primvar data or no matches)",
              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, sharp_pairs: list, seam_pairs: list) -> int:
    """Apply seam+sharp edges from pxr-read primvar data via bmesh.

    sharp_pairs / seam_pairs: list of (v0, v1) vertex index pairs read
    from the USD file via pxr.UsdGeom.PrimvarsAPI.

    Returns 1 if any edge data was applied, 0 otherwise.
    """
    if not sharp_pairs and not seam_pairs:
        return 0

    mesh = obj.data

    # 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(pairs, mark_fn):
        applied = 0
        for v0, v1 in pairs:
            if 0 <= v0 < n_verts and 0 <= v1 < n_verts:
                edge = bm.edges.get([bm.verts[v0], bm.verts[v1]])
                if edge:
                    mark_fn(edge)
                    applied += 1
        return applied

    n_sharp = _apply_pairs(sharp_pairs, lambda e: setattr(e, 'smooth', False))
    n_seam = _apply_pairs(seam_pairs, 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 if (n_sharp or n_seam) else 0