"""Render/localization core for the AR view (no servo/hardware deps, so it can be unit-tested against the sim alone). A Layer = one MuJoCo model + one ArUco board: it measures the board's world frame by render loop-closure, poses itself from joint angles, and renders from a localized camera pose.""" import os import xml.etree.ElementTree as ET import numpy as np import cv2 import mujoco from calib_utils import detect_boards, solve_pose, cam_mj_to_cv, RX180 CENTER_TICK = 2048 TICKS_TO_RAD = 2.0 * np.pi / 4096.0 def ticks_to_rad(tick, sign=1): return sign * (tick - CENTER_TICK) * TICKS_TO_RAD def mat2quat(R): q = np.zeros(4) mujoco.mju_mat2Quat(q, np.ascontiguousarray(R, float).ravel()) return q def look_at_mj(eye, target, up=(0, 0, 1)): """Rotation for a MuJoCo camera at `eye` looking at `target` (MJ cam views along -Z).""" eye = np.asarray(eye, float) fwd = np.asarray(target, float) - eye fwd /= np.linalg.norm(fwd) right = np.cross(fwd, np.asarray(up, float)) right /= np.linalg.norm(right) up2 = np.cross(right, fwd) return eye, np.column_stack([right, up2, -fwd]) def inject_extcam(model_path): """Copy the model with an extra 'extcam' worldbody camera we repose per frame, written next to the original so its meshes/textures still resolve.""" t = ET.parse(model_path) ET.SubElement(t.getroot().find("worldbody"), "camera", {"name": "extcam", "pos": "0 0 1", "xyaxes": "1 0 0 0 1 0", "fovy": "45"}) tmp = model_path.replace(".xml", "_arview.xml") t.write(tmp) return tmp def set_cam(model, cam_id, T_camCV_w, fovy_deg): """Point the extcam at an OpenCV camera world pose with a given vertical FOV.""" model.cam_pos[cam_id] = T_camCV_w[:3, 3] model.cam_quat[cam_id] = mat2quat(T_camCV_w[:3, :3] @ RX180) model.cam_fovy[cam_id] = fovy_deg def board_geom_pose(model, data, mat_name): """World (pos, 3x3 rotation) of the geom carrying the board texture material.""" mid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_MATERIAL, mat_name) for g in range(model.ngeom): if model.geom_matid[g] == mid: return np.array(data.geom_xpos[g]), np.array(data.geom_xmat[g]).reshape(3, 3) raise RuntimeError(f"no geom uses material '{mat_name}'") def measure_board_world(model, data, cam_id, board_spec, pos, xmat): """Loop-closure: look at the board along its normal, detect it, recover its GridBoard frame in the sim world. Tries both normal directions (the textured face may point either way).""" W = H = 1200 fovy = 50.0 f = H / (2 * np.tan(np.radians(fovy) / 2)) K = np.array([[f, 0, W / 2], [0, f, H / 2], [0, 0, 1]]) img = None for sgn in (1.0, -1.0): normal = xmat[:, 2] * sgn eye = pos + normal * 0.22 + xmat[:, 0] * 0.04 p, R_mj = look_at_mj(eye, pos, up=xmat[:, 1]) model.cam_pos[cam_id] = p model.cam_quat[cam_id] = mat2quat(R_mj) model.cam_fovy[cam_id] = fovy mujoco.mj_forward(model, data) T_camCV_w = cam_mj_to_cv(data.cam_xpos[cam_id], data.cam_xmat[cam_id]) with mujoco.Renderer(model, H, W) as ren: ren.update_scene(data, camera=cam_id) img = ren.render() dets = detect_boards(cv2.cvtColor(img, cv2.COLOR_RGB2GRAY), [board_spec]) if board_spec["name"] in dets: T_B_cam, err = solve_pose(*dets[board_spec["name"]], K) return T_camCV_w @ T_B_cam, err if img is not None: cv2.imwrite(f"_measure_debug_{board_spec['name']}.png", cv2.cvtColor(img, cv2.COLOR_RGB2BGR)) raise RuntimeError(f"could not detect '{board_spec['name']}' board in the sim render") def label(img, text, color=(0, 255, 0)): cv2.putText(img, text, (8, 22), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 3) cv2.putText(img, text, (8, 22), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 1) return img def blank(H, W, text=""): """A dark filler panel (keeps the grid aligned) with a dim caption.""" p = np.zeros((H, W, 3), np.uint8) if text: cv2.putText(p, text, (8, 22), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (110, 110, 110), 1) return p def blend_cell(layer, dets, K, real, fovy): """real|sim 50% blend for a layer if its board is observed, else a 'not observed' panel. `dets`/`K` are at the DETECTION resolution (full frame); `real` is the (smaller) display image and `fovy` the resolution-independent vertical FOV to render it at.""" H, W = real.shape[:2] d = dets.get(layer.board_name) if d is None or K is None: return blank(H, W, f"{layer.label}: not observed") T_B_cam, _ = solve_pose(*d, K) T_camCV_w = layer.T_board_w @ np.linalg.inv(T_B_cam) sim = layer.render(T_camCV_w, fovy, H, W) return label(cv2.addWeighted(real, 0.5, sim, 0.5, 0), f"{layer.label} render") class Layer: """One render source: a model localized from one ArUco board.""" def __init__(self, label_, model_path, board_spec, mat_name, map_joint): self.label = label_ self.board_spec = board_spec self.board_name = board_spec["name"] self.map_joint = map_joint self._tmp = inject_extcam(model_path) self.model = mujoco.MjModel.from_xml_path(self._tmp) self.data = mujoco.MjData(self.model) self.extcam = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_CAMERA, "extcam") mujoco.mj_forward(self.model, self.data) pos, xmat = board_geom_pose(self.model, self.data, mat_name) self.T_board_w, self.reproj = measure_board_world( self.model, self.data, self.extcam, board_spec, pos, xmat) self.posing = [] # (servo_id, qpos_adr) self.renderers = {} def build_posing(self, calib): """calib = {servo_id(str): {joint, sign, ...}}; map to this model's joints, keeping each servo's own sign so a shared id can't cross-contaminate layers.""" for sid, c in calib.items(): jn = self.map_joint(c["joint"]) if jn is None: continue jid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, jn) if jid >= 0: self.posing.append((int(sid), int(self.model.jnt_qposadr[jid]), int(c.get("sign", 1)))) def pose(self, ticks): """ticks = {servo_id: raw_tick}; sets each mapped joint from its own sign.""" for sid, adr, sgn in self.posing: t = ticks.get(sid) if t is not None: self.data.qpos[adr] = ticks_to_rad(t, sgn) def render(self, T_camCV_w, fovy, H, W): set_cam(self.model, self.extcam, T_camCV_w, fovy) mujoco.mj_forward(self.model, self.data) r = self.renderers.get((H, W)) if r is None: r = mujoco.Renderer(self.model, H, W) self.renderers[(H, W)] = r r.update_scene(self.data, camera=self.extcam) return cv2.cvtColor(r.render(), cv2.COLOR_RGB2BGR) def close(self): for r in self.renderers.values(): r.close() if os.path.exists(self._tmp): os.remove(self._tmp)