"""Reusable per-ckpt viz panels — kp + height + rot + grip + DINO PCA. One entry point :func:`render_run_panels` consumes a ``viz_inputs`` dict (populated by :mod:`code.render_run_viz`) and returns ONE stacked BGR PNG (numpy uint8). Same code path for every model variant — only the caller dispatches on ``model_kind`` to fill ``viz_inputs``. viz_inputs schema:: { "run_name": str, "model_kind": str, "step": int, "img_size": int, "scene_bgr": (H, W, 3) uint8 — live scene-cam frame at img_size "K_in_scene": (3, 3) float — scene cam intrinsics at img_size "T_w2c_scene": (4, 4) float — world → scene-cam "pred_xyz": (T, 3) float — predicted EE WORLD XYZ "gt_xyz": (T, 3) float — GT EE WORLD XYZ "pred_z_bins": (T,) int or None — discrete height bin (volume models) "gt_z_bins": (T,) int or None "z_lo": float "z_hi": float "n_z": int — number of height bins "grip_logits": (T, n_grip) float "rot_logits": (T, n_rot) float "gt_grip": (T,) int "gt_rot": (T,) int "dino_patches": (D, P, P) float or None — per-pixel DINO features for the PCA panel } """ from __future__ import annotations import io from typing import Any import cv2 import numpy as np import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt from .viz.keypoints import keypoints_overlay from .viz.trajectory import project_world_to_image from .viz.feature_pca import feature_pca_panel from .viz.height_curve import height_curve_panel from .viz.bin_grids import bin_grid_panel # ────────────────────────────────────────────────────────────────────── # Panel builders (uniform 6-row stack) # ────────────────────────────────────────────────────────────────────── def _panel_kp(v: dict) -> np.ndarray: bgr = v["scene_bgr"] H, W = bgr.shape[:2] pred_pix, pred_in = project_world_to_image( v["pred_xyz"], v["K_in_scene"], v["T_w2c_scene"], image_wh=(W, H)) gt_pix, _ = project_world_to_image( v["gt_xyz"], v["K_in_scene"], v["T_w2c_scene"], image_wh=(W, H)) # GT in WHITE outline, pred filled CYAN with trajectory line. out = keypoints_overlay(bgr.copy(), gt_pix, filled=False, draw_line=False, color_override=(255, 255, 255), radius_px=10) out = keypoints_overlay(out, pred_pix, in_frustum=pred_in, draw_line=True, radius_px=6) bar = _label_bar(W, "keypoints · white = GT · rainbow = pred chunk") return np.vstack([bar, out]) def _panel_height(v: dict, out_wh=(960, 280)) -> np.ndarray: """Pred-vs-GT z over T. For volume models we have discrete bins; for regression models we feed continuous z directly by mapping to fake "bins" so the existing height_curve_panel keeps working.""" z_lo, z_hi, n_z = v["z_lo"], v["z_hi"], v["n_z"] if v["pred_z_bins"] is not None and v["gt_z_bins"] is not None: panel = height_curve_panel( np.asarray(v["pred_z_bins"]), np.asarray(v["gt_z_bins"]), z_lo, z_hi, n_z, out_wh=out_wh) else: # Continuous z: bin manually with the same spacing. def _to_bin(z): bw = (z_hi - z_lo) / max(1, n_z) return np.clip(((np.asarray(z) - z_lo) / bw).astype(int), 0, n_z - 1) pred_b = _to_bin(v["pred_xyz"][:, 2]) gt_b = _to_bin(v["gt_xyz"][:, 2]) panel = height_curve_panel(pred_b, gt_b, z_lo, z_hi, n_z, out_wh=out_wh) bar = _label_bar(panel.shape[1], "height · pred vs GT (m)") return np.vstack([bar, panel]) def _panel_grip(v: dict) -> np.ndarray: grid = bin_grid_panel(v["grip_logits"], target=v["gt_grip"]) bar = _label_bar(grid.shape[1], "grip bins · column = T · row = bin") return np.vstack([bar, grid]) def _panel_rot(v: dict) -> np.ndarray: grid = bin_grid_panel(v["rot_logits"], target=v["gt_rot"]) bar = _label_bar(grid.shape[1], "rotation bins · column = T · row = bin") return np.vstack([bar, grid]) def _panel_pca(v: dict) -> np.ndarray | None: if v.get("dino_patches") is None: return None # feature_pca_panel takes (P, P, C); the dump may give (C, P, P) or # (P, P, C). Accept both. feats = np.asarray(v["dino_patches"]) if feats.ndim == 3 and feats.shape[0] < feats.shape[-1]: feats = feats.transpose(1, 2, 0) pca_rgb = feature_pca_panel(feats, out_wh=(512, 512)) pca_bgr = cv2.cvtColor(pca_rgb, cv2.COLOR_RGB2BGR) bar = _label_bar(pca_bgr.shape[1], "DINO patch features · PCA (RGB = top-3 PCs)") return np.vstack([bar, pca_bgr]) def _label_bar(width: int, text: str) -> np.ndarray: bar = np.full((34, width, 3), 22, dtype=np.uint8) cv2.putText(bar, text, (10, 23), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (220, 220, 220), 1, cv2.LINE_AA) return bar def _resize_to_width(img: np.ndarray, w: int) -> np.ndarray: h, w0 = img.shape[:2] new_h = int(round(h * w / w0)) return cv2.resize(img, (w, new_h), interpolation=cv2.INTER_AREA) # ────────────────────────────────────────────────────────────────────── # Top-level # ────────────────────────────────────────────────────────────────────── def render_run_panels(viz_inputs: dict) -> np.ndarray: """Returns one BGR (H, W, 3) uint8 panel stack. All sub-panels are rescaled so their width matches the keypoint image (= the run's ``img_size``), so the whole stack reads as a uniform column.""" panels = [] # Use the keypoint image's native width as the reference. p_kp = _panel_kp(viz_inputs) target_w = p_kp.shape[1] panels.append(p_kp) p_h = _panel_height(viz_inputs, out_wh=(target_w, max(220, target_w // 2))) panels.append(_resize_to_width(p_h, target_w)) p_grip = _panel_grip(viz_inputs) panels.append(_resize_to_width(p_grip, target_w)) p_rot = _panel_rot(viz_inputs) panels.append(_resize_to_width(p_rot, target_w)) p_pca = _panel_pca(viz_inputs) if p_pca is not None: panels.append(_resize_to_width(p_pca, target_w)) # Top header. header = np.full((60, target_w, 3), 14, dtype=np.uint8) cv2.putText( header, f"{viz_inputs['run_name']} · {viz_inputs['model_kind']} · " f"step {viz_inputs.get('step', '?')}", (12, 38), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (235, 235, 255), 2, cv2.LINE_AA) return np.vstack([header] + panels)