isaac/robozon_sorter/sim/run.py

"""Entry point: build the cell, start the belt, run CRE-ROI v2b on every item as it passes
under the stand, and divert class D with the pusher.

    ./python.sh -m robozon_sorter.sim.run                 # windowed, watchable
    ./python.sh -m robozon_sorter.sim.run --headless      # batch, prints the log
    ./python.sh -m robozon_sorter.sim.run --no-vision     # mechanics only, uses ground truth

It also runs inside an already-open Isaac Sim: `from robozon_sorter.sim.run import main`.
"""
from __future__ import annotations

import argparse
import json
import sys


def parse_args(argv=None):
    p = argparse.ArgumentParser(description="Robozon conveyor sorting cell")
    p.add_argument("--headless", action="store_true", help="no window")
    p.add_argument("--no-vision", action="store_true",
                   help="skip CRE-ROI and route on ground truth (mechanics smoke test)")
    p.add_argument("--loops", type=int, default=1, help="passes over the test items")
    p.add_argument("--speed", type=float, default=None, help="override belt speed, m/s")
    p.add_argument("--pusher", type=float, default=None, help="override blade speed, m/s")
    p.add_argument("--log", default=None, help="write the run log here as JSON")
    return p.parse_args(argv)


async def _run(app_utils, stage, args):
    from .. import config as C
    from . import scene as S
    from .mechanics import Cell

    if args.speed:
        C.BELT_SPEED = args.speed
    built = S.build(stage)
    items = built["items"]
    print(f"cell built: {len(items)} test items "
          f"({sorted({m['zone'] for m in items.values()})})")

    vision = None
    if not args.no_vision:
        from ..cv.pipeline import CreRoiV2b
        vision = CreRoiV2b()
        vision.attach_cameras()
        print("CRE-ROI v2b ready; gate pixels:", vision.gate_px)

    await app_utils.update_app_async(steps=40)
    cell = Cell(stage, items.keys())
    cell.park_all()
    await app_utils.update_app_async(steps=15)

    import omni.timeline
    timeline = omni.timeline.get_timeline_interface()
    app_utils.play(commit=True)
    await app_utils.update_app_async(steps=20)

    order = [n for _ in range(args.loops) for n in sorted(items)]
    dt = 1.0 / 60.0
    log, active, done = [], [], set()
    classified, diverted = set(), set()
    nxt, t = 0, 0.0
    print(f"\n{'t':>7}  event")
    while t < 45.0 * args.loops * max(len(order), 1) / 6 and len(done) < len(order):
        await app_utils.update_app_async(steps=2)
        t += 2 * dt

        if nxt < len(order) and (not active or cell.pose(active[-1])[0] < C.SPAWN_X - C.RELEASE_GAP):
            name = order[nxt]
            cell.release(name)
            active.append(name)
            nxt += 1
            print(f"{t:7.2f}  release  {name}")

        for name in list(active):
            x = float(cell.pose(name)[0])

            if name not in classified and abs(x - C.CAM_X) < 0.06:
                gt = items[name]["zone"]
                if vision is not None:
                    was_playing = timeline.is_playing()
                    res = vision.measure()
                    # Replicator's step stops the timeline; resume or the line freezes
                    if was_playing and not timeline.is_playing():
                        timeline.play()
                        await app_utils.update_app_async(steps=2)
                    pred = res["cls"]
                    print(f"{t:7.2f}  vision   {name:18s} pred={pred} gt={gt} "
                          f"{'ok' if pred == gt else 'MISS'} dims={res['dims']} "
                          f"K={res['k']:.2f} views={res['views']} cre={res['cre_ms']}ms")
                    log.append(dict(item=name, gt=gt, **res))
                else:
                    pred = gt
                    print(f"{t:7.2f}  route    {name:18s} class={pred} (ground truth)")
                    log.append(dict(item=name, gt=gt, cls=pred))
                cell.pred = getattr(cell, "pred", {})
                cell.pred[name] = pred
                classified.add(name)

            if (name not in diverted and getattr(cell, "pred", {}).get(name) == "D"
                    and cell.laser() == name):
                took, held = await cell.divert(app_utils, name, speed=args.pusher)
                diverted.add(name)
                msg = f" (retract held {held:.2f}s)" if held > 0.01 else ""
                print(f"{t:7.2f}  divert   {name:18s} cycle {took:.2f}s{msg}")
                t += took

            place = cell.where(name)
            if place in ("bin", "line-end"):
                print(f"{t:7.2f}  done     {name:18s} -> {place}")
                for rec in log:
                    if rec["item"] == name and "outcome" not in rec:
                        rec["outcome"] = place
                active.remove(name)
                done.add(name)

    app_utils.stop()
    await app_utils.update_app_async(steps=15)
    cell.blade_to(C.BLADE_HOME_Y)

    graded = [r for r in log if "cls" in r and r["cls"] != "?"]
    hits = sum(1 for r in graded if r["cls"] == r["gt"])
    print(f"\n  {len(log)} items, {hits}/{len(graded)} agreed with ground truth")
    if vision is not None and graded:
        cre = [r["cre_ms"] for r in graded if r.get("cre_ms")]
        tot = [r["total_ms"] for r in graded if r.get("total_ms")]
        if cre:
            print(f"  CRE batched {sum(cre)/len(cre):.0f} ms/item, "
                  f"end-to-end {sum(tot)/len(tot):.0f} ms/item")
    routed = [r for r in log if r.get("outcome")]
    if routed:
        print("  routing: " + ", ".join(f"{r['item']}->{r['outcome']}" for r in routed))
    if args.log:
        with open(args.log, "w") as fh:
            json.dump(log, fh, indent=2)
        print(f"  log written to {args.log}")
    return log


def main(argv=None):
    args = parse_args(argv)
    try:
        import omni.usd
        stage = omni.usd.get_context().get_stage()
        inside = stage is not None
    except Exception:
        inside = False

    if not inside:
        from isaacsim import SimulationApp
        app = SimulationApp({"headless": args.headless, "width": 1600, "height": 900})
        import omni.usd
        import isaacsim.core.experimental.utils.stage as stage_utils
        stage_utils.create_new_stage()
        stage = omni.usd.get_context().get_stage()
    else:
        app = None

    import asyncio
    import isaacsim.core.experimental.utils.app as app_utils
    loop = asyncio.get_event_loop()
    try:
        return loop.run_until_complete(_run(app_utils, stage, args))
    finally:
        if app is not None:
            app.close()


if __name__ == "__main__":
    sys.exit(0 if main() is not None else 1)