isaac/scripts/run_flow.py

"""Поток товаров с шагом 700 мм: лазерная завеса определяет класс, пушер берёт класс D.

Луч теперь ДАТЧИК, а не преграда - коллизия снята прямо в файле. Раньше он перекрывал всю
ширину полотна на 19 мм над лентой, и товар вставал на x = -3.044, не доходя до пушера.

КАК ДАТЧИК ОПРЕДЕЛЯЕТ КРУГОВОЕ СЕЧЕНИЕ. Одиночный луч даёт только факт прохода. Завеса
из 181 луча поперёк ленты меряет ШИРИНУ товара, а по мере его проезда набирается профиль
ширины вдоль хода. У коробки он прямоугольный - ширина постоянна почти до конца; у тела
кругового сечения он дугообразный, ширина плавно нарастает и спадает.

Различаются они отношением средней ширины к наибольшей. Для прямоугольника оно стремится
к 1.0, для круга даёт площадь полукруга к описанному прямоугольнику, то есть pi/4 = 0.785.
Порог 0.90 разделяет их с запасом и не требует ни камеры, ни обучения - только геометрия.

Спавн идёт по времени, а не расстановкой заранее: при 1 м/с шаг 700 мм это ровно 0.70 с
между выпусками. Точка выпуска x = 1.90, а не C.SPAWN_X = 2.30 - в этой сборке подающая
секция ConveyorTrack_05 кончается на x = 2.001, и 2.30 висит в воздухе.
"""
import sys, math
REPO = "/home/dasha/robozon-sorter"
if REPO not in sys.path:
    sys.path.insert(0, REPO)

import omni.usd, omni.timeline
import isaacsim.core.experimental.utils.app as app_utils
from pxr import Gf, Usd, UsdGeom, UsdPhysics, PhysxSchema, UsdShade
from isaacsim.core.experimental.prims import RigidPrim
from omni.physx import get_physx_scene_query_interface

from robozon_sorter import config as C
from robozon_sorter.sim import plow_cell
from robozon_sorter.sim.mechanics import Cell
from robozon_sorter.sim.plow import Plow

SPEED = 1.0
PITCH = 0.70                     # м между товарами
RELEASE_X = 1.90
SCENE = f"{REPO}/scene/plow_cell_90_45_test.usd"
GATE_X = -3.20                   # где стоит луч
RAYS, Y0, Y1 = 31, -0.45, 0.45
CURTAIN_H = 0.40
ROUND_T = 0.90                   # средняя/наибольшая ширина: ниже - круглое сечение

# что пускаем: цилиндры - класс D (круговое сечение), коробки - не D
PLAN = [("D_cyl_1", "cyl"), ("box_1", "box"), ("D_cyl_2", "cyl"), ("box_2", "box"),
        ("D_cyl_3", "cyl"), ("box_3", "box"), ("D_cyl_4", "cyl"), ("box_4", "box")]

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

killed = [p.GetPath() for p in stage.Traverse()
          if "ConveyorBeltGraph" in p.GetName() or "DiverterAnimGraph" in p.GetName()]
for path in killed:
    stage.RemovePrim(path)
print(f"удалено узлов графов: {len(killed)} {[str(k).split(chr(47))[-1] for k in killed]}")
await app_utils.update_app_async(steps=10)
info = plow_cell.prepare(stage, belt_speed=SPEED, script_control=True, kinematic_arm=True)
for path, intent in (("/World/ConveyorTrack_05/Belt", (-1, 0, 0)),
                     ("/World/ConveyorTrack_06/Belt", (0, 1, 0))):
    pr = stage.GetPrimAtPath(path)
    if pr.IsValid():
        plow_cell.drive_belt(stage, path, intent, SPEED)
for path in list(plow_cell.BELTS) + [plow_cell.BRANCH,
                                     "/World/ConveyorTrack_05/Belt", "/World/ConveyorTrack_06/Belt"]:
    pr = stage.GetPrimAtPath(path)
    if pr.IsValid():
        PhysxSchema.PhysxSurfaceVelocityAPI(pr).CreateSurfaceVelocityEnabledAttr().Set(True)

bb = UsdGeom.BBoxCache(Usd.TimeCode.Default(), [UsdGeom.Tokens.default_, UsdGeom.Tokens.render])
TOP = bb.ComputeWorldBound(stage.GetPrimAtPath("/World/ConveyorTrack_05/Belt")
                           ).ComputeAlignedRange().GetMax()[2]
GRIP = stage.GetPrimAtPath(plow_cell.GRIP_MATERIAL)
print(f"сцена готова: лент {len(info['belts'])}+2, скорость {SPEED} м/с, "
      f"верх ленты z={TOP:.3f}, луч на x={GATE_X}")

ROOT = "/World/_Flow"
if stage.GetPrimAtPath(ROOT).IsValid():
    stage.RemovePrim(ROOT)
stage.DefinePrim(ROOT, "Xform")

def make(name, kind):
    path = f"{ROOT}/{name}"
    if kind == "cyl":
        g = UsdGeom.Cylinder.Define(stage, path)
        g.CreateRadiusAttr().Set(0.045); g.CreateHeightAttr().Set(0.10)
        g.CreateAxisAttr().Set("Z")
        half = 0.05
    else:
        g = UsdGeom.Cube.Define(stage, path); g.CreateSizeAttr().Set(2.0)
        half = 0.045
    xf = UsdGeom.Xformable(g.GetPrim())
    xf.AddTranslateOp().Set(Gf.Vec3d(RELEASE_X, 0.0, TOP + half + 0.006))
    if kind == "box":
        xf.AddScaleOp().Set(Gf.Vec3f(0.045, 0.045, 0.045))
    p = g.GetPrim()
    UsdPhysics.RigidBodyAPI.Apply(p); UsdPhysics.CollisionAPI.Apply(p)
    UsdPhysics.MassAPI.Apply(p).CreateMassAttr().Set(0.4)
    rb = PhysxSchema.PhysxRigidBodyAPI.Apply(p)
    rb.CreateEnableCCDAttr().Set(True); rb.CreateSolverPositionIterationCountAttr().Set(32)
    if GRIP.IsValid():
        UsdShade.MaterialBindingAPI.Apply(p).Bind(
            UsdShade.Material(GRIP), bindingStrength=UsdShade.Tokens.strongerThanDescendants,
            materialPurpose="physics")
    return path

query = get_physx_scene_query_interface()

# Положения тел ЧИТАЮТСЯ ЧЕРЕЗ RigidPrim, а не через BBoxCache. BBoxCache берёт авторские
# трансформы из слоя USD, а физика пишет состояние в Fabric - во время прогона эти два
# источника расходятся, и замер по BBoxCache показывает позы, которых на экране нет.
# Именно из-за этого предыдущий прогон отчитался, что все восемь товаров стоят ровно в
# точках выпуска, хотя таймлайн отработал полные 26 секунд.
_views = {}
def pos_of(path):
    v = _views.get(path)
    if v is None:
        v = RigidPrim(paths=[path]); _views[path] = v
    q = v.get_world_poses()[0].numpy()[0]
    return float(q[0]), float(q[1]), float(q[2])

def curtain_width():
    """ширина того, что сейчас под завесой, в мм. Луч, попавший выше полотна, - товар."""
    z0 = TOP + CURTAIN_H
    hit_y = []
    for i in range(RAYS):
        y = Y0 + (Y1 - Y0) * i / (RAYS - 1)
        h = query.raycast_closest((GATE_X, y, z0), (0.0, 0.0, -1.0), CURTAIN_H - 0.001)
        if h and h.get("hit"):
            zh = z0 - h["distance"]
            if zh > TOP + 0.006:          # выше полотна на 6 мм - значит товар
                hit_y.append(y)
    if not hit_y:
        return 0.0
    return (max(hit_y) - min(hit_y)) * 1000.0 + (Y1 - Y0) / (RAYS - 1) * 1000.0

# Таймлайн сцены кончается на своём endTimeCode и останавливается сам. В прошлом прогоне
# из-за этого прошло только 5.1 с вместо 26: головной товар едва дошёл до створа, и
# завеса не успела ничего измерить. Продлеваем ленту времени под длительность опыта.
fps = stage.GetTimeCodesPerSecond() or 24.0
print(f"таймлайн: {stage.GetStartTimeCode()}..{stage.GetEndTimeCode()} кадров при {fps} к/с "
      f"= {(stage.GetEndTimeCode()-stage.GetStartTimeCode())/fps:.1f} с - продлеваем")
stage.SetEndTimeCode(stage.GetStartTimeCode() + fps * 90.0)
tl.set_end_time(float(stage.GetEndTimeCode()) / fps)
tl.set_looping(False)

cell = Cell(stage, items={})
plow = Plow(stage); plow.target(C.PLOW_PRESET["D"])

tl.play()
await app_utils.update_app_async(steps=20)
t0 = float(tl.get_current_time())
released, profiles, done, order = [], {}, {}, []
next_release = 0.0
pusher_busy_until = -1.0
print(f"\nвыпуск каждые {PITCH/SPEED:.2f} с (шаг {PITCH*1000:.0f} мм при {SPEED} м/с)")
print("\n  событие")
print("  " + "-" * 76)

T_END = 26.0
import time as _wall
_it, _w0 = 0, _wall.time()
while float(tl.get_current_time()) - t0 < T_END:
    t = float(tl.get_current_time()) - t0
    _it += 1
    if _it % 40 == 0:
        print(f"   [цикл] итерация {_it}: сим t={t:5.2f}s, стена {_wall.time()-_w0:5.1f}s, "
              f"играет={tl.is_playing()}, товаров={len(released)}")
    if not tl.is_playing():
        # play() после stop() перематывает в начало и сбрасывает физику: в прошлом прогоне
        # это возвращало все товары в точки выпуска и обесценивало весь замер.
        print(f"   [цикл] ТАЙМЛАЙН ОСТАНОВИЛСЯ САМ на t={t:.2f}s - прерываю, "
              f"перезапуск обнулил бы опыт")
        break
    # выпуск потока
    if len(released) < len(PLAN) and t >= next_release:
        name, kind = PLAN[len(released)]
        path = make(name, kind)
        released.append((name, kind, path))
        order.append(name)
        next_release += PITCH / SPEED
        print(f"  {t:5.2f}s  выпущен {name} ({'цилиндр' if kind=='cyl' else 'коробка'})")
    # завеса: набрать профиль ширины для того, кто сейчас в створе
    w = curtain_width()
    if w > 5.0:
        # чей это профиль - ближайший по x к воротам
        best, bd = None, 1e9
        for name, kind, path in released:
            if not stage.GetPrimAtPath(path).IsValid():
                continue
            d = abs(pos_of(path)[0] - GATE_X)
            if d < bd:
                best, bd = name, d
        if best is not None and bd < 0.20:
            profiles.setdefault(best, []).append(w)
    # решение по классу, когда товар вышел из створа
    for name, kind, path in released:
        if name in done or name not in profiles:
            continue
        if not stage.GetPrimAtPath(path).IsValid():
            continue
        x = pos_of(path)[0]
        if x < GATE_X - 0.09 and len(profiles[name]) >= 3:
            prof = profiles[name]
            ratio = (sum(prof) / len(prof)) / max(prof)
            cls = "D" if ratio < ROUND_T else "B/C"
            done[name] = dict(gt=("D" if kind == "cyl" else "B/C"), pred=cls,
                              ratio=round(ratio, 3), wmax=round(max(prof)),
                              n=len(prof), pushed=False)
            print(f"  {t:5.2f}s  завеса: {name} ширина макс {max(prof):.0f} мм, "
                  f"проб {len(prof)}, ср/макс {ratio:.3f} -> класс {cls}")
    # пушер: взять класс D, когда он дошёл до ножа
    if t > pusher_busy_until:
        for name, kind, path in released:
            d = done.get(name)
            if not d or d["pred"] != "D" or d["pushed"]:
                continue
            if not stage.GetPrimAtPath(path).IsValid():
                continue
            x = pos_of(path)[0]
            if x <= C.PUSH_X + 0.06:
                print(f"  {t:5.2f}s  ПУШЕР берёт {name} на x={x:+.2f}")
                await cell.stroke(app_utils, out=True, speed=1.2)
                await cell.stroke(app_utils, out=False, speed=1.5)
                d["pushed"] = True
                pusher_busy_until = float(tl.get_current_time()) - t0 + 0.15
                break
    await app_utils.update_app_async(steps=2)

# итог
print("\n  ИТОГ")
print(f"  {'товар':10s} {'истина':7s} {'датчик':7s} {'ср/макс':>8s} {'шир,мм':>7s} "
      f"{'пушер':>6s} {'конец X,Y':>16s}  где")
print("  " + "-" * 84)
okc = okp = 0
for name, kind, path in released:
    d = done.get(name, dict(gt=("D" if kind == "cyl" else "B/C"), pred="-", ratio=0,
                            wmax=0, pushed=False))
    if stage.GetPrimAtPath(path).IsValid():
        x, y, z = pos_of(path)
    else:
        x = y = z = float("nan")
    where = ("ВЕТКА пушера" if y > 0.50 else
             "упал" if z < TOP - 0.20 else
             "+Y (угол)" if y > 0.10 else
             "-Y" if y < -0.10 else "прямо")
    if d["pred"] == d["gt"]:
        okc += 1
    if (d["gt"] == "D") == bool(d["pushed"]):
        okp += 1
    print(f"  {name:10s} {d['gt']:7s} {d['pred']:7s} {d['ratio']:8.3f} {d['wmax']:7.0f} "
          f"{'да' if d['pushed'] else 'нет':>6s} ({x:+6.2f},{y:+6.2f})  {where}")
print(f"\n  цикл: {_it} итераций, сим {float(tl.get_current_time())-t0:.2f}s, "
      f"стена {_wall.time()-_w0:.0f}s, играет={tl.is_playing()}")
tl.stop(); await app_utils.update_app_async(steps=5)
print(f"\n  класс определён верно: {okc}/{len(released)}   "
      f"пушер сработал по назначению: {okp}/{len(released)}")