simready-forge 0.1.0__py3-none-any.whl

simready_forge/isaac_review.py

@@ -0,0 +1,420 @@
+"""Review a SimReady asset in Isaac Sim: RTX renders at rest and with every joint driven, plus PhysX checks.
+
+    OMNI_KIT_ACCEPT_EULA=YES <isaac sim python> isaac_review.py <asset.usda> <out_dir>
+
+Needs Isaac Sim's Python (the `isaacsim` pip package, or python.bat / python.sh of an install). Writes:
+  out_dir/review.png    contact sheet: 5 views at rest, 3 with every joint driven to 80% of the far side of its
+                        range, 2 with each two-sided joint driven to 80% of the other side
+  out_dir/view_*.png    the single renders
+  out_dir/review.json   settle (NaN, fell through, drift) and per-joint target vs reached, on each side
+Exit code 0 when the physics checks pass, 1 when one fails, 2 on an error.
+"""
+
+import json
+import math
+import sys
+import traceback
+from pathlib import Path
+
+ASSET = str(Path(sys.argv[1]).resolve()).replace("\\", "/")
+OUT = Path(sys.argv[2]).resolve()
+OUT.mkdir(parents=True, exist_ok=True)
+W, H = 640, 480
+
+from isaacsim import SimulationApp
+
+app = SimulationApp({"headless": True, "width": W, "height": H, "renderer": "RayTracedLighting", "anti_aliasing": 3})
+import carb
+import numpy as np
+import omni.replicator.core as rep
+import omni.timeline
+import omni.usd
+from PIL import Image, ImageDraw
+from pxr import Gf, Sdf, Usd, UsdGeom, UsdLux, UsdPhysics, UsdShade
+
+AXIS = {"X": 0, "Y": 1, "Z": 2}
+VIEWS = [("front", 270, 12), ("front_right", 310, 28), ("right", 0, 12), ("back_left", 135, 28), ("top", 280, 72)]
+
+
+def setup():
+    ctx = omni.usd.get_context()
+    ctx.new_stage()
+    for _ in range(5):
+        app.update()
+    stage = ctx.get_stage()
+    UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.z)
+    UsdGeom.SetStageMetersPerUnit(stage, 1.0)
+    UsdGeom.Xform.Define(stage, "/World")
+    UsdPhysics.Scene.Define(stage, "/World/physicsScene").CreateGravityMagnitudeAttr(9.81)
+    UsdLux.DomeLight.Define(stage, "/World/Dome").CreateIntensityAttr(450.0)
+    sun = UsdLux.DistantLight.Define(stage, "/World/Sun")
+    sun.CreateIntensityAttr(2200.0)
+    sun.CreateAngleAttr(2.0)
+    UsdGeom.Xformable(sun.GetPrim()).AddRotateXYZOp().Set(Gf.Vec3f(-48, 18, 30))
+    checker = OUT / "_ground.png"  # 10 cm grid, darker every metre
+    grid = np.full((2000, 2000, 3), 140, np.uint8)
+    for k in range(0, 2000, 100):  # 1 px = 1 mm: 10 cm lines 8 mm wide, so they survive a 640 px render
+        grid[k : k + 8, :] = grid[:, k : k + 8] = 112
+    for k in range(0, 2000, 1000):
+        grid[k : k + 16, :] = grid[:, k : k + 16] = 72
+    Image.fromarray(grid).save(checker)
+    ground = UsdGeom.Mesh.Define(stage, "/World/Ground")
+    h = 50.0
+    ground.CreatePointsAttr([(-h, -h, 0), (h, -h, 0), (h, h, 0), (-h, h, 0)])
+    ground.CreateFaceVertexCountsAttr([4])
+    ground.CreateFaceVertexIndicesAttr([0, 1, 2, 3])
+    UsdGeom.PrimvarsAPI(ground).CreatePrimvar("st", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.vertex).Set(
+        [(-h / 2, -h / 2), (h / 2, -h / 2), (h / 2, h / 2), (-h / 2, h / 2)]
+    )
+    UsdPhysics.CollisionAPI.Apply(ground.GetPrim())
+    mat = UsdShade.Material.Define(stage, "/World/Looks/Ground")
+    grip = UsdPhysics.MaterialAPI.Apply(mat.GetPrim())  # not `ground`: that is the mesh the material binds to
+    grip.CreateStaticFrictionAttr(0.8)
+    grip.CreateDynamicFrictionAttr(0.6)  # left at 0, sliding friction against the ground was half the asset's own
+    shader = UsdShade.Shader.Define(stage, "/World/Looks/Ground/Shader")
+    shader.CreateIdAttr("UsdPreviewSurface")
+    tex = UsdShade.Shader.Define(stage, "/World/Looks/Ground/Tex")
+    tex.CreateIdAttr("UsdUVTexture")
+    tex.CreateInput("file", Sdf.ValueTypeNames.Asset).Set(str(checker).replace("\\", "/"))
+    tex.CreateInput("wrapS", Sdf.ValueTypeNames.Token).Set("repeat")
+    tex.CreateInput("wrapT", Sdf.ValueTypeNames.Token).Set("repeat")
+    st = UsdShade.Shader.Define(stage, "/World/Looks/Ground/St")
+    st.CreateIdAttr("UsdPrimvarReader_float2")
+    st.CreateInput("varname", Sdf.ValueTypeNames.Token).Set("st")
+    tex.CreateInput("st", Sdf.ValueTypeNames.Float2).ConnectToSource(st.ConnectableAPI(), "result")
+    shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).ConnectToSource(tex.ConnectableAPI(), "rgb")
+    shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(0.85)
+    mat.CreateSurfaceOutput().ConnectToSource(shader.ConnectableAPI(), "surface")
+    UsdShade.MaterialBindingAPI.Apply(ground.GetPrim()).Bind(mat)
+    holder = UsdGeom.Xform.Define(stage, "/World/Asset")
+    holder.GetPrim().GetReferences().AddReference(ASSET)
+    box = (
+        UsdGeom.BBoxCache(Usd.TimeCode.Default(), [UsdGeom.Tokens.default_, UsdGeom.Tokens.render])
+        .ComputeWorldBound(holder.GetPrim())
+        .ComputeAlignedRange()
+    )
+    holder.AddTranslateOp().Set(Gf.Vec3d(0, 0, 0.001 - box.GetMin()[2]))  # lowest point on the ground
+    cam = UsdGeom.Camera.Define(stage, "/World/ReviewCam")
+    cam.CreateFocalLengthAttr(24.0)
+    cam.CreateHorizontalApertureAttr(20.955)
+    cam.CreateVerticalApertureAttr(20.955 * H / W)
+    cam.CreateClippingRangeAttr(Gf.Vec2f(0.005, 1000.0))
+    s = carb.settings.get_settings()
+    s.set("/rtx/post/tonemap/op", 4)
+    s.set("/rtx/post/aa/op", 3)
+    s.set("/rtx/directLighting/sampledLighting/enabled", True)
+    for _ in range(30):
+        app.update()
+    return stage, holder, UsdGeom.Xformable(cam.GetPrim()).AddTransformOp()
+
+
+def bounds(stage):
+    return (
+        UsdGeom.BBoxCache(Usd.TimeCode.Default(), [UsdGeom.Tokens.default_, UsdGeom.Tokens.render])
+        .ComputeWorldBound(stage.GetPrimAtPath("/World/Asset"))
+        .ComputeAlignedRange()
+    )
+
+
+def aim(cam_op, rng, azimuth, elevation):
+    lo, hi = rng.GetMin(), rng.GetMax()
+    centre = (lo + hi) / 2
+    radius = max((hi - lo).GetLength() / 2, 0.02)
+    dist = radius / math.sin(math.atan((20.955 * H / W) / 2 / 24.0)) * 1.05
+    a, e = math.radians(azimuth), math.radians(elevation)
+    eye = centre + Gf.Vec3d(math.cos(e) * math.cos(a), math.cos(e) * math.sin(a), math.sin(e)) * dist
+    cam_op.Set(Gf.Matrix4d().SetLookAt(eye, centre, Gf.Vec3d(0, 0, 1)).GetInverse())
+
+
+def capture(annot, name):
+    rep.orchestrator.step(rt_subframes=8, delta_time=0.0, pause_timeline=True)
+    image = Image.fromarray(np.asarray(annot.get_data())[:, :, :3])
+    image.save(OUT / f"{name}.png")
+    return image
+
+
+def frames(bodies):
+    cache = UsdGeom.XformCache()
+    out = {}
+    for b in bodies:
+        m = cache.GetLocalToWorldTransform(b)
+        out[str(b.GetPath())] = (np.array(m.ExtractRotationMatrix()).T, np.array(m.ExtractTranslation()))
+    return out
+
+
+def quat_matrix(q):
+    """A Gf quaternion as a 3x3 rotation acting on column vectors."""
+    w, (x, y, z) = q.GetReal(), q.GetImaginary()
+    return np.array(
+        [
+            [1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y)],
+            [2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x)],
+            [2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y)],
+        ]
+    )
+
+
+def subtree_load(stage, prim, joints, f):
+    """The mass a joint carries (its child body and everything jointed beyond it) and the lever arm from the joint's
+    anchor to that mass's centre, from the settled body poses."""
+    children = {}
+    for other in joints:
+        j = UsdPhysics.Joint(other)
+        children.setdefault(str(j.GetBody0Rel().GetTargets()[0]), []).append(str(j.GetBody1Rel().GetTargets()[0]))
+    joint = UsdPhysics.Joint(prim)
+    parent = str(joint.GetBody0Rel().GetTargets()[0])
+    todo, seen = [str(joint.GetBody1Rel().GetTargets()[0])], set()
+    while todo:
+        body = todo.pop()
+        if body not in seen and body in f:
+            seen.add(body)
+            todo += children.get(body, [])
+    masses = {}
+    for body in seen:
+        api = UsdPhysics.MassAPI(stage.GetPrimAtPath(body))
+        masses[body] = (api.GetMassAttr().Get() if api else None) or 0.1
+    total = sum(masses.values()) or 0.1
+    centre = sum(f[b][1] * m for b, m in masses.items()) / total
+    r0, t0 = f[parent] if parent in f else (np.eye(3), np.zeros(3))
+    pos0 = joint.GetLocalPos0Attr().Get()
+    anchor = r0 @ np.array(pos0 if pos0 is not None else (0, 0, 0), dtype=float) + t0
+    return total, max(0.02, float(np.linalg.norm(centre - anchor)))
+
+
+class JointGauge:
+    """A joint's position in its own coordinate, the one PhysX's targets and limits use: joint frame 1 against joint
+    frame 0 (each body's pose turned by its localRot), about or along the joint's axis. The rest pose need not be
+    position 0 (a lamp can rest at -37 degrees on its shoulder). A hinge's angle is followed in small steps from
+    the rest value, so travel past 180 degrees (a cap's thread) counts."""
+
+    def __init__(self, prim):
+        joint = UsdPhysics.Joint(prim)
+        self.b0 = str(joint.GetBody0Rel().GetTargets()[0])
+        self.b1 = str(joint.GetBody1Rel().GetTargets()[0])
+
+        def frame(attr, fallback):
+            value = attr.Get()
+            return fallback(value) if value is not None else None
+
+        self.rot0 = frame(joint.GetLocalRot0Attr(), quat_matrix)
+        self.rot1 = frame(joint.GetLocalRot1Attr(), quat_matrix)
+        self.rot0 = np.eye(3) if self.rot0 is None else self.rot0
+        self.rot1 = np.eye(3) if self.rot1 is None else self.rot1
+        self.pos0 = np.array(joint.GetLocalPos0Attr().Get() or (0, 0, 0), dtype=float)
+        self.pos1 = np.array(joint.GetLocalPos1Attr().Get() or (0, 0, 0), dtype=float)
+        self.axis = np.eye(3)[:, AXIS[prim.GetAttribute("physics:axis").Get()]]  # in joint frame 0
+        self.prismatic = prim.IsA(UsdPhysics.PrismaticJoint)
+        self.rest, self.value, self.last = None, None, None
+
+    def _relative(self, f):
+        """Joint frame 1 in joint frame 0: its rotation and, for a slide, its offset."""
+        (r0, t0), (r1, t1) = f[self.b0], f[self.b1]
+        frame0, frame1 = r0 @ self.rot0, r1 @ self.rot1
+        offset = frame0.T @ ((r1 @ self.pos1 + t1) - (r0 @ self.pos0 + t0))
+        return frame0.T @ frame1, offset
+
+    def _angle(self, rotation):
+        angle = math.degrees(math.acos(max(-1.0, min(1.0, (np.trace(rotation) - 1) / 2))))
+        w = np.array(
+            [rotation[2, 1] - rotation[1, 2], rotation[0, 2] - rotation[2, 0], rotation[1, 0] - rotation[0, 1]]
+        )
+        return angle if w @ self.axis >= 0 else -angle
+
+    def sample(self, f):
+        rotation, offset = self._relative(f)
+        if self.prismatic:
+            self.value = float(offset @ self.axis)
+        elif self.last is None:
+            self.value = self._angle(rotation)
+        else:
+            self.value += self._angle(rotation @ self.last.T)
+        self.last = rotation
+        if self.rest is None:
+            self.rest = self.value
+
+
+def verdict(joint):
+    return "FAIL" if not joint["passed"] else "short" if joint["short"] else "over" if joint.get("over") else "ok"
+
+
+def sheet(rest, posed, report):
+    tw, th = 320, 240
+    settle = report["settle"]
+    lines = [
+        f"settle: {'ok' if settle['passed'] else 'FAIL'} (drift {settle['drift_m']} m, moved {settle.get('moved_m')} m)"
+    ]
+    for k, v in report["joints"].items():
+        line = f"{k}: target {v['target']} reached {v['reached']} {verdict(v)}"
+        if "other_side" in v:
+            o = v["other_side"]
+            line += f"; other side {o['target']} reached {o['reached']} {verdict(o)}"
+        lines.append(line)
+    canvas = Image.new("RGB", (tw * 5, th * 2 + 30 + 16 * len(lines) + 12), (22, 24, 28))
+    draw = ImageDraw.Draw(canvas)
+    for i, (name, image) in enumerate(rest + posed):
+        x, y = (i % 5) * tw, (i // 5) * (th + 30)
+        canvas.paste(image.resize((tw, th)), (x, y))
+        draw.rectangle((x, y, x + 8 + 7 * len(name), y + 16), fill=(0, 0, 0))
+        draw.text((x + 4, y + 2), name, fill=(255, 255, 255))
+    draw.text((6, th + 8), "   |   ".join(lines)[:260], fill=(230, 230, 120))
+    for i, line in enumerate(lines):
+        draw.text((6, 2 * th + 36 + i * 16), line, fill=(230, 230, 230))
+    canvas.save(OUT / "review.png")
+
+
+def main():
+    stage, _holder, cam_op = setup()
+    rp = rep.create.render_product("/World/ReviewCam", (W, H))
+    annot = rep.AnnotatorRegistry.get_annotator("rgb")
+    annot.attach([rp])
+    rest_box = bounds(stage)
+    aim(cam_op, rest_box, *VIEWS[0][1:])
+    for _ in range(40):
+        app.update()
+    rest = []
+    for name, az, el in VIEWS:
+        aim(cam_op, rest_box, az, el)
+        rest.append((f"rest {name}", capture(annot, f"view_rest_{name}")))
+    bodies = [p for p in Usd.PrimRange(stage.GetPrimAtPath("/World/Asset")) if p.HasAPI(UsdPhysics.RigidBodyAPI)]
+    joints = [
+        p
+        for p in Usd.PrimRange(stage.GetPrimAtPath("/World/Asset"))
+        if p.IsA(UsdPhysics.RevoluteJoint) or p.IsA(UsdPhysics.PrismaticJoint)
+    ]
+    report = {"asset": ASSET, "bodies": len(bodies), "joints": {}}
+    start = frames(bodies)  # the modelled pose
+    timeline = omni.timeline.get_timeline_interface()
+    timeline.play()
+    for _ in range(90):
+        app.update()
+    mid = frames(bodies)
+    for _ in range(30):
+        app.update()
+    settled = frames(bodies)
+    pos = np.array([t for _, t in settled.values()]) if settled else np.zeros((1, 3))
+    drift = max((np.linalg.norm(settled[k][1] - mid[k][1]) for k in settled), default=0.0)
+    # how far any body ended from where it was modelled: a pile that rolls apart and comes to rest has no drift
+    moved = max((np.linalg.norm(settled[k][1] - start[k][1]) for k in settled if k in start), default=0.0)
+    report["settle"] = {
+        "nan": bool(not np.isfinite(pos).all()),
+        "fell_through": bool(pos[:, 2].min() < -0.5),
+        "drift_m": round(float(drift), 4),
+        "moved_m": round(float(moved), 4),
+        "most_moved": [  # which bodies: the three that ended furthest from their modelled pose
+            {"body": k.rsplit("/", 1)[-1], "moved_m": round(float(np.linalg.norm(settled[k][1] - start[k][1])), 4)}
+            for k in sorted(
+                (k for k in settled if k in start),
+                key=lambda k: -float(np.linalg.norm(settled[k][1] - start[k][1])),
+            )[:3]
+        ],
+    }
+    settle = report["settle"]
+    settle["stable"] = bool(not settle["nan"] and not settle["fell_through"])
+    # a joint chain that sags under gravity still moves after 1.5 s; a loose part that slid or rolled has moved
+    # further than settling explains (half a percent of the asset's size, 1 cm at least)
+    allowed = max(0.01, 0.005 * max(rest_box.GetSize()))
+    settle["moved_allowed_m"] = round(float(allowed), 4)
+    settle["holds_pose"] = bool(drift < 0.005 and moved < allowed)
+    settle["passed"] = settle["stable"] and settle["holds_pose"]
+    plans = {}
+    for prim in joints:
+        joint = (
+            UsdPhysics.RevoluteJoint(prim) if prim.IsA(UsdPhysics.RevoluteJoint) else UsdPhysics.PrismaticJoint(prim)
+        )
+        lo, hi = joint.GetLowerLimitAttr().Get(), joint.GetUpperLimitAttr().Get()
+        revolute = prim.IsA(UsdPhysics.RevoluteJoint)
+        unlimited = not (math.isfinite(lo) and math.isfinite(hi)) or max(abs(lo), abs(hi)) > 1e6
+        if unlimited and not revolute:
+            continue  # an unlimited slide has no range to drive into
+        if unlimited:
+            lo, hi = -90.0, 90.0  # a wheel or a free knob: turn it a quarter
+        far, near = (lo, hi) if abs(lo) > abs(hi) else (hi, lo)
+        sides = [0.8 * far] + ([0.8 * near] if abs(near) > 0.1 * (hi - lo) else [])  # a one-sided range: one side
+        # The drive is sized to the load it carries, as a real actuator is: strong enough to hold the subtree
+        # against gravity within about a third of a degree (or 2% of a slide's range), capped so it cannot
+        # overpower the joint's limits.
+        load, lever = subtree_load(stage, prim, joints, settled)
+        weight = load * 9.81
+        stiffness = 3.0 * weight * lever if revolute else weight / min(0.002, 0.02 * (hi - lo))  # per deg or per m
+        drive = UsdPhysics.DriveAPI.Apply(prim, "angular" if revolute else "linear")
+        drive.CreateStiffnessAttr().Set(stiffness)
+        drive.CreateDampingAttr().Set(stiffness * 0.15)
+        drive.CreateTargetPositionAttr().Set(float(sides[0]))
+        drive.CreateMaxForceAttr().Set(50.0 * weight * lever + 1.0 if revolute else 50.0 * weight + 10.0)
+        plans[prim] = (sides, lo, hi, drive)
+    gauges = {prim: JointGauge(prim) for prim in plans}
+    for gauge in gauges.values():
+        gauge.sample(settled)
+
+    def run(steps):
+        for _ in range(steps):
+            app.update()
+            now = frames(bodies)  # every step: a light, fast joint can turn 180 degrees in a few
+            for gauge in gauges.values():
+                gauge.sample(now)
+
+    def measured(prim, target, lo, hi):
+        reached = gauges[prim].value
+        miss = reached - target
+        if not gauges[prim].prismatic and abs(hi - lo - 360.0) < 0.1:  # one full turn: PhysX may go the short way
+            miss = (miss + 180.0) % 360.0 - 180.0
+        return {
+            "target": round(target, 4),
+            "reached": round(reached, 4),
+            "passed": bool(math.isfinite(reached) and abs(miss) <= 0.15 * (hi - lo)),
+            # within the pass band but well short: something stops it (a collision, the drive's cap): look at it
+            "short": bool(math.isfinite(reached) and abs(miss) > 0.05 * (hi - lo) and miss * target < 0),
+            # well past it: gravity on a light link, a weak hold
+            "over": bool(math.isfinite(reached) and abs(miss) > 0.05 * (hi - lo) and miss * target > 0),
+        }
+
+    run(180)
+    driven = frames(bodies)  # a blow-up while driving leaves review.json looking normal unless it is checked here
+    report["driven_nan"] = bool(not all(np.isfinite(t).all() and np.isfinite(r).all() for r, t in driven.values()))
+    for prim, (sides, lo, hi, _) in plans.items():
+        report["joints"][prim.GetName()] = {
+            "type": "revolute deg" if prim.IsA(UsdPhysics.RevoluteJoint) else "prismatic m",
+            "limits": [round(lo, 4), round(hi, 4)],
+            "rest": round(gauges[prim].rest, 4),
+            **measured(prim, sides[0], lo, hi),
+        }
+    posed = []
+    timeline.pause()
+    posed_box = bounds(stage)
+    for name, az, el in (VIEWS[1], VIEWS[0], VIEWS[4]):  # the top shows a cap or dial's turn
+        aim(cam_op, posed_box, az, el)
+        caption = "joints driven" if plans else "after 5 s"  # no joints: the same scene, simulated
+        posed.append((f"{caption} {name}", capture(annot, f"view_posed_{name}")))
+    two_sided = {prim: plan for prim, plan in plans.items() if len(plan[0]) > 1}
+    if two_sided:
+        for sides, _, _, drive in two_sided.values():
+            drive.GetTargetPositionAttr().Set(float(sides[1]))
+        timeline.play()
+        run(240)  # across the whole range, so a little longer
+        timeline.pause()
+        for prim, (sides, lo, hi, _) in two_sided.items():
+            report["joints"][prim.GetName()]["other_side"] = measured(prim, sides[1], lo, hi)
+        other_box = bounds(stage)
+        for name, az, el in (VIEWS[1], VIEWS[0]):
+            aim(cam_op, other_box, az, el)
+            posed.append((f"other side {name}", capture(annot, f"view_other_{name}")))
+    timeline.stop()
+    report["passed"] = bool(
+        report["settle"]["passed"]
+        and not report["driven_nan"]
+        and all(j["passed"] and j.get("other_side", {"passed": True})["passed"] for j in report["joints"].values())
+    )
+    sheet(rest, posed, report)
+    report["contact_sheet"] = str(OUT / "review.png")
+    (OUT / "review.json").write_text(json.dumps(report, indent=1), encoding="utf-8")
+    print("REVIEW " + json.dumps(report, indent=1))
+    return report["passed"]
+
+
+try:
+    code = 0 if main() else 1
+except Exception:  # noqa: BLE001  Kit swallows uncaught exceptions and exits 0
+    traceback.print_exc()
+    code = 2
+app.close()
+sys.exit(code)