From 8d61fc30bec8b3697eea0b62d8d0c956152f5c88 Mon Sep 17 00:00:00 2001 From: zjli2013 Date: Thu, 22 Oct 2026 21:51:24 +0800 Subject: [PATCH] feat(assets): a snapshot import places the robot at any heading in the room The only open floor beside a reconstructed living room's coffee table faces the robot toward room -Y; the importer assumed +X. Positions and bounds now go through one room-to-robot transform, and kept objects turn with it. Co-authored-by: Cursor --- robotsmith/assets/producers/aha3d_snapshot.py | 33 +++++++++++++------ tests/test_aha3d_snapshot.py | 12 +++++++ 2 files changed, 35 insertions(+), 10 deletions(-) diff --git a/robotsmith/assets/producers/aha3d_snapshot.py b/robotsmith/assets/producers/aha3d_snapshot.py index 7275424..6ac91d1 100644 --- a/robotsmith/assets/producers/aha3d_snapshot.py +++ b/robotsmith/assets/producers/aha3d_snapshot.py @@ -1,7 +1,7 @@ """AHa-3D scene export -> snapshot scenario. -The robot takes one spot in the room, facing +X; the furniture it works at becomes its table, -and the roots it can use become assets at their recorded positions: +The robot takes one spot in the room, facing ``--yaw`` degrees from room +X; the furniture it works +at becomes its table, and the roots it can use become assets at their recorded positions: python -m robotsmith.assets.producers.aha3d_snapshot --support furniture/desk \\ --base 0.62 -1.30 --prefix office48 --scenario scenarios/office48_desk.yaml --task task.yaml @@ -30,13 +30,25 @@ FLOOR_TOL_M = 0.05 -def plan(scene: dict, support: str, base_xy) -> dict: +def to_robot(xy, base_xy, yaw_deg: float = 0.0) -> np.ndarray: + """Room-frame xy point(s) in the robot frame: origin at the base, +X along the robot's heading.""" + c, s = np.cos(np.radians(yaw_deg)), np.sin(np.radians(yaw_deg)) + return (np.asarray(xy, dtype=float) - np.asarray(base_xy, dtype=float)) @ np.array([[c, -s], [s, c]]) + + +def box_to_robot(lo, hi, base_xy, yaw_deg: float = 0.0) -> tuple[np.ndarray, np.ndarray]: + """Robot-frame xy bounds of a room-frame box (bounds of its rotated corners).""" + corners = to_robot([[lo[0], lo[1]], [lo[0], hi[1]], [hi[0], lo[1]], [hi[0], hi[1]]], base_xy, yaw_deg) + return corners.min(axis=0), corners.max(axis=0) + + +def plan(scene: dict, support: str, base_xy, yaw_deg: float = 0.0) -> dict: """Each root's fate: ``table`` (replaced by the robot's table), ``support`` (rests on it), ``floor`` or ``dropped`` with a reason; positions in the robot frame.""" base = np.asarray(base_xy, dtype=float) roots = {r["instance_id"]: r for r in scene["roots"]} desk = roots[support] - rel = {i: (np.asarray(r["lo"][:2]) - base, np.asarray(r["hi"][:2]) - base) for i, r in roots.items() if "lo" in r} + rel = {i: box_to_robot(r["lo"], r["hi"], base, yaw_deg) for i, r in roots.items() if "lo" in r} out = {} for iid, r in roots.items(): if iid == support: @@ -59,21 +71,21 @@ def plan(scene: dict, support: str, base_xy) -> dict: out[iid] = ({"decision": "floor"} if np.linalg.norm(gap) <= FLOOR_REACH_M else {"decision": "dropped", "reason": f"more than {FLOOR_REACH_M} m from the robot"}) if "origin" in r: - out[iid]["anchor_xy"] = [round(float(v), 4) for v in np.asarray(r["origin"][:2]) - base] + out[iid]["anchor_xy"] = [round(float(v), 4) + 0.0 for v in to_robot(r["origin"][:2], base, yaw_deg)] size = table_size([rel[i] for i, d in out.items() if d["decision"] == "support"]) cx, cy = table_center_xy(DEFAULT_WORKSPACE_XY) t_lo, t_hi = np.array([cx - size[0] / 2, cy - size[1] / 2]), np.array([cx + size[0] / 2, cy + size[1] / 2]) for iid, d in out.items(): if d["decision"] == "floor" and np.all(rel[iid][0] < t_hi) and np.all(rel[iid][1] > t_lo): out[iid] = {"decision": "dropped", "reason": "overlaps the robot's table", "anchor_xy": d["anchor_xy"]} - return {"support": support, "base_xy_room": [float(v) for v in base], "yaw_deg": 0.0, + return {"support": support, "base_xy_room": [float(v) for v in base], "yaw_deg": float(yaw_deg), "table_xy": [t_lo.round(3).tolist(), t_hi.round(3).tolist()], "roots": out} def import_snapshot(export_dir: Path, *, support: str, base_xy, prefix: str, assets: Path, scenario: Path, - task: dict, provenance: dict) -> dict: + task: dict, provenance: dict, yaw_deg: float = 0.0) -> dict: scene = json.loads((export_dir / "scene.json").read_text(encoding="utf-8")) - mapping = plan(scene, support, base_xy) + mapping = plan(scene, support, base_xy, yaw_deg) roots = {r["instance_id"]: r for r in scene["roots"]} objects = [] for iid, d in mapping["roots"].items(): @@ -85,7 +97,7 @@ def import_snapshot(export_dir: Path, *, support: str, base_xy, prefix: str, ass description=f"{iid} from an AHa-3D scene ({r['semantic_class']})", provenance=provenance | {"instance_id": iid, "semantic_class": r["semantic_class"]}) objects.append({"id": r["item"], "asset": d["asset"], "on": d["decision"], "anchor_xy": d["anchor_xy"], - "yaw_deg": 0.0}) + "yaw_deg": 0.0 - yaw_deg}) doc = {"scenario": scenario.stem, "robot": "franka", "scene": {"support": "table_simple", "objects": objects}, "task": task} scenario.write_text(yaml.safe_dump(doc, sort_keys=False, default_flow_style=None), encoding="utf-8", newline="\n") @@ -98,6 +110,7 @@ def main() -> None: ap.add_argument("export_dir", type=Path) ap.add_argument("--support", required=True, help="instance_id of the furniture the robot works at") ap.add_argument("--base", type=float, nargs=2, required=True, help="robot base xy in the room frame") + ap.add_argument("--yaw", type=float, default=0.0, help="robot heading in the room frame, degrees from +X") ap.add_argument("--prefix", required=True) ap.add_argument("--assets", type=Path, default=Path("assets")) ap.add_argument("--scenario", type=Path, required=True) @@ -107,7 +120,7 @@ def main() -> None: prov = json.loads(a.provenance.read_text(encoding="utf-8")) if a.provenance else {} mapping = import_snapshot(a.export_dir, support=a.support, base_xy=a.base, prefix=a.prefix, assets=a.assets, scenario=a.scenario, task=yaml.safe_load(a.task.read_text(encoding="utf-8")), - provenance=prov) + provenance=prov, yaw_deg=a.yaw) kept = {i: d for i, d in mapping["roots"].items() if d["decision"] in ("support", "floor")} print(json.dumps(kept, indent=1)) diff --git a/tests/test_aha3d_snapshot.py b/tests/test_aha3d_snapshot.py index 0ab18a1..6aa5a42 100644 --- a/tests/test_aha3d_snapshot.py +++ b/tests/test_aha3d_snapshot.py @@ -43,3 +43,15 @@ def test_plan_sorts_every_root_and_keeps_room_positions(): def test_plan_expresses_positions_relative_to_the_robot_base(): roots = plan(SCENE, "furniture/desk", (0.5, 0.0))["roots"] assert roots["decor/cup"]["anchor_xy"] == [0.74, 0.14] + + +def test_plan_follows_the_robot_heading(): + def turn(r): + if "lo" not in r: + return r + (x0, y0), (x1, y1) = r["lo"][:2], r["hi"][:2] + return _root(r["instance_id"], r["semantic_class"], [-y1, x0, r["lo"][2]], [-y0, x1, r["hi"][2]]) + + turned = plan({"roots": [turn(r) for r in SCENE["roots"]]}, "furniture/desk", (0.0, 0.0), yaw_deg=90.0) + assert turned["roots"] == plan(SCENE, "furniture/desk", (0.0, 0.0))["roots"] + assert turned["yaw_deg"] == 90.0