isaac/robozon_sorter/sim/plow_vision.py

"""Vision stack on top of the plow cell: infeed belt, item feeder, laser gate and the
CRE-ROI v2b decision that tells the pusher what to divert.

Purely additive. `sim/plow_cell.py` still owns the belts and the plow, and nothing here
edits the authored kinematics — the DiverterAnimGraph, the plow hinge and the pusher's own
drive are left exactly as `plow_cell.prepare()` leaves them.

The layout the scene augmentation produced:

    ConveyorTrack_05  x  0.00 .. +2.00   infeed, items are released at x=+1.70
    ConveyorTrack_02  x -2.00 ..  0.00   camera portal straddles x=-0.75
    ConveyorTrack_03  x -6.00 .. -2.00   laser gate at x=-3.74, pusher at x=-3.90
    Belt_01           branch to the bin

Goods run -X at the configured belt speed, so an item is released, measured under the
portal, and reaches the gate about 3.4 s later at 1 m/s.
"""
from __future__ import annotations

import json
from pathlib import Path

from pxr import Gf, PhysxSchema, UsdGeom, UsdPhysics, UsdShade

from .. import config as C
from . import plow_cell as _cell
from . import scene as _scene

# the conveyor added by scripts/add_vision_to_plow_cell.py
INFEED = "/World/ConveyorTrack_05/Belt"
INFEED_TRACK = "/World/ConveyorTrack_05"
INFEED_X0, INFEED_X1 = 0.0, 2.0

# release point: on the infeed belt, clear of its upstream edge so the item settles before
# it reaches the transfer to ConveyorTrack_02
SPAWN_X = 1.70

ITEMS_ROOT = _scene.ITEMS_ROOT
LASER_GATE = "/World/SortingRig/LaserGate"


def configure_infeed(stage, speed=None):
    """drive the added conveyor the same way as the rest of the line.

    Its local X is +X in world (unlike the branch, which is rotated), so the surface
    velocity is simply -speed on X.
    """
    speed = speed if speed is not None else C.BELT_SPEED
    prim = stage.GetPrimAtPath(INFEED)
    if not prim.IsValid():
        raise RuntimeError(
            f"{INFEED} missing - run scripts/add_vision_to_plow_cell.py first")

    if not prim.HasAPI(UsdPhysics.RigidBodyAPI):
        UsdPhysics.RigidBodyAPI.Apply(prim)
        UsdPhysics.RigidBodyAPI(prim).CreateKinematicEnabledAttr().Set(True)
    PhysxSchema.PhysxSurfaceVelocityAPI.Apply(prim)
    PhysxSchema.PhysxSurfaceVelocityAPI(prim).CreateSurfaceVelocityAttr().Set(
        Gf.Vec3f(-speed, 0.0, 0.0))

    grip = stage.GetPrimAtPath(_cell.GRIP_MATERIAL)
    if grip.IsValid():
        api = UsdShade.MaterialBindingAPI.Apply(prim)
        api.Bind(UsdShade.Material(grip),
                 bindingStrength=UsdShade.Tokens.strongerThanDescendants,
                 materialPurpose="physics")

    # the added track brings its own conveyor graph; it carries no speed and would only
    # fight the explicit surface velocity
    for suffix in ("", "_01"):
        g = stage.GetPrimAtPath(f"{INFEED_TRACK}/ConveyorBeltGraph{suffix}")
        if g.IsValid():
            g.SetActive(False)
    return prim


def load_items(stage, meshes_dir=None):
    """the bundled per-class test meshes, as dynamic rigid bodies parked off the line"""
    meshes_dir = Path(meshes_dir or C.MESHES)
    manifest = json.loads((meshes_dir / "manifest.json").read_text())
    UsdGeom.Xform.Define(stage, ITEMS_ROOT)
    items = {}
    for i, (name, meta) in enumerate(sorted(manifest.items())):
        usd = meshes_dir / f"{name}.usd"
        if not usd.exists():
            continue
        prim = UsdGeom.Xform.Define(stage, f"{ITEMS_ROOT}/{name}").GetPrim()
        refs = prim.GetReferences()
        refs.ClearReferences()
        refs.AddReference(str(usd))
        # Meshes flattened out of the working scene bring their own xformOp:translate at
        # float precision. ClearXformOpOrder() drops the *order*, not the attribute, so
        # adding a fresh double-precision op collides with what is already there and USD
        # raises. Match whatever precision the prim already carries.
        xf = UsdGeom.Xformable(prim)
        xf.ClearXformOpOrder()
        park = (9.0 + 1.2 * i, 5.0, 0.4)
        # Items exported from the working scene carry translate as float3, and
        # ClearXformOpOrder() drops the ORDER but keeps the attribute. AddTranslateOp() then
        # warns-as-raises about the precision mismatch (it still succeeds), and a retry hits
        # "already exists". Reuse the attribute that is there instead of adding anything.
        attr = prim.GetAttribute("xformOp:translate")
        if attr:
            op = UsdGeom.XformOp(attr)
            op.Set(Gf.Vec3f(*park) if str(attr.GetTypeName()) == "float3" else Gf.Vec3d(*park))
            xf.SetXformOpOrder([op])
        else:
            xf.AddTranslateOp().Set(Gf.Vec3d(*park))
        UsdPhysics.RigidBodyAPI.Apply(prim)
        # exported meshes arrive kinematic and hidden; both make them inert
        UsdPhysics.RigidBodyAPI(prim).CreateKinematicEnabledAttr().Set(False)
        UsdPhysics.MassAPI.Apply(prim).CreateMassAttr().Set(0.6)
        px = PhysxSchema.PhysxRigidBodyAPI.Apply(prim)
        px.CreateEnableCCDAttr().Set(True)
        px.CreateSolverPositionIterationCountAttr().Set(24)
        px.CreateSleepThresholdAttr().Set(0.0)      # a settled item must stay draggable
        # without this a blade sweeping into the item separates them at whatever speed
        # PhysX picks, which fires the item off the line instead of deflecting it
        px.CreateMaxDepenetrationVelocityAttr().Set(C.MAX_DEPENETRATION)
        UsdGeom.Imageable(prim).MakeVisible()
        items[name] = meta
    return items


def prepare(stage, belt_speed=None, script_control=True, meshes_dir=None):
    """plow_cell.prepare() plus the infeed belt, the items and the camera housekeeping"""
    info = _cell.prepare(stage, belt_speed=belt_speed, script_control=script_control)
    configure_infeed(stage, belt_speed)
    hidden = _scene.hide_aim_markers(stage)
    items = load_items(stage, meshes_dir)

    # mechanics.Cell releases at C.SPAWN_X; the plow cell's infeed is shorter than the
    # sorter's, so point it at this belt. run.py already sets C.BELT_SPEED the same way.
    C.SPAWN_X = SPAWN_X

    calib_path = C.CONFIG / "calib.json"
    info.update(items=items, aim_markers_hidden=hidden, spawn_x=SPAWN_X,
                calib=json.loads(calib_path.read_text()) if calib_path.exists() else None)
    return info


def load(usd_path=None, belt_speed=None, script_control=True, meshes_dir=None):
    stage = _cell.open_scene(usd_path)
    return stage, prepare(stage, belt_speed, script_control, meshes_dir)