isaac/control_test/flow_capture.py

"""STAGE 1 of the real-time flow bench (inside Isaac, no torch).

Items ride the real belt at PITCH spacing, exactly like run_pipeline.py. When one
crosses the inspection gate the timeline is paused, all six cameras are rendered, and the
run continues. Nothing about the item is written into the capture except its frames -
class and size are what stage 2 has to predict.

Meshes in items/ are 2.0-2.8x smaller than the catalogue dims they are labelled with, so
each is scaled up uniformly to catalogue scale first. Left small, every item fits the
450x320x320 envelope and class C becomes geometrically unreachable - the class metric
would be measuring nothing. Ground truth is the bbox actually in the scene after scaling.
"""
import asyncio, json, os, pathlib, sys, time
import numpy as np
import omni.timeline, omni.usd
import omni.kit.viewport.utility as vp
import isaacsim.core.experimental.utils.app as app_utils
from pxr import Gf, Usd, UsdGeom, UsdPhysics, PhysxSchema
from isaacsim.core.experimental.prims import RigidPrim

HERE = pathlib.Path("/home/dasha/robozon-sorter/control_test")
for extra in (str(HERE), "/home/dasha/robozon-sorter"):
    if extra not in sys.path:
        sys.path.insert(0, extra)
for _m in [k for k in list(sys.modules)
           if k.startswith(("robozon_sorter", "cell", "classify", "cam_configs"))]:
    del sys.modules[_m]
import importlib; importlib.invalidate_caches()

import cell
import cam_configs as CC
import classify as CL
from robozon_sorter import config as C

try:    PITCH = float(pitch)
except NameError: PITCH = 0.70
try:    SPEED = float(speed)
except NameError: SPEED = 1.0
try:    ITEMS = list(items)
except NameError:
    ITEMS = ["bag", "backpack", "lunchbox", "helmet", "pillow",
             "detergent", "bucket", "box_400x400x300", "box_300x200x200"]

GATE_X = float(CC.TARGET[0])          # inspection point, items travel -X past it
OUT = str(HERE / "captures" / "flow")
os.makedirs(OUT, exist_ok=True)
ROOT = "/World/FlowItems"

stage = omni.usd.get_context().get_stage()
tl = omni.timeline.get_timeline_interface()
if tl.is_playing():
    tl.stop(); await app_utils.update_app_async(steps=10)

info = await cell.prepare(stage, belt_speed=SPEED, script_control=True)
_gl = cell.add_ground_and_light(stage)
print(f"prepare: belts={len(info['belts'])} | пол {_gl['ground']}, купол {_gl['light']}")

calib = CC.apply_config(stage, CC.DEFAULT)
print(f"камеры {CC.DEFAULT}: " + ", ".join(
    f"{n.replace('_Left','')} h={c['height_mm']:.0f} d={c['standoff_mm']:.0f}"
    for n, c in calib.items() if n.endswith("_Left")))

# extra fill so the side views do not sit in shadow; the belt cell itself is unlit metal
for nm, pos in (("K0", (-0.75, 1.6, 2.6)), ("K1", (-0.75, -1.6, 2.6)),
                ("K2", (0.6, 0.0, 2.6)), ("K3", (-2.1, 0.0, 2.6))):
    p = f"/World/_CapLight_{nm}"
    if not stage.GetPrimAtPath(p).IsValid():
        from pxr import UsdLux
        sl = UsdLux.SphereLight.Define(stage, p)
        sl.CreateRadiusAttr().Set(0.25); sl.CreateIntensityAttr().Set(90000.0)
        UsdGeom.Xformable(sl.GetPrim()).AddTranslateOp().Set(Gf.Vec3d(*pos))

lib = {r["name"]: r for r in CL.load_library(str(HERE / "items_flow")) if "error" not in r}
ITEMS = [n for n in ITEMS if n in lib]
BB = lambda: UsdGeom.BBoxCache(Usd.TimeCode.Default(),
                               [UsdGeom.Tokens.default_, UsdGeom.Tokens.render])

UsdGeom.Xform.Define(stage, ROOT)
GT, ORDER = {}, []
for i, name in enumerate(ITEMS):
    path = f"{ROOT}/{name}"
    if stage.GetPrimAtPath(path).IsValid():
        stage.RemovePrim(path)
    # items_flow/ meshes are already catalogue-scale and already seated on z=0 (see
    # scale_items.py), so this is run_pipeline's proven load path verbatim: reference on
    # the body prim, one translate op, nothing nested.
    prim = UsdGeom.Xform.Define(stage, path).GetPrim()
    prim.GetReferences().ClearReferences()
    prim.GetReferences().AddReference(lib[name]["path"])
    xf = UsdGeom.Xformable(prim); xf.ClearXformOpOrder()
    xf.AddTranslateOp(precision=UsdGeom.XformOp.PrecisionDouble).Set(
        Gf.Vec3d(9.0 + 1.5 * i, 5.0, 0.4))
    UsdPhysics.RigidBodyAPI.Apply(prim)
    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.CreateSolverVelocityIterationCountAttr().Set(8)
    px.CreateSleepThresholdAttr().Set(0.0)
    px.CreateMaxDepenetrationVelocityAttr().Set(C.MAX_DEPENETRATION)
    for d in Usd.PrimRange(prim):
        if d.HasAPI(UsdPhysics.CollisionAPI):
            pc = PhysxSchema.PhysxCollisionAPI.Apply(d)
            pc.CreateContactOffsetAttr().Set(0.004)
            pc.CreateRestOffsetAttr().Set(0.001)
    await app_utils.update_app_async(steps=4)
    r = BB().ComputeWorldBound(prim).ComputeAlignedRange()
    dims = sorted([(r.GetMax()[k] - r.GetMin()[k]) * 1000.0 for k in range(3)], reverse=True)
    GT[name] = dict(dims_mm=[round(v, 1) for v in dims], k=lib[name]["k"],
                    zone_label=lib[name]["cls"],
                    zone_scene=CL.classify(dims, lib[name]["k"]))
    UsdGeom.Imageable(prim).MakeInvisible()
    ORDER.append(name)
    print(f"  {name:18s} {[round(v) for v in dims]} мм  метка {lib[name]['cls']}, "
          f"по геометрии сцены {GT[name]['zone_scene']}")

view = RigidPrim(paths=[f"{ROOT}/{n}" for n in ORDER])
plow_home = None
try:
    from robozon_sorter.sim.plow import Plow
    plow = Plow(stage, kinematic=True); plow.home()
except BaseException:
    pass

w = vp.get_active_viewport(); orig_cam = w.camera_path
CAMS = list(calib.keys())
manifest = dict(config=CC.DEFAULT, target=[float(v) for v in CC.TARGET],
                pitch_m=PITCH, speed_mps=SPEED, calib=calib, items={})

def activate(name):
    prim = stage.GetPrimAtPath(f"{ROOT}/{name}")
    for op in UsdGeom.Xformable(prim).GetOrderedXformOps():
        if op.GetOpType() == UsdGeom.XformOp.TypeTranslate:
            op.Set(Gf.Vec3d(cell.ENTRY_X, cell.ENTRY_Y, C.BELT_Z + 0.02)); break
    UsdGeom.Imageable(prim).MakeVisible()

async def shoot(name, x_at):
    tl.pause()
    for _ in range(3):
        await app_utils.update_app_async(steps=2)
    files = {}
    for cam in CAMS:
        w.camera_path = (f"/RigRS/{cam}" if stage.GetPrimAtPath(f"/RigRS/{cam}").IsValid()
                         else CC._cam(stage, cam).GetPath())
        await app_utils.update_app_async(steps=18); await asyncio.sleep(0)
        f = f"{OUT}/{name}__{cam}.png"
        vp.capture_viewport_to_file(w, file_path=f)
        await app_utils.update_app_async(steps=10); await asyncio.sleep(0)
        files[cam] = f
    manifest["items"][name] = dict(files=files, x_at=round(float(x_at), 4), gt=GT[name])
    tl.play()
    print(f"    снят {name} на x={x_at:+.3f}")

tl.play()
await app_utils.update_app_async(steps=10)
print(f"\n===== ПОТОК: {len(ORDER)} товаров, шаг {PITCH*1000:.0f} мм @ {SPEED} м/с "
      f"(интервал {PITCH/SPEED:.2f} с) =====")

released, shot, prev_x = [], set(), {}
t_next = float(tl.get_current_time())
idx = 0
t_end = t_next + (len(ORDER) + 1) * PITCH / SPEED + 25.0
while float(tl.get_current_time()) < t_end and len(shot) < len(ORDER):
    now = float(tl.get_current_time())
    if idx < len(ORDER) and now >= t_next:
        activate(ORDER[idx]); released.append(ORDER[idx])
        print(f"  {now - (t_end - (len(ORDER)+1)*PITCH/SPEED - 25.0):6.2f}с  выпущен {ORDER[idx]}")
        idx += 1; t_next = now + PITCH / SPEED
    try:
        pos = view.get_world_poses()[0].numpy()
    except BaseException:
        pos = None
    if pos is not None:
        for j, n in enumerate(ORDER):
            if n in shot or n not in released:
                continue
            x = float(pos[j][0])
            if prev_x.get(n, 9e9) > GATE_X >= x and \
                    abs(float(pos[j][2]) - C.BELT_Z) < 0.5:
                shot.add(n)
                await shoot(n, x)
            prev_x[n] = x
    await app_utils.update_app_async(steps=2)
    await asyncio.sleep(0)

tl.stop()
w.camera_path = orig_cam
await app_utils.update_app_async(steps=10)
json.dump(manifest, open(f"{OUT}/manifest.json", "w"), indent=1)
print(f"\nснято {len(manifest['items'])}/{len(ORDER)} -> {OUT}/manifest.json")
missed = [n for n in ORDER if n not in manifest["items"]]
if missed:
    print("не прошли ворота:", ", ".join(missed))