# Created by BVH, Jul - Aug 2025. # Collects metrics and creates visualizations (with automatic sorting between input / prediction / ground truth), for validation / testing, and optionally also during training. import os import time import gc from datetime import datetime import imageio import lovely_tensors import matplotlib import matplotlib.pyplot as plt import numpy as np import torch from einops import rearrange from lovely_numpy import lo from rich import print from imaginaire.callbacks.manual_gc import trigger_manual_gc from imaginaire.utils import log, misc lovely_tensors.monkey_patch() np.set_printoptions(precision=3, suppress=True) torch.set_printoptions(precision=3, sci_mode=False, threshold=1000) from custom.any4d.a4d_logistics import detach_dict_recursive from custom.utils.utils import resilient_subroutine from custom.utils.io import save_json from custom.utils.logging import to_json_safe from custom.eval.visuals import ( cached_decode_latent_lowdim_auto, cached_decode_latent_video_auto, cached_decode_viz_latent_video_auto, flex_save_video, get_per_pixel_frame_mask_means, make_gallery, simple_horizontal_concat, viz_custom, paint_label, paint_mask_indicator, tensor_to_thwc_uint8, unpack_nested, sorted_camslast, ) import re from custom.eval.wm_visuals import draw_actions, get_traj_bounds, draw_traj, viz_anydrive_gaia from anydata.geometry.camera import Camera def paint_actions( cur_pred_viz, cur_gt_viz, k, has_output, has_supervise, output_dict, data_batch, meta_dl_batch, a4d_vae, batch_idx, cams, cam_idx, ): ''' Paint (2D-projected) input + predicted + target actions onto cur_pred_viz / cur_gt_viz in place (via slice assignment). Called once per rgb entry when proprio/action exists. ''' start_time = time.time() # low pri TODO(bvh): prefer lowdim stuff (action/traj/etc) to stay on CPU instead if 'proprio' in output_dict['x0_entries'] and 'proprio' in data_batch['a4d_latent']: proprio_in = cached_decode_latent_lowdim_auto(output_dict['x0_entries'], 'proprio', a4d_vae)[batch_idx] # (T, C). proprio_in_mask = data_batch['a4d_latent']['proprio_input_mask'][batch_idx, :, 0].to(proprio_in.device) # (T). else: proprio_in = None if 'action' in output_dict['x0_entries'] and 'action' in data_batch['a4d_latent']: action_in = cached_decode_latent_lowdim_auto(output_dict['x0_entries'], 'action', a4d_vae)[batch_idx] # (T, C). action_in_mask = data_batch['a4d_latent']['action_input_mask'][batch_idx, :, 0].to(action_in.device) # (T). else: action_in = None if 'action' in output_dict['y0_pred_entries'] and 'action' in data_batch['a4d_latent']: action_pred = cached_decode_latent_lowdim_auto(output_dict['y0_pred_entries'], 'action', a4d_vae)[batch_idx] # (T, C). action_pred_mask = data_batch['a4d_latent']['action_output_mask'][batch_idx, :, 0].to(action_pred.device) # (T). action_gt = cached_decode_latent_lowdim_auto(output_dict['x0_entries'], 'action', a4d_vae)[batch_idx] # (T, C). action_gt_mask = data_batch['a4d_latent']['action_supervise_mask'][batch_idx, :, 0].to(action_gt.device) # (T). else: action_pred = None action_gt = None if proprio_in is not None: mixed_in = proprio_in * proprio_in_mask[:, None] # (T, C) mixed_in_mask = proprio_in_mask elif action_in is not None: mixed_in = action_in * action_in_mask[:, None] # (T, C) mixed_in_mask = action_in_mask else: raise RuntimeError(f'No action conditioning at all found even though either proprio and/or action exists in x0_entries?') # We want "action conditioning" (typically proprio key) and "action prediction" (typically action key) # to be shown together (white and colored respectively), even if they correspond to different streams. # If nothing is predicted, just show the input anyway (in the prediction and GT columns). if action_pred is not None and action_gt is not None: mixed_pred = mixed_in + action_pred * action_pred_mask[:, None] # (T, C) mixed_gt = mixed_in + action_gt * action_gt_mask[:, None] # (T, C) else: mixed_pred = mixed_in mixed_gt = mixed_in T_orig = None # If orig_action exists (i.e., action was perturbed), use it for GT visualization # This shows perturbed actions on PRED video and clean/original actions on GT video if 'orig_action' in meta_dl_batch: orig_action = torch.zeros_like(action_in) # (T_extended, C) T_orig = meta_dl_batch['orig_action'][batch_idx].shape[0] # e.g., 41 # Fill first T_orig frames with original (clean) action, and keep the rest as zeros (no action) orig_action[:T_orig, :] = meta_dl_batch['orig_action'][batch_idx] # For GT: use original (clean) action instead of perturbed mixed_gt_viz = mixed_in + orig_action * action_gt_mask[:, None] if action_gt is not None else orig_action else: mixed_gt_viz = mixed_gt Ta = mixed_gt.shape[0] Tp = cur_gt_viz.shape[1] # log.debug(f'Ta: {Ta}, Tp: {Tp}') # assert Ta >= Tp, f'Insufficient actions: Ta ({Ta}) < Tp ({Tp})' # Get original action for coord text color (show deviation from clean) orig_action_for_viz = None if 'orig_action' in meta_dl_batch: orig_action_for_viz = torch.zeros_like(action_in) # (T_extended, C) T_orig = meta_dl_batch['orig_action'][batch_idx].shape[0] # e.g., 41 # Fill first T_orig frames with original action, and keep the rest as zeros (no action) orig_action_for_viz[:T_orig, :] = meta_dl_batch['orig_action'][batch_idx] # Fill first 41 frames for t in range(min(Ta, Tp)): t_idx = t if T_orig is None or t < T_orig else T_orig - 1 if cams is not None: cur_cam = cams[(t_idx, cam_idx)][batch_idx] else: # Synthesize a virtual/fake egocentric-style camera for datasets without extrinsics # (e.g. YAM XDOF, HumanoidEveryday, etc.). # NOTE: EE action xyz are in a body frame: X=forward, # Y=left, Z=up, origin at the torso, hands ~0.2-0.4m forward. # NOTE: could tweak elev for other views (90=top-down, 0=forward, ~60=agibot-like). # NOTE: AnyData Camera() takes Tcw as cam->world (cols = cam axes in world). H_img, W_img = cur_gt_viz.shape[2], cur_gt_viz.shape[3] elev, dist = np.deg2rad(45.0), 0.9 # 1.2 C = np.array([-dist * np.cos(elev), 0.0, dist * np.sin(elev)]) z = -C / np.linalg.norm(C) # optical axis (look at origin) x = np.cross(z, [0.0, 0.0, 1.0]); x /= np.linalg.norm(x) # image right y = np.cross(z, x) # image down Tcw = np.eye(4, dtype=np.float32) Tcw[:3, 0], Tcw[:3, 1], Tcw[:3, 2], Tcw[:3, 3] = x, y, z, C cur_cam = Camera(Tcw=torch.from_numpy(Tcw).unsqueeze(0), hw=(H_img, W_img)) if has_output and mixed_pred is not None: # PRED video: show perturbed actions with coord text showing deviation cur_pred_viz[:, t] = draw_actions( cur_pred_viz[:, t], cur_cam, action=mixed_pred[None], condition=mixed_in_mask[None], t=t, orig_action=orig_action_for_viz[None] if orig_action_for_viz is not None else None, show_coords=True) if has_supervise and mixed_gt_viz is not None: # GT video: show clean actions (deviation = 0, so always green) cur_gt_viz[:, t] = draw_actions( cur_gt_viz[:, t], cur_cam, action=mixed_gt_viz[None], condition=mixed_in_mask[None], t=t, orig_action=mixed_gt_viz[None], # Same as action => deviation = 0 show_coords=True) log.debug(f'Drawing actions onto {k} ' f'took {time.time() - start_time:.2f} seconds', rank0_only=False) def paint_traj( cur_pred_viz, cur_gt_viz, k, has_output, has_supervise, output_dict, data_batch, meta_dl_batch, a4d_vae, batch_idx, ): ''' Paint (2D-projected) input + predicted + target trajectories onto cur_pred_viz / cur_gt_viz in place (via slice assignment). Called once per rgb0 entry when traj exists. ''' start_time = time.time() traj_in = cached_decode_latent_lowdim_auto(output_dict['x0_entries'], 'traj', a4d_vae)[batch_idx] # (T, 2). traj_in_mask = data_batch['a4d_latent']['traj_input_mask'][batch_idx, :, 0].to(traj_in.device) # (T). mixed_in = traj_in * traj_in_mask[:, None] # (T, 2) if 'traj' in output_dict['y0_pred_entries']: traj_pred = cached_decode_latent_lowdim_auto(output_dict['y0_pred_entries'], 'traj', a4d_vae)[batch_idx] # (T, 2). traj_gt = traj_in traj_pred_mask = data_batch['a4d_latent']['traj_output_mask'][batch_idx, :, 0].to(traj_pred.device) # (T). traj_gt_mask = data_batch['a4d_latent']['traj_supervise_mask'][batch_idx, :, 0].to(traj_gt.device) # (T). mixed_pred = mixed_in + traj_pred * traj_pred_mask[:, None] # (T, 2) mixed_gt = mixed_in + traj_gt * traj_gt_mask[:, None] # (T, 2) else: mixed_pred = mixed_in mixed_gt = mixed_in if 'orig_traj' in meta_dl_batch: # When perturb_traj is active, show both clean and perturbed trajectories on both PRED and GT orig_traj = meta_dl_batch['orig_traj'][batch_idx] # (T, 2). Ta = orig_traj.shape[0] else: orig_traj = None Ta = mixed_pred.shape[0] Tp = cur_gt_viz.shape[1] for t in range(min(Ta, Tp)): traj_bounds = get_traj_bounds(orig_traj, mixed_pred, mixed_gt) if orig_traj is not None: # clean = dark gray, small => text on GT only cur_pred_viz[:, t] = draw_traj( cur_pred_viz[:, t], orig_traj, None, True, t, traj_bounds, 1, show_coords=False) cur_gt_viz[:, t] = draw_traj( cur_gt_viz[:, t], orig_traj, None, True, t, traj_bounds, 1, show_coords=True) if has_output and mixed_pred is not None: # perturbed = colorful, medium => text on PRED only cur_pred_viz[:, t] = draw_traj( cur_pred_viz[:, t], mixed_pred, traj_in_mask, False, t, traj_bounds, 2, show_coords=True) if has_supervise and mixed_gt is not None: # perturbed = colorful, medium cur_gt_viz[:, t] = draw_traj( cur_gt_viz[:, t], mixed_gt, traj_in_mask, False, t, traj_bounds, 2, show_coords=False) log.debug(f'Drawing traj onto {k} ' f'took {time.time() - start_time:.2f} seconds', rank0_only=False) def collect_sample_lowdim(data_batch, output_dict, a4d_vae, batch_idx, uses_traj, uses_proprio_or_action): ''' Curated, eval-ready lowdim for one sample: 2D trajectory (clean orig + perturbed input + predicted delta + combined) with masks, plus action/proprio when present. Decoded to real units (trajectory = (lateral, forward) meters relative to first frame). Mirrors paint_traj / paint_actions extraction; returns raw tensors (the caller runs _to_json_safe over the whole export). ''' out = {} meta_dl_batch = data_batch['meta'] al = data_batch['a4d_latent'] if uses_traj: traj_in = cached_decode_latent_lowdim_auto(output_dict['x0_entries'], 'traj', a4d_vae)[batch_idx] # (T, 2) perturbed input traj_in_mask = al['traj_input_mask'][batch_idx, :, 0] mixed_in = traj_in * traj_in_mask.to(traj_in.device)[:, None] traj = dict( units='(lateral, forward) meters, relative to first frame', orig=meta_dl_batch['orig_traj'][batch_idx] if 'orig_traj' in meta_dl_batch else None, # clean input_perturbed=traj_in, input_mask=traj_in_mask, ) if 'traj' in output_dict['y0_pred_entries']: traj_pred = cached_decode_latent_lowdim_auto(output_dict['y0_pred_entries'], 'traj', a4d_vae)[batch_idx] traj_out_mask = al['traj_output_mask'][batch_idx, :, 0] traj['pred_delta'] = traj_pred traj['perturbed_plus_pred'] = mixed_in + traj_pred * traj_out_mask.to(traj_pred.device)[:, None] traj['output_mask'] = traj_out_mask traj['supervise_mask'] = al['traj_supervise_mask'][batch_idx, :, 0] out['trajectory'] = traj if uses_proprio_or_action: ld = {} if 'action' in output_dict['x0_entries'] and 'action' in al: ld['action_input'] = cached_decode_latent_lowdim_auto(output_dict['x0_entries'], 'action', a4d_vae)[batch_idx] if 'action' in output_dict['y0_pred_entries']: ld['action_pred'] = cached_decode_latent_lowdim_auto(output_dict['y0_pred_entries'], 'action', a4d_vae)[batch_idx] if 'proprio' in output_dict['x0_entries'] and 'proprio' in al: ld['proprio_input'] = cached_decode_latent_lowdim_auto(output_dict['x0_entries'], 'proprio', a4d_vae)[batch_idx] out['lowdim'] = ld # Curated ego_pose as a clean (T, 4, 4) array (the per-(t,cam) cams dict is too bulky to keep whole). anydata = data_batch.get('anydata', data_batch.get('vidar', {})) ego = anydata.get('ego_pose') if isinstance(anydata, dict) else None if isinstance(ego, dict) and len(ego): try: out['ego_pose'] = torch.stack([ego[t][batch_idx] for t in sorted(ego.keys())], dim=0) # (T, 4, 4) except Exception: pass return out @torch.no_grad() @resilient_subroutine(max_failure_rate=0.2, grace_period=10, default=dict) def create_save_visuals_a4d( phase: str = 'train', train_step: int = -1, data_batch: dict[str, torch.Tensor] = None, output_dict: dict[str, torch.Tensor] = None, a4d_vae=None, config=None, kendall_loss=None, metrics_dict=None, dst_dp: str = None, detail: int = 2, quality_pix: int = 7, quality_lat: int = 8, val_step: int = -1, batch_idx: int = None, gc_collect: bool = True, **leftover, ): ''' NOTE: This function can be called on ALL ranks, but processes only the FIRST sample within the batch by default. This can be changed by setting batch_idx to 'all'. :param phase (str): train / val / test. :param train_step (int): Batch index, like training step but changes along with gradient accumulation. Starts from 1. :param data_batch (dict): Input batch, contains anydata, fps, prompt, etc. :param output_dict (dict): Output batch, contains both encoded (latent) and decoded (raw) input/output/GT entries. :param a4d_vae (VAE). :param config (Any4DConfig). :param kendall_loss (float): Total mean loss value for the ENTIRE batch. :param metrics_dict (dict): Metrics dictionary, as returned by calculate_metrics_a4d(). :param dst_dp (str): Output directory. :param detail (int): 1 / 2 / 3. :param quality_pix (int): Quality of exported pixel MP4 video files. :param quality_lat (int): Quality of exported visualized latent MP4 video files. :param val_step (int): Validation step index, since training step will be frozen during validation. :param batch_idx (int): Batch index (None / int / 'all). :param gc_collect (bool): Whether to collect garbage after saving visuals. :return nice_videos (dict): Dictionary mapping short descriptions to exported MP4 video paths. ''' output_dict = detach_dict_recursive(output_dict) B = output_dict['x0_streams']['v0'].shape[0] # =========================== # ======== Recursion ======== # =========================== if batch_idx == 'all' and B > 1: nice_videos_all = dict() for b in range(B): nice_videos = create_save_visuals_a4d( phase, train_step, data_batch, output_dict, a4d_vae, config, kendall_loss, metrics_dict, dst_dp, detail, quality_pix, quality_lat, val_step, b, gc_collect) for k in nice_videos.keys(): # nice_videos_all[f'{k}_b{b}'] = nice_videos[k] nice_videos_all[k] = nice_videos[k] # should already differentiate by return nice_videos_all if isinstance(batch_idx, int): log_batch_idx = True else: batch_idx = 0 log_batch_idx = False # validation calls this with 'all' but is often B=1, so don't log # ============================= # ======== Preparation ======== # ============================= assert phase in ['train', 'val', 'test'] assert detail in [1, 2, 3] data_key = 'anydata' if 'anydata' in data_batch else 'vidar' anydata_batch = data_batch[data_key] meta_ds_batch = anydata_batch['metadata'] # from dataSET (though vidar batches don't have it) meta_dl_batch = data_batch['meta'] # from dataLOADER # has_actions = 'anydata' in data_batch and 'action' in data_batch['anydata'] uses_proprio_or_action = ('proprio' in output_dict['x0_entries'] and 'proprio' in data_batch['a4d_latent']) or \ ('action' in output_dict['x0_entries'] and 'action' in data_batch['a4d_latent']) has_cams = 'cams_orig' in anydata_batch uses_traj = 'traj' in output_dict['x0_entries'] and 'traj' in data_batch['a4d_latent'] used_streams = sorted([k for k in output_dict['x0_streams'].keys() if not(k.endswith('_mask') or k.endswith('_cdec'))]) used_entries = sorted([k for k in output_dict['x0_entries'].keys() if not(k.endswith('_mask') or k.endswith('_cdec'))]) my_rank_td = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 exp_tag = config.job.undated_name dl_idx = unpack_nested(data_batch, 'dl_idx', batch_idx=batch_idx, dtype=int, optional=True) # int dl_key = data_batch['dl_key'] if 'dl_key' in data_batch and data_batch['dl_key'] is not None else '' # str dset_name = unpack_nested(data_batch, 'dset_name', batch_idx=batch_idx, dtype=str, optional=True) # str sample_id = unpack_nested(data_batch, 'sample_id', batch_idx=batch_idx, dtype=str, optional=True) # str sample_file = unpack_nested(meta_ds_batch, 'path', batch_idx=batch_idx, dtype=str, optional=True) or '' sample_file = os.path.basename(sample_file) # str, NOTE: webdataset tarfile name only sample_tags = meta_ds_batch.get('tags', [['unknown']] * (batch_idx + 1))[batch_idx] # vidar fallback prompt = unpack_nested(data_batch, 'prompt', batch_idx=batch_idx, dtype=str, optional=True) # str fps_pix = unpack_nested(data_batch, 'fps', batch_idx=batch_idx, dtype=float, optional=False) # float fps_lat = fps_pix / a4d_vae.vae_ratio[0] # float V_max = config.num_views V_used = [v for v in range(V_max) if f'v{v}' in used_streams] used_cams = unpack_nested(meta_dl_batch, 'used_cams', batch_idx=batch_idx, dtype=str, optional=True) # str meta_tasks = unpack_nested(meta_dl_batch, 'tasks', batch_idx=batch_idx, dtype=str, optional=True) # str if used_cams is not None: used_cams = [int(c) for c in used_cams.split(',')] # NOTE(bvh): camera NAMES come from dataset metadata and exist regardless of whether the cams # (pose) modality is loaded -- do NOT gate name resolution on has_cams, else pose-free splits # (e.g. YAMxdof ILA) fall back to bare indices in filenames and info.json. cam_names_meta = meta_ds_batch.get('cameras', None) if hasattr(meta_ds_batch, 'get') else None if used_cams is not None and cam_names_meta is not None: view_idx_to_cam_idx = used_cams # e.g. [0, 3, 1] # list() copy: += would extend the metadata list in place and pollute it with 'Z' entries cam_idx_to_cam_name = list(cam_names_meta[batch_idx]) + ['Z'] * V_max # 'Z' = padded views sel_cam_names = [str(cam_idx_to_cam_name[i]) for i in view_idx_to_cam_idx] # e.g. ['ext5', 'Z', 'ext1'] elif used_cams is not None: sel_cam_names = [str(c) for c in used_cams] # in dataloader space else: sel_cam_names = None # Build per-view, per-frame camera labels: "cam_name (x, y, z)" # cam_labels[view_idx] = list of T strings, one per raw frame. cam_labels = None if sel_cam_names is not None and has_cams: try: cams_dict = anydata_batch['cams'] avail_times_vis = sorted(set(k[0] for k in cams_dict.keys())) avail_cams_vis = sorted(set(k[1] for k in cams_dict.keys())) cam_labels = {} for v_idx, cam_name in enumerate(sel_cam_names): cam_idx = used_cams[v_idx] if used_cams is not None and v_idx < len(used_cams) else v_idx labels = [] for t in avail_times_vis: key = (t, avail_cams_vis[cam_idx] if cam_idx < len(avail_cams_vis) else 0) if key in cams_dict: xyz = cams_dict[key].Tcw[batch_idx, :3, 3] labels.append(f'{cam_name} ({xyz[0]:.2f}, {xyz[1]:.2f}, {xyz[2]:.2f})') else: labels.append(cam_name) cam_labels[v_idx] = labels except Exception: cam_labels = None if sel_cam_names is not None: cams_friendly = ','.join(sel_cam_names) else: cams_friendly = None # Subsampled (kept) frame indices from temporal stride/trim (set in DataTransforms). frame_indices = anydata_batch.get('frame_indices') if frame_indices is None: frames_friendly = None elif len(frame_indices) <= 10: frames_friendly = str(frame_indices) else: frames_friendly = f'[{frame_indices[0]}, {frame_indices[1]}, ..., {frame_indices[-1]}] ({len(frame_indices)})' tags_friendly = '[' + ', '.join(str(t) for t in sample_tags) + ']' if 'train' in phase: loss = kendall_loss.item() # single float. sigmas = {k: v[batch_idx].item() for (k, v) in output_dict['sigmas'].items()} # single float. sigma_v0 = sigmas['v0'] if 'train' in phase or 'val' in phase: prefix_both = f'{dst_dp}/{exp_tag}_{phase}' if train_step >= 0: prefix_both += f'_s{train_step}' if dl_key is not None and dl_key != '': prefix_both += f'_{dl_key}--{dset_name}' else: prefix_both += f'_{dset_name}' if 'train' in phase: prefix_both += f'_n{sigma_v0:.2f}_l{loss:.3f}_r{my_rank_td}' elif 'val' in phase: prefix_both += f'_d{dl_idx}_i{val_step}_r{my_rank_td}' if log_batch_idx: prefix_both += f'_b{batch_idx}' if cams_friendly is not None and len(cams_friendly) <= 16: prefix_both += f'_c{cams_friendly}' elif 'test' in phase: # TODO(bvh) this branch is not actually used yet, infer.py becomes phase='val' because it calls model.validation_step() prefix_both = f'{dst_dp}/{exp_tag}_{phase}' if train_step >= 0: prefix_both += f'_s{train_step}' prefix_both += f'_d{dl_idx}' if cams_friendly is not None and len(cams_friendly) <= 16: prefix_both += f'_c{cams_friendly}' prefix_detail = f'{prefix_both}_all/' prefix_nice = f'{prefix_both}_' log.debug(f'Saving {phase} visuals ' f'with detail level {detail} to: {prefix_both}...', rank0_only=False) start_time = time.time() # Persist a per-sample JSON sidecar next to the visuals: curated conditioning (orig + perturbed 2D # trajectory, masks, action/proprio, ego_pose) plus the data_batch / output_dict pruned of large # tensors and many-key dicts (small extrinsics / ego kept). For later reference / metrics / eval. if getattr(config, 'viz_export_info_json', True): try: export = dict( sample_id=sample_id, dset_name=dset_name, dl_idx=dl_idx, dl_key=dl_key, phase=phase, train_step=train_step, val_step=val_step, sample_file=sample_file, sample_tags=sample_tags, prompt=prompt, fps_pix=fps_pix, fps_lat=fps_lat, tasks=meta_tasks, used_cams=used_cams, sel_cam_names=sel_cam_names, frame_indices=frame_indices, metrics=metrics_dict, **collect_sample_lowdim(data_batch, output_dict, a4d_vae, batch_idx, uses_traj, uses_proprio_or_action), data_batch_pruned=data_batch, output_dict_pruned=output_dict, ) save_json(to_json_safe(export), f'{prefix_nice}info.json', indent=1) except Exception as e: log.warning(f'info json export failed: {e}', rank0_only=False) # ======================================== # ======== Process latent streams ======== # ======================================== nice_videos = dict() # These should (also) go to wandb def decode_viz_fn(d, k): return cached_decode_viz_latent_video_auto(d, k, a4d_vae, config, first=True) def save_vid_lat_fn(x, p, d): return flex_save_video( x, p=p, u=a4d_vae.video_tokenizer.spatial_compression_factor, f=fps_lat + 1, h=2, q=quality_lat, e=(detail >= d), v=(1 if train_step <= 1 else 0)) def save_vid_pix_fn(x, p, d): return flex_save_video( x, p=p, u=1, f=fps_pix, h=4, q=quality_pix, e=(detail >= d), v=(1 if train_step <= 1 else 0)) if detail >= 3: for v in V_used: cur_xt_lat = output_dict['xt_streams'][f'v{v}'][batch_idx] # (C, T, H, W) cur_y0_pred_lat = output_dict['y0_pred_streams'][f'v{v}'][batch_idx] # (C, T, H, W) cur_x0_lat = output_dict['x0_streams'][f'v{v}'][batch_idx] # (C, T, H, W) save_vid_lat_fn( cur_xt_lat, f'{prefix_detail}stream_xt_v{v}_pca_c{cur_xt_lat.shape[0]}.mp4', 3) save_vid_lat_fn( cur_y0_pred_lat, f'{prefix_detail}stream_y0_pred_v{v}_pca_c{cur_y0_pred_lat.shape[0]}.mp4', 3) save_vid_lat_fn( cur_x0_lat, f'{prefix_detail}stream_x0_v{v}_pca_c{cur_x0_lat.shape[0]}.mp4', 3) # =========================================================== # ======== Process raw entries & construct galleries ======== # =========================================================== annot_vids = dict() # For decorated visualizations. annot_vids['in'] = dict() # Assembled conditioning input samples. annot_vids['pred'] = dict() # Predicted output samples. annot_vids['gt'] = dict() # Clean target samples. clean_vids = dict() # For raw pixel videos. clean_vids['in'] = dict() clean_vids['pred'] = dict() clean_vids['gt'] = dict() frame_mm = dict() frame_mm['input'] = dict() frame_mm['output'] = dict() frame_mm['supervise'] = dict() for k in used_entries: if k not in config.all_highdim_entries: continue if k not in data_batch['a4d_latent']: continue frame_mm['input'][k] = get_per_pixel_frame_mask_means( data_batch['a4d_latent'][k + '_input_mask'][batch_idx], a4d_vae) frame_mm['output'][k] = get_per_pixel_frame_mask_means( data_batch['a4d_latent'][k + '_output_mask'][batch_idx], a4d_vae) frame_mm['supervise'][k] = get_per_pixel_frame_mask_means( data_batch['a4d_latent'][k + '_supervise_mask'][batch_idx], a4d_vae) has_input = frame_mm['input'][k].sum() > 0 has_output = frame_mm['output'][k].sum() > 0 has_supervise = frame_mm['supervise'][k].sum() > 0 cur_xt_pix = decode_viz_fn(output_dict['xt_entries'], k) # (C, T, H, W) tensor cur_y0_pred_pix = decode_viz_fn(output_dict['y0_pred_entries'], k) # (C, T, H, W) tensor cur_x0_pix = decode_viz_fn(output_dict['x0_entries'], k) # (C, T, H, W) tensor # Export raw pixel MP4s up front, before any annotation. Side effect only; # the annotation pipeline below runs on independent tensor clones. if has_input: save_vid_pix_fn(cur_xt_pix, f'{prefix_detail}entry_xt_{k}_pix.mp4', 2) if has_output: save_vid_pix_fn(cur_y0_pred_pix, f'{prefix_detail}entry_y0_pred_{k}_pix.mp4', 2) if has_supervise: save_vid_pix_fn(cur_x0_pix, f'{prefix_detail}entry_x0_{k}_pix.mp4', 2) # Annotation working tensors (numeric overlays stay in tensor space). cur_in_viz = cur_xt_pix.clone() cur_pred_viz = cur_y0_pred_pix.clone() cur_gt_viz = cur_x0_pix.clone() # Draw borders to distinguish between input (green) vs output (blue) vs supervise (red). paint_mask_indicator( cur_in_viz, frame_mm, k, [[0, 0, 0], [1, 0, 0], [0, 0, 0]], config.viz_mask_border_width) paint_mask_indicator( cur_pred_viz, frame_mm, k, [[0, 0, 0], [1, 0, 0], [0, 1, 0]], config.viz_mask_border_width) paint_mask_indicator( cur_gt_viz, frame_mm, k, [[0, 0, 1], [0, 0, 0], [0, 0, 0]], config.viz_mask_border_width) # Paint (2D projected) input & output & target actions onto respective pred & GT videos # (all views to depict/understand 3D consistency) if 'rgb' in k and uses_proprio_or_action: view_idx = [c[0] for c in enumerate(config.video_entries) if k in c[1]][0] cam_idx = used_cams[view_idx] cams = anydata_batch['cams_orig'] if has_cams else None paint_actions( cur_pred_viz, cur_gt_viz, k, has_output, has_supervise, output_dict, data_batch, meta_dl_batch, a4d_vae, batch_idx, cams, cam_idx) # Paint (BEV / non-projected) perturbed & original trajectories # onto pred & GT video (first/front view only) if 'rgb0' in k and uses_traj: paint_traj( cur_pred_viz, cur_gt_viz, k, has_output, has_supervise, output_dict, data_batch, meta_dl_batch, a4d_vae, batch_idx) # Single tensor -> numpy transition per entry. annot AND clean variants. cur_in_clean_arr = tensor_to_thwc_uint8(cur_xt_pix) cur_pred_clean_arr = tensor_to_thwc_uint8(cur_y0_pred_pix) cur_gt_clean_arr = tensor_to_thwc_uint8(cur_x0_pix) cur_in_arr = tensor_to_thwc_uint8(cur_in_viz) cur_pred_arr = tensor_to_thwc_uint8(cur_pred_viz) cur_gt_arr = tensor_to_thwc_uint8(cur_gt_viz) # 'cams' tile height-subsample on annot only (moved out of make_gallery so # labels land at final tile resolution); clean stays at original height. if 'cams' in k: cur_in_arr = cur_in_arr[:, ::3, :, :] cur_pred_arr = cur_pred_arr[:, ::3, :, :] cur_gt_arr = cur_gt_arr[:, ::3, :, :] # Top-right cam name / per-frame cam label (annot only; clean stays unlabeled). m = re.search(r'\d+', k) if m is not None and (sel_cam_names is not None or cam_labels is not None): idx = int(m.group()) if cam_labels is not None and idx in cam_labels: frame_labels = cam_labels[idx] for arr in (cur_in_arr, cur_pred_arr, cur_gt_arr): T_cur = arr.shape[0] for t in range(T_cur): lbl = frame_labels[t] if t < len(frame_labels) else frame_labels[-1] paint_label(arr[t:t+1], lbl, font_scale=0.5, font_color=(210, 210, 210), position='top_right') elif sel_cam_names is not None and idx < len(sel_cam_names): for arr in (cur_in_arr, cur_pred_arr, cur_gt_arr): paint_label(arr, sel_cam_names[idx], font_scale=0.5, font_color=(210, 210, 210), position='top_right') # Top-left entry-key + count title (column-specific count text; raw k stays as dict key). # 'in' tile only kept when this entry is NOT a mixed input/output -- this mirrors # the dedup that the old post-loop rename block was performing. count_in = (frame_mm['input'][k] > 0.5).sum().item() count_out = (frame_mm['output'][k] > 0.5).sum().item() count_gt = (frame_mm['supervise'][k] > 0.5).sum().item() keep_in = has_input \ and not any(k.startswith(e) for e in config.viz_input_blacklist) \ and count_out == 0 if keep_in: paint_label(cur_in_arr, f'{k} ({count_in} in)') clean_vids['in'][k] = cur_in_clean_arr annot_vids['in'][k] = cur_in_arr if has_output: paint_label(cur_pred_arr, f'{k} ({count_in} in, {count_out} out)') clean_vids['pred'][k] = cur_pred_clean_arr annot_vids['pred'][k] = cur_pred_arr if has_supervise: paint_label(cur_gt_arr, f'{k} ({count_gt} gt)') clean_vids['gt'][k] = cur_gt_clean_arr annot_vids['gt'][k] = cur_gt_arr # Sort but put cams at the bottom since often less important. for col in ('in', 'pred', 'gt'): sorted_keys = sorted_camslast(list(annot_vids[col].keys())) annot_vids[col] = {k: annot_vids[col][k] for k in sorted_keys} clean_vids[col] = {k: clean_vids[col][k] for k in sorted_keys if k in clean_vids[col]} headers = [] footers = [] if prompt is not None: headers.append(prompt) metrics_str = ' '.join([f'{k}: {v[batch_idx]}' for (k, v) in metrics_dict['friendly'].items() if v[batch_idx] != '']) footers.append(f'{meta_tasks} {metrics_str}') if 'sampling_params' in output_dict: num_steps = output_dict['sampling_params']['num_steps'] else: num_steps = None # NOTE(bvh): val_step and batch_idx are ignored here because # already in filename + should be captured / distinguished by dl_idx: date_str = datetime.now().strftime('%m/%d') footers.append(f'{exp_tag} {date_str} {phase} step: {train_step} ' + (f'rank: {my_rank_td} {dl_key} {dset_name} dl_idx: {dl_idx} ') + (f'diff: {num_steps} ' if num_steps is not None else '') # + (f'cam: {cams_friendly} ' if cams_friendly is not None else '') + (f'sig_v0: {sigma_v0:.2f} loss: {loss:.3f}' if 'train' in phase else '')) footers.append(f'tags: {tags_friendly} ' # + (f'cams: {cams_friendly} ' if cams_friendly is not None else '') + (f'frames: {frames_friendly} ' if frames_friendly is not None else '') + f'fps: {fps_pix:.1f}') # if 'sampling_params' in output_dict: # sampling_params = output_dict['sampling_params'] # params_str = ' '.join([f'{k}: {v}' for (k, v) in sampling_params.items()]) # footers.append(f'{params_str}') footers.append(f'{sample_file} {sample_id}') # ================================================ # ======== Export generic Any4D galleries ======== # ================================================ if len(annot_vids['in']) == 0: columns = ['pred', 'gt'] else: columns = ['in', 'pred', 'gt'] # Default simple horizontal concat (easier to extract raw frames / figures from) simple_cat = simple_horizontal_concat(annot_vids, columns) wandb_cat_key = f'{dl_key}--{dset_name}/CAT--{dl_idx}' (nice_videos[wandb_cat_key], _) = flex_save_video( simple_cat, p=f'{prefix_nice}cat.mp4', d=3, u=1, f=fps_pix, h=4, q=quality_pix, v=(1 if train_step <= 1 else 0)) gal_all_pix = make_gallery(annot_vids, columns, headers, footers) wandb_gal_key = f'{dl_key}--{dset_name}/GAL--{dl_idx}' (nice_videos[wandb_gal_key], _) = flex_save_video( gal_all_pix, p=f'{prefix_nice}gal.mp4', d=3, u=1, f=fps_pix, h=4, q=quality_pix, v=(1 if train_step <= 1 else 0)) log.info(f'Saved {phase} visuals with detail level {detail}', rank0_only=False) log.info(f'Took {time.time() - start_time:.2f} seconds', rank0_only=False) # Gallery without GT column (useful for inference where GT is unavailable or stub data) if 'nogt' in config.viz_extra_modes: columns_nogt = [c for c in columns if c != 'gt'] gal_nogt = make_gallery(annot_vids, columns_nogt, headers, footers) wandb_nogt_key = f'{dl_key}--{dset_name}/GALNOGT--{dl_idx}' (nice_videos[wandb_nogt_key], _) = flex_save_video( gal_nogt, p=f'{prefix_nice}gal_nogt.mp4', d=3, u=1, f=fps_pix, h=4, q=quality_pix, v=(1 if train_step <= 1 else 0)) # ====================================================== # ======== Export simplified / custom galleries ======== # ====================================================== # Extra visualization modes below, tailored for specific use cases if 'anyact1' in config.viz_extra_modes: # Robotics world model: pred views (rgb0, rgb1) only, prompt as header custom_vid = viz_custom(annot_vids, ['rgb0', 'rgb1'], gt=False, headers=[prompt] if prompt else []) if custom_vid is not None: wandb_key = f'{dl_key}--{dset_name}/ACT--{dl_idx}' (nice_videos[wandb_key], _) = flex_save_video( custom_vid, p=f'{prefix_nice}act.mp4', d=3, u=1, f=fps_pix, h=4, q=quality_pix, v=(1 if train_step <= 1 else 0)) if 'anydrive1' in config.viz_extra_modes: # Driving world model: pred views (rgb1, rgb0, rgb2) on top, GT on bottom custom_vid = viz_custom(annot_vids, ['rgb1', 'rgb0', 'rgb2'], gt=True, headers=[]) if custom_vid is not None: wandb_key = f'{dl_key}--{dset_name}/DRV--{dl_idx}' (nice_videos[wandb_key], _) = flex_save_video( custom_vid, p=f'{prefix_nice}drv.mp4', d=3, u=1, f=fps_pix, h=4, q=quality_pix, v=(1 if train_step <= 1 else 0)) # NOTE(bvh): keep the GAIA stuff synced with unified_anydrive dataloader. if 'gaia1' in config.viz_extra_modes: # rgb0-4 = front_wide, left_side_pvm, right_side_pvm, rear_pvm, front_pvm. # desired = front_wide | front/left | rear/right. # entries_2d = [['rgb0'], ['rgb4', 'rgb1'], ['rgb3', 'rgb2']] # NOTE(bvh): ^ matches Yiting's TSS4 MDC visualization pred_rgb_vids = {k: v for k, v in annot_vids['pred'].items() if k.startswith('rgb')} num_rgb_views = len(pred_rgb_vids) # ^ could be different from V_used and/or V_max! assert num_rgb_views in [5, 11], f'GAIA{num_rgb_views} not supported' custom_vid = viz_anydrive_gaia(pred_rgb_vids, sel_cam_names=sel_cam_names) if custom_vid is not None: wandb_key = f'{dl_key}--{dset_name}/GAIA{num_rgb_views}--{dl_idx}' (nice_videos[wandb_key], _) = flex_save_video( custom_vid, p=f'{prefix_nice}gaia{num_rgb_views}.mp4', d=3, u=1, f=fps_pix, h=4, q=quality_pix, v=(1 if train_step <= 1 else 0)) # if 'gaia11_1' in config.viz_extra_modes: # rgb0-10 = front_wide, left_side_pvm, right_side_pvm, rear_pvm, front_pvm, # left_side_forward, right_side_forward, rear_tele, front_tele, # left_side_rearward, right_side_rearward. # desired = | front_wide | pvm(front/left/rear/right) | wide(l_f/l_r/r_f/r_r) | tele(front/rear) # entries_2d = [['rgb0'], ['rgb10', 'rgb1'], ['rgb9', 'rgb2'], ['rgb8', 'rgb3'], ['rgb7', 'rgb4'], ['rgb6', 'rgb5']] # BVH shape notes (any4d): # nice_videos: {'pix': '/basile/ws/repos/cpred2-basile/logs_a4d/a4d2/debug/08-13-16-18_b2/visuals/b2_val_s1_r0_lbmv12_d885044_v1,0_all.mp4'} if gc_collect: trigger_manual_gc(source='create_save_visuals_a4d') return nice_videos