''' BVH, Apr 2026. Any4D inference with donor view splicing for anydata datasets. Runs two dataloaders -- a source (e.g., HumEv) and a donor (e.g., AgiBot) -- and injects donor cameras (extrinsics + intrinsics) and donor RGB into each source batch to produce multi-view inputs for datasets that have fewer views than the model expects. For the common single-dataset case (optionally with a lightweight single-camera RGB override from a video file), use infer_anydata.py instead. ''' import os import sys sys.path.insert(0, os.getcwd()) # Library imports import argparse import copy import time import traceback import warnings import numpy as np import torch from rich import print from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TimeRemainingColumn from torch.utils.data import DataLoader # Internal imports from custom.dataloader.collate import any4d_collate from custom.dataset.anydata_dataset import AnyDataset from custom.eval.infer_anydata import ( add_base_args, load_dataset, setup_runtime, process_one_sample, finalize_run, ) np.set_printoptions(precision=3, suppress=True) warnings.filterwarnings('ignore', category=FutureWarning) def parse_args(): parser = argparse.ArgumentParser( description='Run Any4D inference with donor camera + RGB splicing from a second dataset') add_base_args(parser) # Donor dataset for view splicing (for datasets with fewer views than the model expects) parser.add_argument('--donor_dset_cfg', type=str, default=None, help='Path to donor dataset yaml (e.g., AgiBot) for camera + RGB injection') parser.add_argument('--inject_num_views', type=int, default=3, help='Target number of views after injection (default 3)') # Legacy: static .pt file with camera params (no RGB, single sample) parser.add_argument('--inject_cameras', type=str, default=None, help='(legacy) Path to .pt file with camera params, use --donor_dset_cfg instead') args = parser.parse_args() return args def inject_cameras_into_batch(data_batch, donor_batch, target_views=3): ''' Splice a single-view dataset (e.g., HumEv) with a multi-view donor (e.g., AgiBot). Uses donor's cameras (extrinsics, intrinsics) and wrist RGB for the injected views, while keeping the source dataset's head RGB as the conditioned input. :param data_batch: collated data batch from the source (single-view) dataset. :param donor_batch: collated data batch from the donor (multi-view) dataset. :param target_views: number of views to produce. :return: modified data_batch with spliced views. ''' anydata = data_batch['anydata'] donor = donor_batch['anydata'] ts_keys = list(anydata['timestep'].keys()) avail_times = sorted(set(k[0] for k in ts_keys)) src_cam = sorted(set(k[1] for k in ts_keys))[0] T = len(avail_times) donor_ts_keys = list(donor['timestep'].keys()) donor_times = sorted(set(k[0] for k in donor_ts_keys)) donor_cams = sorted(set(k[1] for k in donor_ts_keys)) T_use = min(T, len(donor_times)) B = anydata['rgb'][(avail_times[0], src_cam)].shape[0] if 'rgb' in anydata else 1 donor_cam_names = donor['metadata'].get('cameras', [f'cam{i}' for i in range(len(donor_cams))]) if isinstance(donor_cam_names, list) and len(donor_cam_names) > 0 and isinstance(donor_cam_names[0], list): donor_cam_names = donor_cam_names[0] # unwrap batch dim from collation print(f'[cyan]Splicing: src {len(set(k[1] for k in ts_keys))} view(s) + donor {len(donor_cams)} views -> {target_views} views') print(f'[cyan] src T={T}, donor T={len(donor_times)}, using T={T_use}') print(f'[cyan] donor cam names: {donor_cam_names}') for new_cam_idx in range(1, target_views): if new_cam_idx >= len(donor_cams): break donor_cam = donor_cams[new_cam_idx] for t_idx in range(T_use): t_src = avail_times[t_idx] t_donor = donor_times[t_idx] key_new = (t_src, new_cam_idx) # Use donor's wrist RGB as GT for this view donor_rgb_key = (t_donor, donor_cam) if 'rgb' in donor and donor_rgb_key in donor['rgb']: anydata.setdefault('rgb', {})[key_new] = donor['rgb'][donor_rgb_key] elif 'rgb' in anydata: # Fallback: clone source RGB anydata['rgb'][key_new] = anydata['rgb'][(t_src, src_cam)].clone() # Donor extrinsics donor_ext_key = (t_donor, donor_cam) if 'extrinsics' in donor and donor_ext_key in donor['extrinsics']: anydata.setdefault('extrinsics', {})[key_new] = donor['extrinsics'][donor_ext_key] # Donor intrinsics if 'intrinsics' in donor and donor_ext_key in donor['intrinsics']: anydata.setdefault('intrinsics', {})[key_new] = donor['intrinsics'][donor_ext_key] # Timestep anydata['timestep'][key_new] = anydata['timestep'][(t_src, src_cam)] # Inject extrinsics for source camera (cam 0) from donor's cam 0 if missing has_src_ext = 'extrinsics' in anydata and any(k[1] == src_cam for k in anydata['extrinsics']) if not has_src_ext and 'extrinsics' in donor: anydata.setdefault('extrinsics', {}) for t_idx in range(T_use): t_src = avail_times[t_idx] t_donor = donor_times[t_idx] donor_key = (t_donor, donor_cams[0]) if donor_key in donor['extrinsics']: anydata['extrinsics'][(t_src, src_cam)] = donor['extrinsics'][donor_key] # Inject intrinsics for source camera if completely missing has_src_intr = 'intrinsics' in anydata and (avail_times[0], src_cam) in anydata['intrinsics'] if not has_src_intr and 'intrinsics' in donor: anydata.setdefault('intrinsics', {}) for t_idx in range(T_use): t_src = avail_times[t_idx] t_donor = donor_times[t_idx] donor_key = (t_donor, donor_cams[0]) if donor_key in donor['intrinsics']: anydata['intrinsics'][(t_src, src_cam)] = donor['intrinsics'][donor_key] # extrinsics_orig for visualization if 'extrinsics' in anydata: anydata['extrinsics_orig'] = copy.deepcopy(anydata['extrinsics']) # Update metadata cameras list. # After collation, cameras is [[cam0, cam1, ...]] (batch of lists). # We need to update the inner list for batch element 0. cameras = anydata['metadata'].get('cameras', [['0']]) if isinstance(cameras[0], list): orig_cam_name = cameras[0][0] else: orig_cam_name = cameras[0] new_cam_names = [orig_cam_name] + list(donor_cam_names[1:target_views]) # Store as batch-of-lists to match collation format anydata['metadata']['cameras'] = [new_cam_names] print(f'[green] Spliced {target_views - 1} donor views. Camera names: {new_cam_names}') print(f'[green] rgb entries: {len(anydata.get("rgb", {}))}') print(f'[green] extrinsics entries: {len(anydata.get("extrinsics", {}))}') return data_batch def main(args): start_time = time.time() model, config, lpips_loss, device, exp_cfg = setup_runtime(args) # Load donor dataset for view splicing (e.g., AgiBot cameras + wrist RGB for HumEv) donor_dataloader = None if args.donor_dset_cfg: print(f'[cyan]Loading donor dataset from {args.donor_dset_cfg}...') donor_ds = AnyDataset( dataset_dir=args.donor_dset_cfg, phase='test', single_gpu=True, shuffle=False, data_overrides={'subsample': args.stop_after} if args.stop_after > 0 else None, ) donor_dataloader = DataLoader( donor_ds, batch_size=args.batch_size, shuffle=False, num_workers=0, pin_memory=True, collate_fn=any4d_collate, ) print(f'[green]Donor dataset loaded: {len(donor_ds)} samples') # Load source dataset dataloader = load_dataset(args, exp_cfg=exp_cfg) # Determine how many samples to process num_total = len(dataloader) num_to_process = num_total if args.stop_after < 0 else min(args.stop_after, num_total) print(f'[cyan]Processing {num_to_process} samples...') all_scene_metrics = [] with Progress( SpinnerColumn(), TextColumn('[progress.description]{task.description}'), BarColumn(), TextColumn('[progress.percentage]{task.percentage:>3.0f}%'), TimeRemainingColumn(), ) as progress: task = progress.add_task('[cyan]Processing samples...', total=num_to_process) # Zip with donor dataloader if available, otherwise iterate source only if donor_dataloader is not None: data_iter = zip(dataloader, donor_dataloader) else: data_iter = ((batch, None) for batch in dataloader) for iteration, (data_batch, donor_batch) in enumerate(data_iter): if iteration >= num_to_process: break try: # Splice donor views (cameras + wrist RGB) into source batch inject_directives = None if donor_batch is not None: data_batch = inject_cameras_into_batch( data_batch, donor_batch, target_views=args.inject_num_views) # View ordering: injected wrist views first (predicted), # source head view last (conditioned). DVS conditions the last view(s). # cam 0 = head (input), cam 1,2 = wrists (predicted) # -> used_cams=[1, 2, 0] so rgb0=hand_left, rgb1=hand_right, rgb2=head(input) n_pred = args.inject_num_views - 1 inject_directives = { 'used_cams': list(range(1, args.inject_num_views)) + [0], 'num_pred_views': n_pred, } # NOTE(bvh): legacy --inject_cameras (.pt file) is no longer supported, # use --donor_dset_cfg instead for proper zipped iteration. process_one_sample( args, config, model, data_batch, lpips_loss, device, iteration, all_scene_metrics, extra_directives=inject_directives) progress.update(task, advance=1) except Exception as e: print(f'[red]Error processing sample {iteration}: {e}') print(f'[red]Traceback: {traceback.format_exc()}') print(f'[yellow]Skipping...') progress.update(task, advance=1) continue finalize_run(args, all_scene_metrics, num_to_process, start_time) if __name__ == '__main__': args = parse_args() main(args)