# Written by yams_any4d agent, 2026-07-10. """ V4Head training visuals, ported from yam_local_train lib/viz/ (feature_pca.py, keypoints.py, bin_grids.py) so the wandb panels match the v4 trainer's: - viz PCA of per-frame DiT features (shared basis across frames) - GT (white open circles) vs pred (rainbow dots + polyline) keypoints over all frames - gripper / rotation bin grids (JET softmax rows, black outline = GT bin) All pure numpy + cv2; saved as pngs and returned as {key: path} for nice_videos. """ import cv2 import numpy as np def _pca_rgb(frames_feat: np.ndarray, out_wh: tuple) -> list: """(T, C, H, W) -> list of T HxW RGB uint8 PCA images, shared basis.""" (T, C, H, W) = frames_feat.shape flat = frames_feat.transpose(0, 2, 3, 1).reshape(-1, C).astype(np.float32) flat = flat - flat.mean(0, keepdims=True) # SVD on a subsample for speed if huge sub = flat[:: max(1, len(flat) // 20000)] _, _, Vt = np.linalg.svd(sub, full_matrices=False) pcs = flat @ Vt[:3].T pcs -= pcs.min(0, keepdims=True) pcs /= pcs.max(0, keepdims=True) + 1e-8 imgs = (pcs.reshape(T, H, W, 3) * 255).astype(np.uint8) return [cv2.resize(imgs[t], out_wh, interpolation=cv2.INTER_NEAREST) for t in range(T)] def pca_strip(frames_feat: np.ndarray, cell_w: int = 160, max_frames: int = 11) -> np.ndarray: """Horizontal strip of per-frame PCA panels.""" (T, C, H, W) = frames_feat.shape step = max(1, T // max_frames) sel = frames_feat[::step][:max_frames] cell_h = int(cell_w * H / W) cells = _pca_rgb(sel, (cell_w, cell_h)) pad = np.full((cell_h, 2, 3), 255, np.uint8) out = [] for c in cells: out.extend([c, pad]) return np.concatenate(out[:-1], axis=1) def keypoints_overlay(bg_bgr: np.ndarray, gt_uv: np.ndarray, pred_uv: np.ndarray) -> np.ndarray: """v4 style: GT white open circles (no line); pred rainbow filled dots + white polyline.""" img = bg_bgr.copy() T = len(gt_uv) gt_pts = np.round(gt_uv).astype(np.int32) for t in range(T - 1): # GT gets its own thin trajectory line for visibility cv2.line(img, tuple(gt_pts[t]), tuple(gt_pts[t + 1]), (200, 200, 200), 1, cv2.LINE_AA) for t in range(T): cv2.circle(img, tuple(gt_pts[t]), 9, (255, 255, 255), 2, lineType=cv2.LINE_AA) pts = np.round(pred_uv).astype(np.int32) for t in range(T - 1): cv2.line(img, tuple(pts[t]), tuple(pts[t + 1]), (255, 255, 255), 1, cv2.LINE_AA) for t in range(T): hue = int(170 * t / max(T - 1, 1)) color = cv2.cvtColor(np.uint8([[[hue, 255, 255]]]), cv2.COLOR_HSV2BGR)[0, 0] cv2.circle(img, tuple(pts[t]), 5, tuple(int(c) for c in color), -1, cv2.LINE_AA) return img def bin_grid_panel(logits: np.ndarray, target=None, cell_w: int = 8, cell_h: int = 8) -> np.ndarray: """(T, K) logits + (T,) GT -> JET grid, black outline on GT bin per row.""" (T, K) = logits.shape x = logits.astype(np.float32) x -= x.max(axis=1, keepdims=True) e = np.exp(x) e /= e.sum(axis=1, keepdims=True) n = e - e.min(axis=1, keepdims=True) n /= (n.max(axis=1, keepdims=True) - n.min(axis=1, keepdims=True) + 1e-8) color = cv2.applyColorMap((n * 255).astype(np.uint8), cv2.COLORMAP_JET) grid = cv2.resize(color, (K * cell_w, T * cell_h), interpolation=cv2.INTER_NEAREST) for t in range(T if target is not None else 0): k = int(target[t]) cv2.rectangle(grid, (k * cell_w, t * cell_h), ((k + 1) * cell_w - 1, (t + 1) * cell_h - 1), (0, 0, 0), 1) return grid def marginal_heatmap_grid(vol_tzpp: np.ndarray, bg_frames: list, cell: int = 160, alpha: float = 0.5, max_cells: int = 8) -> np.ndarray: """v4 viz/heatmap.py port: horizontal strip of T cells; each = bg frame blended with JET of the Z-marginal softmax (joint softmax over (Z,P,P), summed over Z).""" (T, Z, P, _) = vol_tzpp.shape step = max(1, T // max_cells) idxs = list(range(0, T, step))[:max_cells] cells = [] pad = np.full((cell, 2, 3), 255, np.uint8) for t in idxs: x = vol_tzpp[t].astype(np.float32).reshape(-1) e = np.exp(x - x.max()) pr = (e / e.sum()).reshape(Z, P, P).sum(0) # (P, P) # visibility: sharp single-voxel peaks vanish after resize+blend; a small # blur turns the true peak into a visible blob without moving it pr = cv2.GaussianBlur(pr, (3, 3), 0.7) pr = (pr - pr.min()) / (pr.max() - pr.min() + 1e-8) hm = cv2.applyColorMap( cv2.resize((pr * 255).astype(np.uint8), (cell, cell), interpolation=cv2.INTER_LINEAR), cv2.COLORMAP_JET) bg = cv2.resize(bg_frames[min(t, len(bg_frames) - 1)], (cell, cell), interpolation=cv2.INTER_AREA) cells.extend([cv2.addWeighted(hm, alpha, bg, 1 - alpha, 0), pad]) return np.concatenate(cells[:-1], axis=1) def make_v4head_visuals(head, head_out: dict, data_batch: dict, out_dir: str, train_step: int, batch_idx: int = 0) -> dict: """Build + save the three v4-style panels for one batch element. Returns {nice_videos_key: png_path}. Never raises (returns {} on failure).""" import os import torch try: # Rank-0 only: concurrent same-path writes from both ranks corrupt the pngs. if torch.distributed.is_available() and torch.distributed.is_initialized() \ and torch.distributed.get_rank() != 0: return {} os.makedirs(out_dir, exist_ok=True) b = batch_idx nice = {} # --- 1. PCA of scene DiT features over frames --- p0 = head_out['scene_feats_per_frame'][b].detach().float().cpu().numpy() # (Tl, C, Hp, Wp) strip = pca_strip(p0) p_pca = f'{out_dir}/v4h_pca_scene_{train_step:06d}.png' cv2.imwrite(p_pca, cv2.cvtColor(strip, cv2.COLOR_RGB2BGR)) nice['v4head/pca_scene'] = p_pca # --- 2. GT vs pred keypoints over all frames (scene view) --- vol = head_out['volume_logits'][b:b + 1].detach().float() vpos_b = head_out['voxel_positions'][b:b + 1].detach().float() pred_xyz = head.decode_pred_xyz(vol, vpos_b)[0] # (T, 3) gt_xyz = head_out['gt_xyz'][b] # (T, 3) hw = head_out['pixel_hw'] pred_uv = head.project_world_to_pix(pred_xyz, hw).cpu().numpy() gt_uv = head.project_world_to_pix(gt_xyz, hw).cpu().numpy() bg = None rgb0 = data_batch.get('a4d_raw', {}).get('rgb0', None) if isinstance(rgb0, torch.Tensor): # (B, C, T, H, W) in [-1,1]-ish fr = rgb0[b, :, 0].detach().float().cpu().numpy().transpose(1, 2, 0) fr = (fr - fr.min()) / (fr.max() - fr.min() + 1e-8) bg = cv2.cvtColor((fr * 255).astype(np.uint8), cv2.COLOR_RGB2BGR) bg = cv2.resize(bg, (int(hw[1]), int(hw[0]))) if bg is None: bg = np.full((int(hw[0]), int(hw[1]), 3), 80, np.uint8) # Render at 2x for marker visibility (v4head panels were unreadably small at 320px). bg = cv2.resize(bg, (bg.shape[1] * 2, bg.shape[0] * 2), interpolation=cv2.INTER_CUBIC) kp = keypoints_overlay(bg, gt_uv * 2.0, pred_uv * 2.0) p_kp = f'{out_dir}/v4h_kp_{train_step:06d}.png' cv2.imwrite(p_kp, kp) nice['v4head/kp_gt_vs_pred'] = p_kp # --- 2b. PCA strip of the UPSAMPLED + refined maps, ALL latent frames # (shared basis) — one cell per latent step, mirroring the per-timestep # volume instantiation. Falls back to latent-0 only if maps_all absent. --- if 'feat_maps_all' in head_out: fma = head_out['feat_maps_all'][b].detach().float().cpu().numpy() # (Tl,Fd,P,P) else: fma = head_out['feat_map'][b].detach().float().cpu().numpy()[None] up_pca = pca_strip(fma, cell_w=160, max_frames=fma.shape[0]) p_up = f'{out_dir}/v4h_pca_up_{train_step:06d}.png' cv2.imwrite(p_up, cv2.cvtColor(up_pca, cv2.COLOR_RGB2BGR)) nice['v4head/pca_upsampled'] = p_up # --- 2b2. wrist upsampled PCA + wrist-view kp panel --- if 'feat_map_wrist' in head_out: fmw = head_out['feat_map_wrist'][b].detach().float().cpu().numpy() wpca = _pca_rgb(fmw[None], (fmw.shape[2] * 5, fmw.shape[1] * 5))[0] p_wup = f'{out_dir}/v4h_pca_up_wrist_{train_step:06d}.png' cv2.imwrite(p_wup, cv2.cvtColor(wpca, cv2.COLOR_RGB2BGR)) nice['v4head/pca_upsampled_wrist'] = p_wup if 'T_w2c_wrist' in head_out: try: Kw = head_out['K_wrist_scaled'].detach().float() Tw = head_out['T_w2c_wrist'][b:b + 1].detach().float() uv_gt_w, _ = head._project(Kw, Tw, head_out['gt_xyz'][b:b + 1].float()) uv_pr_w, _ = head._project(Kw, Tw, pred_xyz[None].float()) rgb1 = data_batch.get('a4d_raw', {}).get('rgb1', None) if isinstance(rgb1, torch.Tensor): frw = rgb1[b, :, 0].detach().float().cpu().numpy().transpose(1, 2, 0) frw = (frw - frw.min()) / (frw.max() - frw.min() + 1e-8) bgw = cv2.cvtColor((frw * 255).astype(np.uint8), cv2.COLOR_RGB2BGR) bgw = cv2.resize(bgw, (int(hw[1]) * 2, int(hw[0]) * 2), interpolation=cv2.INTER_CUBIC) kpw = keypoints_overlay(bgw, uv_gt_w[0].cpu().numpy() * 2.0, uv_pr_w[0].cpu().numpy() * 2.0) p_kpw = f'{out_dir}/v4h_kp_wrist_{train_step:06d}.png' cv2.imwrite(p_kpw, kpw) nice['v4head/kp_gt_vs_pred_wrist'] = p_kpw except Exception as e: print(f'[v4head_viz] wrist kp panel failed: {e}') # --- 2c. v4-style marginal heatmap grid over timesteps --- vol_np = head_out['volume_logits'][b].detach().float().cpu().numpy() # (T, NZ, P, P) if getattr(head, 'use_wrist', True) and vol_np.shape[1] == 2 * getattr(head, 'n_z', vol_np.shape[1]): vol_np = vol_np[:, :vol_np.shape[1] // 2] # scene slabs only in 2-view mode # scene-only heads: FULL volume (old unconditional halving cut z>0.25m mass # and made trained heatmaps look like noise — Cameron caught it 2026-07-11) bg_frames = [] if isinstance(rgb0, torch.Tensor): Trgb = rgb0.shape[2] Tv = vol_np.shape[0] for t in range(Tv): fr_t = rgb0[b, :, min(int(round(t * (Trgb - 1) / max(Tv - 1, 1))), Trgb - 1)] fr_t = fr_t.detach().float().cpu().numpy().transpose(1, 2, 0) fr_t = (fr_t - fr_t.min()) / (fr_t.max() - fr_t.min() + 1e-8) bg_frames.append(cv2.cvtColor((fr_t * 255).astype(np.uint8), cv2.COLOR_RGB2BGR)) else: bg_frames = [np.full((64, 64, 3), 80, np.uint8)] * vol_np.shape[0] p_hm = f'{out_dir}/v4h_heatmap_grid_{train_step:06d}.png' cv2.imwrite(p_hm, marginal_heatmap_grid(vol_np, bg_frames)) nice['v4head/heatmap_grid'] = p_hm # --- 3. Grip / rot bin grids --- gl = head_out['grip_logits'][b].detach().float().cpu().numpy() rl = head_out['rot_logits'][b].detach().float().cpu().numpy() gt_g = head_out['tgt_grip'][b].cpu().numpy() gt_r = head_out['tgt_rot'][b].cpu().numpy() p_g = f'{out_dir}/v4h_grip_grid_{train_step:06d}.png' p_r = f'{out_dir}/v4h_rot_grid_{train_step:06d}.png' cv2.imwrite(p_g, bin_grid_panel(gl, gt_g)) cv2.imwrite(p_r, bin_grid_panel(rl, gt_r)) nice['v4head/grip_grid'] = p_g nice['v4head/rot_grid'] = p_r # Verify every png is PIL-readable (wandb.Image loads via PIL); drop any that isn't. from PIL import Image as _PILImage for k in list(nice.keys()): try: _im = _PILImage.open(nice[k]); _im.load() except Exception: print(f'[v4head_viz] WARNING: dropping unreadable panel {nice[k]}') del nice[k] return nice except Exception as e: # visuals must never kill training print(f'[v4head_viz] WARNING: visuals failed at step {train_step}: {e}') return {}