"""Rerun panel construction for ``deploy.py``. Two functions, one per model kind. Both call the same shared viz primitives that ``train.py`` and ``train_baseline.py`` use for their wandb panels, so the rerun viewer at inference matches the wandb viewer during training. Volume model (``DinoVolumeScene``) panels: - ``scene/{rgb,composite,mask_overlay}`` — live + mujoco overlay - ``pred/kp`` — rainbow predicted keypoints (no GT at inference) - ``pred/heatmap_grid`` — per-T marginal heatmap blended over RGB - ``pred/heatmap_grid_raw`` — same heatmap with no RGB - ``pred/pca_dino_mlp`` — PCA of refined per-pixel features - ``pred/{z,grip,rot}_grid`` — per-(T, bin) head probabilities - ``status`` — text panel Baseline (``DinoCLSXYZBaseline``) panels: - ``scene/{rgb,composite,mask_overlay}`` — same as above - ``pred/kp_projected`` — predicted (T, 3) world XYZ projected to image - ``pred/{grip,rot}_grid`` — same bin grids - ``status`` — text panel """ from __future__ import annotations import cv2 import numpy as np import rerun as rr import torch from .render import render_arm_with_exo from .rerun_helpers import to_jpeg from .viz.bin_grids import bin_grid_panel from .viz.feature_pca import feature_pca_panel, shared_feature_pca_panels from .viz.heatmap import marginal_heatmap_grid from .viz.keypoints import keypoints_overlay from .viz.overlay import compose_mask_over_live, compose_render_over_live from .viz.trajectory import project_world_to_image def _render_cam_check(joints7, T_cam_in_rbase, K_render, Wr, Hr, W, H): """Mujoco exo+arm render scaled back to the live (W, H). Returns ``(render_bgr, mask)`` ready for the overlay composers.""" rgb_lo, mask_lo = render_arm_with_exo(joints7, T_cam_in_rbase, K_render, Wr, Hr) bgr_lo = cv2.cvtColor(rgb_lo, cv2.COLOR_RGB2BGR) render_bgr = cv2.resize(bgr_lo, (W, H), interpolation=cv2.INTER_LINEAR) mask = cv2.resize(mask_lo, (W, H), interpolation=cv2.INTER_NEAREST) return render_bgr, mask def _log_scene_overlay(bgr, joints7, T_cam_in_rbase, K_render, Wr, Hr, W, H): """Log the three shared scene panels (rgb, composite, mask_overlay).""" render_bgr, mask = _render_cam_check( joints7, T_cam_in_rbase, K_render, Wr, Hr, W, H) rr.log("scene/rgb", to_jpeg(bgr)) rr.log("scene/composite", to_jpeg(compose_render_over_live(bgr, render_bgr, mask, alpha=0.5))) rr.log("scene/mask_overlay", to_jpeg(compose_mask_over_live(bgr, mask, color=(0, 0, 255), alpha=0.5))) def log_volume_panels(frame_idx, bgr, joints7, decoded, T_cam_in_rbase, K_render, Wr, Hr, W, H, img_size, input_bgr=None): """Panel set for the factorized-volume model — mirrors ``train.py``. Hot path optimisation: the joint softmax over ``(Z*P*P)`` is computed ONCE on the GPU (cheap), then we ship only the small (T,P,P) and (T,Z) marginals to the CPU. Avoids the previous triple-CPU-softmax over a 268 MB tensor that cost ~1.5 s per ``r``-press at T=32, Z=128, P=256. """ raw = decoded["raw_out"] vol_gpu = raw["volume_logits"][0] # (T, Z, P, P) on GPU T_dim, Z, P_, _ = vol_gpu.shape # GPU softmax — one cudaLaunch, ~5 ms even at T=32, Z=128, P=256. flat_gpu = vol_gpu.reshape(T_dim, -1) probs_flat = torch.softmax(flat_gpu, dim=1) probs_3d = probs_flat.reshape(T_dim, Z, P_, P_) probs_yx = probs_3d.sum(dim=1).cpu().numpy() # (T, P, P) ≈ 8 MB probs_z = probs_3d.sum(dim=(2, 3)).cpu().numpy() # (T, Z) ≈ 16 KB # Heatmap-grid helper takes 3D logits; pass log(probs_yx) so its # internal softmax round-trips back to probs_yx (cheap, ~10 ms now). log_probs_yx = np.log(probs_yx + 1e-12) feats = raw["F_refined"][0].cpu().numpy() # ``input_bgr`` is the image the model actually saw (= primary view for # multi-view runs). Falls back to ``bgr`` (= scene cam) for legacy # single-view models so the overlay lands on the right pixels. bg405 = cv2.resize(input_bgr if input_bgr is not None else bgr, (img_size, img_size), interpolation=cv2.INTER_AREA) rr.set_time("step", sequence=frame_idx) _log_scene_overlay(bgr, joints7, T_cam_in_rbase, K_render, Wr, Hr, W, H) # Predicted keypoints on the model-input view (no GT at inference). pix_uv = decoded["pix_uv"].astype(np.float32) kp = keypoints_overlay(bg405.copy(), pix_uv, draw_line=True, radius_px=6) rr.log("pred/kp", to_jpeg(kp)) # Heatmap grids (with + without RGB underneath). HEATMAP_CELL_PX = 256 rr.log("pred/heatmap_grid", to_jpeg(marginal_heatmap_grid(bg405, log_probs_yx, cell_size=HEATMAP_CELL_PX, alpha=0.5), target_w=None)) rr.log("pred/heatmap_grid_raw", to_jpeg(marginal_heatmap_grid(bg405, log_probs_yx, cell_size=HEATMAP_CELL_PX, alpha=1.0), target_w=None)) # PCA of refined features. pca = feature_pca_panel(feats.transpose(1, 2, 0)) rr.log("pred/pca_dino_mlp", to_jpeg(cv2.cvtColor(pca, cv2.COLOR_RGB2BGR))) # Per-(T, bin) grids — uses the GPU-computed marginals. grip_l = raw["grip_logits"][0].cpu().numpy() # (T, n_grip) rot_l = raw["rot_logits"][0].cpu().numpy() # (T, n_rot) rr.log("pred/z_grid", to_jpeg(bin_grid_panel(np.log(probs_z + 1e-12)))) rr.log("pred/grip_grid", to_jpeg(bin_grid_panel(grip_l))) rr.log("pred/rot_grid", to_jpeg(bin_grid_panel(rot_l))) rr.log("status", rr.TextDocument( f"step: {frame_idx}\n" f"model: volume (factorized YX × Z × T)\n" f"pred_pix[0]: ({pix_uv[0, 0]:.1f}, {pix_uv[0, 1]:.1f}) img-coord\n" f"height[0]: {decoded['height_m'][0]*1000:.0f}mm\n" f"grip[0]: {decoded['grip'][0]:.3f}\n" f"rot_bin[0]: {decoded['rot_bin'][0]}\n" )) def log_volume_panels_multiview( frame_idx, bgr_scene, bgrs_by_view, K_by_view, T_w2c_by_view, joints7, decoded, T_cam_in_rbase, K_render, Wr, Hr, W, H, img_size, views, ): """Multi-view sibling of ``log_volume_panels`` for the volumetric multi-view model. For every view we log the live image + predicted-XYZ keypoints projected through THAT view's K + T_w2c, plus a jointly-fit PCA panel (one basis fit across all views' feature maps, each painted with the shared basis — same as wandb during training). ``scene/{rgb,composite,mask_overlay}`` keep the scene-cam mujoco overlay regardless of which views the model consumes — it's the cam-check overlay, not a model-input panel. """ raw = decoded["raw_out"] xyz_pred = decoded["xyz_world"] # (T, 3) rr.set_time("step", sequence=frame_idx) _log_scene_overlay(bgr_scene, joints7, T_cam_in_rbase, K_render, Wr, Hr, W, H) # Jointly-fit PCA over every view's (B, F, P, P) → (P, P, F) feature map. view_feats_hwc = [v[0].detach().cpu().numpy().transpose(1, 2, 0) for v in raw["view_feat_maps"]] pca_imgs = shared_feature_pca_panels(view_feats_hwc, out_wh=(img_size, img_size)) # Per-view marginal heatmap from the JOINT softmax over the combined # (T, NZ, P, P) volume — each view's slab contributes a (T, P, P) # probability mass that lives in that view's image-pixel grid. vol_gpu = raw["volume_logits"][0] # (T, NZ, P, P) T_dim, NZ, P_, _ = vol_gpu.shape Z_per_view = NZ // max(1, len(views)) probs_full = torch.softmax(vol_gpu.reshape(T_dim, -1), dim=1) \ .reshape(T_dim, NZ, P_, P_) probs_yx_per_view = [] # each (T, P, P) for k in range(len(views)): slab = probs_full[:, k * Z_per_view:(k + 1) * Z_per_view] probs_yx_per_view.append(slab.sum(dim=1).cpu().numpy()) # Per-view: live RGB, kp overlay, PCA. for k, v_name in enumerate(views): bgr_v = bgrs_by_view.get(v_name) if bgr_v is None: continue K_v = K_by_view[v_name] T_w2c_v = T_w2c_by_view[v_name] bg_sq = cv2.resize(bgr_v, (img_size, img_size), interpolation=cv2.INTER_AREA) pix = [] for t in range(xyz_pred.shape[0]): p_w = xyz_pred[t].astype(np.float64) cam = T_w2c_v[:3, :3] @ p_w + T_w2c_v[:3, 3] z = max(float(cam[2]), 1e-6) uv = K_v @ cam pix.append([uv[0] / z, uv[1] / z]) pix = np.asarray(pix, dtype=np.float32) kp = keypoints_overlay(bg_sq.copy(), pix, draw_line=True, radius_px=6) rr.log(f"{v_name}/rgb", to_jpeg(bg_sq)) rr.log(f"pred/kp_{v_name}", to_jpeg(kp)) rr.log(f"pred/pca_{v_name}", to_jpeg(cv2.cvtColor(pca_imgs[k], cv2.COLOR_RGB2BGR))) # Marginal YX heatmap blended over this view's RGB + raw (no RGB). log_probs_yx = np.log(probs_yx_per_view[k] + 1e-12) rr.log(f"pred/heatmap_{v_name}", to_jpeg(marginal_heatmap_grid(bg_sq, log_probs_yx, cell_size=256, alpha=0.5), target_w=None)) rr.log(f"pred/heatmap_raw_{v_name}", to_jpeg(marginal_heatmap_grid(bg_sq, log_probs_yx, cell_size=256, alpha=1.0), target_w=None)) # Z marginal across the COMBINED volume (sums slabs from all views). vol = raw["volume_logits"][0] # (T, NZ, P, P) T_dim = int(vol.shape[0]) probs = torch.softmax(vol.reshape(T_dim, -1), dim=-1).reshape(vol.shape) probs_z = probs.sum(dim=(2, 3)).cpu().numpy() # (T, NZ) rr.log("pred/z_grid", to_jpeg(bin_grid_panel(np.log(probs_z + 1e-12)))) grip_l = raw["grip_logits"][0].cpu().numpy() rot_l = raw["rot_logits"][0].cpu().numpy() rr.log("pred/grip_grid", to_jpeg(bin_grid_panel(grip_l))) rr.log("pred/rot_grid", to_jpeg(bin_grid_panel(rot_l))) rr.log("status", rr.TextDocument( f"step: {frame_idx}\n" f"model: volumetric multi-view views={views}\n" f"xyz[0]: ({xyz_pred[0, 0]:.3f}, {xyz_pred[0, 1]:.3f}, " f"{xyz_pred[0, 2]:.3f}) m\n" f"grip[0]: {decoded['grip'][0]:.3f}\n" f"rot_bin[0]: {decoded['rot_bin'][0]}\n" )) def log_volume_panels_v3( frame_idx, bgr_scene, bgrs_by_view, K_by_view, T_w2c_by_view, joints7, decoded, T_cam_in_rbase, K_render, Wr, Hr, W, H, img_size, views, ): """Inference panel set for the v3 factorized model. Same layout as ``log_volume_panels_multiview`` but using ``v3_per_view_heatmaps`` (per-view log-marginal over (y, x)) and ``v3_view_dominance`` (per-view log-marginal over t) instead of indexing into a materialized ``volume_logits``. Per-timestep heatmaps would be too many panels — we render only t=0 here. """ from .model_volumetric_v3 import (v3_per_view_heatmaps_all_t, v3_view_dominance, v3_z_marginal) raw = decoded["raw_out"] xyz_pred = decoded["xyz_world"] # (T, 3) rr.set_time("step", sequence=frame_idx) _log_scene_overlay(bgr_scene, joints7, T_cam_in_rbase, K_render, Wr, Hr, W, H) view_feats_hwc = [v[0].detach().cpu().numpy().transpose(1, 2, 0) for v in raw["view_feat_maps"]] pca_imgs = shared_feature_pca_panels(view_feats_hwc, out_wh=(img_size, img_size)) log_p_per_view_all_t = [hp[0].detach().cpu().numpy() # (T, P, P) for hp in v3_per_view_heatmaps_all_t(raw)] dom = v3_view_dominance(raw)[0].detach().cpu().numpy() # (N, T) for k, v_name in enumerate(views): bgr_v = bgrs_by_view.get(v_name) if bgr_v is None: continue K_v = K_by_view[v_name] T_w2c_v = T_w2c_by_view[v_name] bg_sq = cv2.resize(bgr_v, (img_size, img_size), interpolation=cv2.INTER_AREA) pix = [] for t in range(xyz_pred.shape[0]): p_w = xyz_pred[t].astype(np.float64) cam = T_w2c_v[:3, :3] @ p_w + T_w2c_v[:3, 3] z = max(float(cam[2]), 1e-6) uv = K_v @ cam pix.append([uv[0] / z, uv[1] / z]) pix = np.asarray(pix, dtype=np.float32) kp = keypoints_overlay(bg_sq.copy(), pix, draw_line=True, radius_px=6) rr.log(f"{v_name}/rgb", to_jpeg(bg_sq)) rr.log(f"pred/kp_{v_name}", to_jpeg(kp)) rr.log(f"pred/pca_{v_name}", to_jpeg(cv2.cvtColor(pca_imgs[k], cv2.COLOR_RGB2BGR))) log_p_yx = log_p_per_view_all_t[k] # (T, P, P) rr.log(f"pred/heatmap_{v_name}", to_jpeg(marginal_heatmap_grid(bg_sq, log_p_yx, cell_size=256, alpha=0.5), target_w=None)) rr.log(f"pred/heatmap_raw_{v_name}", to_jpeg(marginal_heatmap_grid(bg_sq, log_p_yx, cell_size=256, alpha=1.0), target_w=None)) # Per-timestep view dominance bar: softmax over N → stacked column. dom -= dom.max(axis=0, keepdims=True) dom_p = np.exp(dom) dom_p /= np.maximum(dom_p.sum(axis=0, keepdims=True), 1e-9) T_ = dom_p.shape[1] col_w = 24; H_bar = 140 bar_w = col_w * T_ bar_img = np.full((H_bar, bar_w, 3), 30, dtype=np.uint8) palette = [(220, 200, 60), (200, 80, 220), (80, 220, 120), (220, 120, 80)] for t in range(T_): x0 = t * col_w + 2; x1 = (t + 1) * col_w - 2 y_cur = H_bar for k in range(dom_p.shape[0]): h_k = int(round(dom_p[k, t] * (H_bar - 4))) if h_k < 1: continue col = palette[k % len(palette)] cv2.rectangle(bar_img, (x0, y_cur - h_k), (x1, y_cur), col, -1) y_cur -= h_k leg = np.full((26, bar_w, 3), 30, dtype=np.uint8) x = 8 for k, v in enumerate(views[:dom_p.shape[0]]): col = palette[k % len(palette)] cv2.rectangle(leg, (x, 6), (x + 14, 20), col, -1) cv2.putText(leg, v, (x + 18, 18), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (240, 240, 240), 1, cv2.LINE_AA) x += 18 + 12 + 8 * len(v) rr.log("pred/view_dominance", to_jpeg(np.vstack([leg, bar_img]))) grip_l = raw["grip_logits"][0].cpu().numpy() rot_l = raw["rot_logits"][0].cpu().numpy() rr.log("pred/grip_grid", to_jpeg(bin_grid_panel(grip_l))) rr.log("pred/rot_grid", to_jpeg(bin_grid_panel(rot_l))) # Per-t marginal over Z bins. z_log = v3_z_marginal(raw)[0].detach().cpu().numpy() # (Z, T) rr.log("pred/z_grid", to_jpeg(bin_grid_panel(z_log.T))) rr.log("status", rr.TextDocument( f"step: {frame_idx}\n" f"model: volumetric v3 (factorized) views={views}\n" f"xyz[0]: ({xyz_pred[0, 0]:.3f}, {xyz_pred[0, 1]:.3f}, " f"{xyz_pred[0, 2]:.3f}) m\n" f"grip[0]: {decoded['grip'][0]:.3f}\n" f"rot_bin[0]: {decoded['rot_bin'][0]}\n" )) def log_baseline_panels(frame_idx, bgr, joints7, decoded, T_cam_in_rbase, K_in, K_render, Wr, Hr, W, H, img_size, T_world_to_cam): """Panel set for the CLS→XYZ baseline — mirrors ``train_baseline.py``. No volume / heatmap / PCA — the baseline doesn't have a spatial head. Instead we project the predicted (T, 3) world XYZ into the scene image with the same code train_baseline.py uses for its ``viz/kp_gt_vs_pred`` panel. """ raw = decoded["raw_out"] grip_l = raw["grip_logits"][0].cpu().numpy() rot_l = raw["rot_logits"][0].cpu().numpy() xyz = decoded["xyz_world"] # (T, 3) bg405 = cv2.resize(bgr, (img_size, img_size), interpolation=cv2.INTER_AREA) rr.set_time("step", sequence=frame_idx) _log_scene_overlay(bgr, joints7, T_cam_in_rbase, K_render, Wr, Hr, W, H) # Projected XYZ keypoints — exact same call as train_baseline.py. pred_pix, pred_in = project_world_to_image( xyz, K_in, T_world_to_cam, image_wh=(img_size, img_size)) kp = keypoints_overlay(bg405.copy(), pred_pix, in_frustum=pred_in, draw_line=True, radius_px=6) rr.log("pred/kp_projected", to_jpeg(kp)) rr.log("pred/grip_grid", to_jpeg(bin_grid_panel(grip_l))) rr.log("pred/rot_grid", to_jpeg(bin_grid_panel(rot_l))) rr.log("status", rr.TextDocument( f"step: {frame_idx}\n" f"model: baseline (CLS → global XYZ regression)\n" f"xyz[0]: ({xyz[0,0]:.3f}, {xyz[0,1]:.3f}, {xyz[0,2]:.3f}) m world\n" f"grip[0]: {decoded['grip'][0]:.3f}\n" f"rot_bin[0]: {decoded['rot_bin'][0]}\n" ))