''' BVH, Nov 2025. Shared utility functions for Any4D inference and Gradio app. This module contains helper functions that are used by both: - infer.py: High-quality batch inference script - gradio_app.py: Interactive Gradio web interface ''' import attrs import os import sys # Library imports import copy import cv2 import glob import imageio import importlib import lovely_tensors import numpy as np import pathlib import time import torch import warnings import zlib from einops import rearrange from lovely_numpy import lo from PIL import Image from rich import print # Internal imports from cosmos_predict2.models.video2world_model import Predict2ModelManagerConfig from custom.utils.io import JsonNumpyEncoder, load_json, save_json lovely_tensors.monkey_patch() np.set_printoptions(precision=3, suppress=True) warnings.filterwarnings('ignore', category=FutureWarning) class DictToObject: '''Simple wrapper to access dict values as attributes.''' def __init__(self, d): self.__dict__.update(d) def derive_scene_seed(seed, data_batch, fallback_id): ''' Per-scene seed for paired / reproducible inference: (seed + crc32(sample_id)) % 10M. Derived from sample_id (content-stable identity), NOT the loop index, because webdataset scene order varies across runs/shards -- this keeps seeds paired across checkpoints and shard layouts. Shared by infer_anydata and infer_vidar (must stay identical for pairing). None seed -> None. ''' if seed is None: return None from custom.eval.visuals import unpack_nested sid = unpack_nested(data_batch, 'sample_id', dtype=str, optional=True) or str(fallback_id) return (seed + zlib.crc32(sid.encode())) % 10_000_000 def load_experiment_config(experiment_config_path): ''' Load experiment configuration from python file. :param experiment_config_path (str): Path to experiment config file :return (dict): Experiment configuration dictionary ''' print(f'[cyan]Loading experiment config from {experiment_config_path}...') # Import the experiment config module config_path = experiment_config_path.replace('.py', '').replace('/', '.') if config_path.startswith('.'): config_path = config_path[1:] config_module = importlib.import_module(config_path) # Extract relevant parts any4d_config = config_module.any4d_config model_config = dict(config_module.model) # Make a copy to safely override training-specific settings for inference without modifying the original module config # Override training-specific config values for inference # skip_first_validation should always be False (not True or 'delay') if model_config.get('skip_first_validation') in [True, 'delay']: print(f'[yellow]Overriding skip_first_validation from {model_config["skip_first_validation"]} to False for inference') model_config['skip_first_validation'] = False # max_iter should be low for inference (not 50000+ for training) if model_config.get('max_iter', 0) > 10: print(f'[yellow]Overriding max_iter from {model_config["max_iter"]} to 1 for inference') model_config['max_iter'] = 1 print(f'[green]Loaded experiment: {config_module.job["name"]}') print(f'[green]Any4D config keys: {list(any4d_config.keys())}') return { 'any4d_config': any4d_config, 'model': model_config, 'job': config_module.job, } def build_model(experiment_config, ckpt_path=None, num_steps=None, guidance_scale=None, sigma_max=None, sigma_min=None, cfg_zero_entries=None): ''' Build model instance from experiment configuration. :param experiment_config (dict): Experiment configuration dictionary :param ckpt_path (str): Optional checkpoint path override :param num_steps (int): Optional number of diffusion steps override :param guidance_scale (float): Optional CFG scale override :param sigma_max (float): Optional maximum noise level override :param sigma_min (float): Optional minimum noise level override :param cfg_zero_entries (str): Optional comma-separated list of entries to zero for CFG :return (tuple): (model, config) ''' print(f'[cyan]Building model architecture...') # Import model class and configs from custom.any4d.a4d_model import Any4DModel, VidarModel from custom.any4d.a4d_config import Any4DConfig from cosmos_predict2.configs.base.defaults.model import ANY4D_PIPELINE_2B # Build Any4D config any4d_cfg = experiment_config['any4d_config'] model_cfg = experiment_config['model'] # Determine if this is Any4D or just Vidar is_any4d = any4d_cfg.get('any4d_active', True) # Get checkpoint path (prioritize argument over config) ckpt_path = ckpt_path if ckpt_path else model_cfg.get('dit_path') # Create a copy of the pipeline config with guardrails disabled for inference import copy import torch.nn from cosmos_predict2.configs.base.config_video2world import CosmosGuardrailConfig, CosmosReason1Config pipe_cfg = copy.deepcopy(ANY4D_PIPELINE_2B) pipe_cfg.guardrail_config = CosmosGuardrailConfig(checkpoint_dir='', enabled=False) pipe_cfg.prompt_refiner_config = CosmosReason1Config(checkpoint_dir='', enabled=False) # Ensure extra_nets is not None (it may be None in the default config) if pipe_cfg.extra_nets is None: pipe_cfg.extra_nets = torch.nn.ModuleDict() # Create full model config dict (don't include job here - will add later) config_dict = { 'vidar_active': True, 'any4d_active': is_any4d, **any4d_cfg, 'pipe_config': pipe_cfg, 'model_manager_config': Predict2ModelManagerConfig( dit_path=ckpt_path, text_encoder_path=model_cfg.get('text_encoder_path', ''), vae_path=model_cfg.get('vae_path'), tokenizer_chunk_duration=model_cfg.get('tokenizer_chunk_duration', 41), ), } # Override sampling parameters if provided if num_steps is not None: config_dict['val_num_steps'] = num_steps if guidance_scale is not None: config_dict['val_cfg_scale'] = guidance_scale if sigma_max is not None: config_dict['val_sigma_max'] = sigma_max if sigma_min is not None: config_dict['val_sigma_min'] = sigma_min if cfg_zero_entries is not None: config_dict['val_cfg_zero_entries'] = cfg_zero_entries.split(',') if isinstance(cfg_zero_entries, str) else cfg_zero_entries # Filter out keys not recognized by Any4DConfig (e.g. vidar2a4d, transforms) valid_fields = {a.name for a in attrs.fields(Any4DConfig)} config_dict = {k: v for k, v in config_dict.items() if k in valid_fields} # Create config object model_config = Any4DConfig(**config_dict) # Add job info as OmegaConf DictConfig (compatible with omegaconf validation) from omegaconf import OmegaConf job_dict = experiment_config.get('job', {'name': 'inference'}) if 'undated_name' not in job_dict: job_dict['undated_name'] = job_dict.get('name', 'inference') object.__setattr__(model_config, 'job', OmegaConf.create(job_dict)) # Create model instance (Any4D or Vidar depending on config) if is_any4d: model = Any4DModel(model_config) print(f'[green]Any4D model architecture built') else: model = VidarModel(model_config) print(f'[green]Vidar model architecture built') return model, model_config def setup_model(model, config, ckpt_path=None): ''' Initialize all model modules (DiT weights, VAE, text encoder, transforms, etc.). The checkpoint path should already be configured via model_manager_config.dit_path during build_model(), but can be overridden here. :param model: Model instance :param config: Model configuration :param ckpt_path (str): Optional checkpoint path override :return: Model with all modules initialized ''' ckpt_path = ckpt_path if ckpt_path else config.model_manager_config.dit_path print(f'[cyan]Setting up model modules (weights from {ckpt_path})...') # Handle S3 paths or wildcards if '*' in ckpt_path and not ckpt_path.startswith('s3://'): matching_ckpts = sorted(glob.glob(ckpt_path)) if matching_ckpts: ckpt_path = matching_ckpts[-1] print(f'[yellow]Using latest matching weights: {ckpt_path}') # Initialize all model modules (loads DiT weights, VAE, text encoder, transforms) model.setup_modules() print(f'[green]Model setup complete') return model def load_model_bundle(experiment_config_path, ckpt_path=None, device='cuda', num_steps=None, guidance_scale=None): ''' Load complete model bundle (config + model + checkpoint). This is the main entry point for loading a trained model. :param experiment_config_path (str): Path to experiment config file :param ckpt_path (str): Optional checkpoint path override :param device (str): Device to load model on :param num_steps (int): Optional number of diffusion steps override :param guidance_scale (float): Optional CFG scale override :return (tuple): (model, config, experiment_config) ''' # Load experiment configuration experiment_config = load_experiment_config(experiment_config_path) # Build model model, config = build_model(experiment_config, ckpt_path=ckpt_path, num_steps=num_steps, guidance_scale=guidance_scale) model = model.to(device) model.device = device # Initialize all model modules (DiT weights, VAE, text encoder, etc.) model = setup_model(model, config, ckpt_path=ckpt_path) model.eval() return model, config, experiment_config def save_video(dst_fp, frames, fps=10, quality=8): ''' Save video frames to MP4 file. :param dst_fp (str): Destination file path :param frames: Frames to save (T, H, W, C) array or tensor :param fps (int): Frame rate :param quality (int): Quality level (0-10) ''' if torch.is_tensor(frames): frames = frames.detach().cpu().numpy() # Handle (T, C, H, W) -> (T, H, W, C) if frames.ndim == 4 and frames.shape[1] in [1, 3, 4]: frames = rearrange(frames, 'T C H W -> T H W C') # Squeeze single-channel images if frames.shape[-1] == 1: frames = frames.squeeze(-1) # Convert float to uint8 if needed if frames.dtype.kind == 'f': frames = (frames * 255.0).astype(np.uint8) # Ensure frames is a list for imageio if isinstance(frames, np.ndarray): frames = list(frames) # Create parent directory os.makedirs(os.path.dirname(dst_fp), exist_ok=True) # Save video imageio.mimwrite(dst_fp, frames, format='ffmpeg', fps=float(fps), macro_block_size=8, quality=int(quality)) def download_video_from_s3(s3_path, local_cache_dir='/tmp/video_cache'): ''' Download video from S3 to local cache. :param s3_path (str): S3 path (s3://bucket/path/to/video.mp4) :param local_cache_dir (str): Local directory for caching :return (str): Local path to downloaded video ''' import boto3 import hashlib # Parse S3 path if not s3_path.startswith('s3://'): raise ValueError(f'Invalid S3 path: {s3_path}') s3_path_clean = s3_path[5:] # Remove 's3://' bucket, key = s3_path_clean.split('/', 1) # Create cache directory os.makedirs(local_cache_dir, exist_ok=True) # Create local filename from S3 key hash key_hash = hashlib.md5(key.encode()).hexdigest()[:8] file_ext = os.path.splitext(key)[1] local_path = os.path.join(local_cache_dir, f's3_{key_hash}{file_ext}') # Download if not cached if not os.path.exists(local_path): print(f'[cyan]Downloading from S3: {s3_path}') s3 = boto3.client('s3') s3.download_file(bucket, key, local_path) print(f'[green]Downloaded to: {local_path}') else: print(f'[green]Using cached video: {local_path}') return local_path def download_video_from_youtube(youtube_url, local_cache_dir='/tmp/video_cache'): ''' Download video from YouTube to local cache. :param youtube_url (str): YouTube URL :param local_cache_dir (str): Local directory for caching :return (str): Local path to downloaded video ''' try: import yt_dlp except ImportError: print('[red]Error: yt-dlp not installed. Install with: pip install yt-dlp') return None import hashlib # Create cache directory os.makedirs(local_cache_dir, exist_ok=True) # Create local filename from URL hash url_hash = hashlib.md5(youtube_url.encode()).hexdigest()[:8] local_path = os.path.join(local_cache_dir, f'yt_{url_hash}.mp4') # Download if not cached if not os.path.exists(local_path): print(f'[cyan]Downloading from YouTube: {youtube_url}') ydl_opts = { 'format': 'best[ext=mp4]/best', 'outtmpl': local_path, 'quiet': False, 'no_warnings': False, } try: with yt_dlp.YoutubeDL(ydl_opts) as ydl: ydl.download([youtube_url]) print(f'[green]Downloaded to: {local_path}') except Exception as e: print(f'[red]Error downloading from YouTube: {e}') return None else: print(f'[green]Using cached video: {local_path}') return local_path def load_video_from_source(video_source, target_T=None, target_H=None, target_W=None): ''' Load video from various sources: local file, S3, or YouTube. :param video_source (str): Video source (local path, s3:// URL, or YouTube URL) :param target_T (int): Optional target number of frames :param target_H (int): Optional target height :param target_W (int): Optional target width :return: (T, C, H, W) tensor of float32 in [0, 1] ''' # Determine source type and get local path if video_source.startswith('s3://'): local_path = download_video_from_s3(video_source) elif 'youtube.com' in video_source or 'youtu.be' in video_source: local_path = download_video_from_youtube(video_source) if local_path is None: return None else: # Assume local file local_path = video_source if not os.path.exists(local_path): print(f'[red]Error: Video file not found: {local_path}') return None # Load video from local path return load_video_from_file(local_path, target_T, target_H, target_W) def load_data_from_tarfile(tar_path, num_timesteps=41, cam_indices=[0, 1], remap_camera_indices=False): ''' Load RGB images and camera data from a tarfile (vidar format). Similar to dvs7.py tarfile loading. :param tar_path (str): Path to tar archive (local or s3://) :param num_timesteps (int): Number of timesteps to load :param cam_indices (list): List of camera indices to load :param remap_camera_indices (bool): If True, remap camera indices to [0, 1] (vidar convention) :return: dict with 'rgb', 'pose', 'intrinsics' in vidar format ''' import io import tarfile from pathlib import Path from PIL import Image print(f'[cyan]Loading data from tarfile: {tar_path}') # Handle S3 paths - download to tmp first if tar_path.startswith('s3://'): print(f'[cyan]Downloading tarfile from S3...') tar_path = download_video_from_s3(tar_path, local_cache_dir='/tmp/tarfile_cache') print(f'[green]Downloaded to: {tar_path}') p = Path(tar_path) if not p.exists(): print(f'[red]Error: Tarfile not found: {tar_path}') return None try: tar = tarfile.open(str(p), 'r') all_files = tar.getnames() # Helper function to find files matching pattern def find_files_in_tar(all_files, pattern, cam_idx): '''Find files matching pattern for specific camera''' matching = [] for fname in all_files: if f',{cam_idx})' in fname and pattern in fname: # Extract time index from filename # Format: ...pattern.(time_idx,cam_idx).ext try: parts = fname.split(f',{cam_idx})') before_cam = parts[0] if '(' in before_cam: time_str = before_cam.split('(')[-1] time_idx = int(time_str) matching.append((time_idx, fname)) except: continue matching.sort() return matching # Load RGB images for all cameras rgb_dict = {} for cam_idx in cam_indices: rgb_files = find_files_in_tar(all_files, 'rgb', cam_idx) print(f'[cyan] Found {len(rgb_files)} RGB files for camera {cam_idx}') for time_idx, fname in rgb_files[:num_timesteps]: member = tar.getmember(fname) img_bytes = tar.extractfile(member).read() # Load image from bytes img = Image.open(io.BytesIO(img_bytes)) img_array = np.array(img).astype(np.float32) / 255.0 # (H, W, 3) in [0, 1] # Convert to tensor (H, W, C) -> (C, H, W) img_tensor = torch.from_numpy(img_array).permute(2, 0, 1) # (3, H, W) # Add batch dimension and store rgb_dict[(time_idx, cam_idx)] = img_tensor.unsqueeze(0) # (1, 3, H, W) # Load camera poses for all cameras pose_dict = {} intrinsics_dict = {} for cam_idx in cam_indices: # Find pose files pose_files = find_files_in_tar(all_files, 'pose', cam_idx) if pose_files: # Load first pose file which contains all timesteps first_time_idx, first_file = pose_files[0] member = tar.getmember(first_file) f = tar.extractfile(member) data = np.load(io.BytesIO(f.read()), allow_pickle=True) pose_data = data['data'].item() # Dict with (time, cam) keys # Extract poses for this camera keys = sorted([k for k in pose_data.keys() if k[1] == cam_idx]) print(f'[cyan] Found {len(keys)} poses for camera {cam_idx}') for key in keys[:num_timesteps]: pose_matrix = torch.from_numpy(pose_data[key]).float() # (4, 4) pose_dict[key] = pose_matrix.unsqueeze(0) # (1, 4, 4) # Find intrinsics files intrinsics_files = find_files_in_tar(all_files, 'intrinsics', cam_idx) if intrinsics_files: # Load first intrinsics file first_time_idx, first_file = intrinsics_files[0] member = tar.getmember(first_file) f = tar.extractfile(member) data = np.load(io.BytesIO(f.read()), allow_pickle=True) intrinsics_data = data['data'].item() # Dict with (time, cam) keys # Extract intrinsics for this camera keys = sorted([k for k in intrinsics_data.keys() if k[1] == cam_idx]) print(f'[cyan] Found {len(keys)} intrinsics for camera {cam_idx}') for key in keys[:num_timesteps]: intrinsics_matrix = torch.from_numpy(intrinsics_data[key]).float() # (4, 4) intrinsics_dict[key] = intrinsics_matrix.unsqueeze(0) # (1, 4, 4) tar.close() print(f'[green]Loaded {len(rgb_dict)} RGB frames, {len(pose_dict)} poses, {len(intrinsics_dict)} intrinsics from tarfile') # Remap time indices to start at 0 (needed for convert_absolute_to_vidar_format) # Find minimum time index across all dictionaries all_time_indices = set() for key in rgb_dict.keys(): all_time_indices.add(key[0]) for key in pose_dict.keys(): all_time_indices.add(key[0]) for key in intrinsics_dict.keys(): all_time_indices.add(key[0]) if all_time_indices: min_time = min(all_time_indices) if min_time != 0: print(f'[cyan]Remapping time indices: min_time={min_time} -> 0') # Remap RGB dict rgb_dict_remapped = {} for (t, c), value in rgb_dict.items(): rgb_dict_remapped[(t - min_time, c)] = value rgb_dict = rgb_dict_remapped # Remap pose dict pose_dict_remapped = {} for (t, c), value in pose_dict.items(): pose_dict_remapped[(t - min_time, c)] = value pose_dict = pose_dict_remapped # Remap intrinsics dict intrinsics_dict_remapped = {} for (t, c), value in intrinsics_dict.items(): intrinsics_dict_remapped[(t - min_time, c)] = value intrinsics_dict = intrinsics_dict_remapped print(f'[green]Remapped {len(all_time_indices)} time indices to start at 0') # Remap camera indices to VIDAR convention (0=output, 1=input) if requested if remap_camera_indices and len(cam_indices) == 2: # Create mapping: cam_indices[0] -> 0 (output), cam_indices[1] -> 1 (input) cam_mapping = {cam_indices[0]: 0, cam_indices[1]: 1} print(f'[cyan]Remapping camera indices: {cam_indices[0]}->0 (out), {cam_indices[1]}->1 (in)') # Remap RGB dict rgb_dict_remapped = {} for (t, c), value in rgb_dict.items(): if c in cam_mapping: rgb_dict_remapped[(t, cam_mapping[c])] = value rgb_dict = rgb_dict_remapped # Remap pose dict pose_dict_remapped = {} for (t, c), value in pose_dict.items(): if c in cam_mapping: pose_dict_remapped[(t, cam_mapping[c])] = value pose_dict = pose_dict_remapped # Remap intrinsics dict intrinsics_dict_remapped = {} for (t, c), value in intrinsics_dict.items(): if c in cam_mapping: intrinsics_dict_remapped[(t, cam_mapping[c])] = value intrinsics_dict = intrinsics_dict_remapped print(f'[green]Remapped camera indices to VIDAR convention (0=out, 1=in)') return { 'rgb': rgb_dict, 'pose': pose_dict, 'intrinsics': intrinsics_dict, } except Exception as e: print(f'[red]Error loading tarfile: {e}') import traceback traceback.print_exc() return None def load_video_from_file(video_path, target_T=None, target_H=None, target_W=None): ''' Load video frames from MP4 file. :param video_path (str): Path to video file :param target_T (int): Optional target number of frames :param target_H (int): Optional target height :param target_W (int): Optional target width :return: (T, C, H, W) tensor of float32 in [0, 1] ''' # Clean and validate video path video_path = str(video_path).strip() print(f'[yellow]Loading video from {video_path}') # Check if file exists if not os.path.exists(video_path): print(f'[red]Error: Video file does not exist: {video_path}') return None # Check if file is readable if not os.path.isfile(video_path): print(f'[red]Error: Path is not a file: {video_path}') return None # Try reading with imageio first (more robust for various formats) try: import imageio print('[cyan]Attempting to read video with imageio...') reader = imageio.get_reader(video_path, 'ffmpeg') frames = [] for frame in reader: # Convert to RGB float32 in [0, 1] if frame.dtype == np.uint8: frame = frame.astype(np.float32) / 255.0 elif frame.dtype == np.uint16: frame = frame.astype(np.float32) / 65535.0 else: frame = frame.astype(np.float32) # Ensure RGB (handle grayscale or RGBA) if len(frame.shape) == 2: frame = np.stack([frame] * 3, axis=-1) elif frame.shape[-1] == 4: frame = frame[:, :, :3] frames.append(frame) reader.close() if len(frames) > 0: print(f'[green]Successfully read {len(frames)} frames with imageio') except Exception as e: print(f'[yellow]imageio failed ({e}), trying OpenCV...') frames = [] # Fallback to OpenCV cap = cv2.VideoCapture(video_path) # Check if video opened successfully if not cap.isOpened(): print(f'[red]Error: OpenCV could not open video file: {video_path}') cap.release() return None # Get video properties frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) fps = cap.get(cv2.CAP_PROP_FPS) print(f'[cyan]Video properties: {frame_count} frames, {fps:.2f} fps') while True: ret, frame = cap.read() if not ret: break # Convert BGR to RGB and normalize to [0, 1] frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frame = frame.astype(np.float32) / 255.0 frames.append(frame) cap.release() if len(frames) == 0: print(f'[red]Error: No frames read from {video_path}') print(f'[yellow]File size: {os.path.getsize(video_path)} bytes') return None print(f'[green]Loaded {len(frames)} frames from video') # Convert to tensor (T, H, W, C) -> (T, C, H, W) frames_array = np.stack(frames, axis=0) frames_tensor = torch.from_numpy(frames_array).permute(0, 3, 1, 2) # Resize if targets specified if target_H is not None and target_W is not None: import torch.nn.functional as F frames_tensor = F.interpolate( frames_tensor, size=(int(target_H), int(target_W)), mode='bilinear', align_corners=False ) # Trim or repeat frames to match target length if target_T is not None: if frames_tensor.shape[0] > target_T: frames_tensor = frames_tensor[:target_T] elif frames_tensor.shape[0] < target_T: # Repeat last frame last_frame = frames_tensor[-1:].repeat(target_T - frames_tensor.shape[0], 1, 1, 1) frames_tensor = torch.cat([frames_tensor, last_frame], dim=0) return frames_tensor def create_input_from_video(raw_video_path, num_frames=14, frame_width=512, frame_height=288, center_crop=True, frame_stride=1, frame_offset=0): ''' Process raw video file into model input format. :param raw_video_path (str): Path to raw video file :param num_frames (int): Number of frames to extract :param frame_width (int): Target width :param frame_height (int): Target height :param center_crop (bool): Whether to center crop :param frame_stride (int): Frame stride for subsampling :param frame_offset (int): Starting frame offset :return: (T, C, H, W) tensor of float32 in [0, 1] ''' # Load video frames = load_video_from_file(raw_video_path) if frames is None: return None # Apply frame stride and offset clip_frames = np.arange(num_frames) * frame_stride + frame_offset clip_frames = clip_frames[clip_frames < frames.shape[0]] if len(clip_frames) < num_frames: print(f'[yellow]Warning: Only {len(clip_frames)} frames available, padding to {num_frames}') # Pad with last frame last_idx = clip_frames[-1] if len(clip_frames) > 0 else 0 while len(clip_frames) < num_frames: clip_frames = np.append(clip_frames, last_idx) frames = frames[clip_frames] # Center crop if requested if center_crop: _, _, H, W = frames.shape aspect_target = frame_width / frame_height aspect_current = W / H if aspect_current > aspect_target: # Too wide, crop width new_W = int(H * aspect_target) left = (W - new_W) // 2 frames = frames[:, :, :, left:left+new_W] else: # Too tall, crop height new_H = int(W / aspect_target) top = (H - new_H) // 2 frames = frames[:, :, top:top+new_H, :] # Resize to target resolution import torch.nn.functional as F frames = F.interpolate( frames, size=(frame_height, frame_width), mode='bilinear', align_corners=False ) return frames def get_random_stub_video(test_stub): ''' Get a random test video path based on test_stub parameter. Searches common locations for test videos. :param test_stub (int): Test stub ID (0 = disabled, 1+ = various test videos) :return (str): Path to test video file, or None if disabled ''' if test_stub == 0: return None # Define potential test video search paths search_paths = [ '/datasets/basile/test_videos/**/*.mp4', '/datasets/basile/kubric/**/*.mp4', '/datasets/basile/**/*sample*.mp4', './test_videos/**/*.mp4', '../test_videos/**/*.mp4', ] # Find all available test videos test_videos = [] for search_pattern in search_paths: matches = glob.glob(search_pattern, recursive=True) test_videos.extend(matches) # Remove duplicates and sort test_videos = sorted(list(set(test_videos))) if len(test_videos) == 0: print('[yellow]Warning: No test videos found. Will skip test inference.') return None print(f'[cyan]Found {len(test_videos)} test videos') # If specific stub requested, use modulo to cycle through available videos if test_stub > 0: idx = (test_stub - 1) % len(test_videos) selected = test_videos[idx] print(f'[cyan]Selected test video: {selected}') return selected return None def calculate_metrics_per_scene(all_samples, data_batch, config, lpips_loss, model, device, metric_resolution=None): ''' Calculate PSNR, SSIM, LPIPS metrics for a scene across multiple samples, per rgb view. Decodes latent predictions to pixel space using the VAE. Scores every tracked rgb entry: config.track_metrics (rgb* keys) when given, else a dynamic fallback = every rgb* present in the predictions. Mirrors calculate_metrics_a4d's has_output skip (fully-conditioned entries with an all-zero output_mask are dropped). The native-res 11-view block in each sample's info.json is the default source of truth; this in-code path is only the optional --infer_metrics one and is the only place metric_resolution applies. :param all_samples: List of output dictionaries from multiple samples :param data_batch: Input data batch with ground truth :param config: Any4DConfig (for track_metrics); None => dynamic rgb* discovery :param lpips_loss: LPIPS loss module :param model: Model with VAE for decoding :param device: Device for computation :param metric_resolution: Optional anydata resize spec (e.g. [-16, 512]) applied to BOTH pred and gt before scoring, for fair comparison across checkpoints with different native resolutions. :return: Dict {sample_id, psnr:{rgb0..rgbN, avg}, ssim:{...}, lpips:{...}, *_per_sample:{per entry}} ''' import skimage.metrics from custom.eval.visuals import cached_decode_latent_video_auto, unpack_nested a4d_vae = model.vae def score_pair(pred_b, gt_b): ''' Per-frame PSNR/SSIM/LPIPS averaged over non-near-black frames; None if none valid. pred_b, gt_b are (3, T, H, W) in [0, 1]. ''' # Optional: score pred & gt at a common resolution (fair cross-model-resolution compare). # Reuses the anydata augmentation.resize spec (e.g. [-16, 512] = longest side 512, /16 stride). # Bilinear stays in [0, 1] so SSIM data_range=1.0 and LPIPS *2-1 remain valid; same target # dims for both -> fair by construction. if metric_resolution is not None: import torch.nn.functional as F from anydata.augmentations.resize import parse_resize_params H0, W0 = pred_b.shape[-2:] th, tw = parse_resize_params((W0, H0), metric_resolution) # -> (H, W) def _rs(v): v = F.interpolate(v.permute(1, 0, 2, 3).float(), size=(th, tw), mode='bilinear', align_corners=False) return v.permute(1, 0, 2, 3).clamp(0.0, 1.0) pred_b, gt_b = _rs(pred_b), _rs(gt_b) T = pred_b.shape[1] frame_psnr, frame_ssim, frame_lpips = [], [], [] for t in range(T): pred_frame = pred_b[:, t] # (3, H, W) gt_frame = gt_b[:, t] # (3, H, W) # Skip near-black frames if gt_frame.mean() < 0.01: continue # PSNR mse = torch.mean((pred_frame - gt_frame) ** 2) frame_psnr.append(-10.0 * torch.log10(mse).item() if mse > 0 else 100.0) # SSIM (convert to float32 first, as numpy doesn't support bfloat16) pred_np = pred_frame.float().permute(1, 2, 0).cpu().numpy() # (H, W, 3) gt_np = gt_frame.float().permute(1, 2, 0).cpu().numpy() # (H, W, 3) frame_ssim.append(skimage.metrics.structural_similarity( pred_np, gt_np, data_range=1.0, channel_axis=2)) # LPIPS (expects (1, 3, H, W) in [-1, 1]) pred_lpips = pred_frame.unsqueeze(0).to(device) * 2.0 - 1.0 gt_lpips = gt_frame.unsqueeze(0).to(device) * 2.0 - 1.0 with torch.no_grad(): frame_lpips.append(lpips_loss(pred_lpips, gt_lpips).item()) if len(frame_psnr) == 0: return None return float(np.mean(frame_psnr)), float(np.mean(frame_ssim)), float(np.mean(frame_lpips)) # Which rgb entries to score: config.track_metrics (rgb* only) if available, else dynamic # discovery of every rgb* present in the predictions. Sorted by view index (rgb0..rgb10). if config is not None and getattr(config, 'track_metrics', None): entry_names = [k for k in config.track_metrics if k.startswith('rgb')] else: entry_names = sorted( {k for s in all_samples for k in s.get('y0_pred_entries', {}) if k.startswith('rgb')}, key=lambda s: int(s[3:])) psnr, ssim, lpips_m = {}, {}, {} psnr_ps, ssim_ps, lpips_ps = {}, {}, {} for k in entry_names: sp, ss, sl = [], [], [] for sample in all_samples: yp = sample.get('y0_pred_entries', {}) x0 = sample.get('x0_entries', {}) if k not in yp or k not in x0: continue # entry tracked but absent from this sample # Skip fully-conditioned entries (no predicted output), mirroring calculate_metrics_a4d. omask = x0.get(k + '_output_mask') if omask is not None and float(omask.sum()) <= 0.0: continue # Decode to pixel space: (B, 3, Tp, Hp, Wp) in [0, 1]; score the first batch element. pred = cached_decode_latent_video_auto(yp, k, a4d_vae) gt = cached_decode_latent_video_auto(x0, k, a4d_vae) scored = score_pair(pred[0].cpu(), gt[0].cpu()) if scored is None: continue sp.append(scored[0]) ss.append(scored[1]) sl.append(scored[2]) if len(sp) == 0: continue psnr[k] = float(np.mean(sp)) ssim[k] = float(np.mean(ss)) lpips_m[k] = float(np.mean(sl)) psnr_ps[k], ssim_ps[k], lpips_ps[k] = sp, ss, sl # avg = mean over the scored views (NO bare scalar). if psnr: psnr['avg'] = float(np.mean(list(psnr.values()))) ssim['avg'] = float(np.mean(list(ssim.values()))) lpips_m['avg'] = float(np.mean(list(lpips_m.values()))) scene_metrics = { 'sample_id': unpack_nested(data_batch, 'sample_id', dtype=str, optional=True), 'psnr': psnr, 'ssim': ssim, 'lpips': lpips_m, 'psnr_per_sample': psnr_ps, 'ssim_per_sample': ssim_ps, 'lpips_per_sample': lpips_ps, } return scene_metrics def calculate_uncertainty_heatmaps(all_samples, entry_names=['rgb0'], a4d_vae=None): ''' Calculate spatial uncertainty/diversity heatmap across multiple samples. Process: 1. Decode each sample's latent predictions to pixel space individually 2. Calculate per-pixel standard deviation across all decoded samples 3. Overlay uncertainty heatmap on darkened sample 1 RGB video :param all_samples: List of output dictionaries from multiple samples :param entry_names: List of spatial entry names to calculate uncertainty for :param a4d_vae: VAE for decoding latents to pixel space :return: Tuple of (uncertainty_videos, diversity_metrics) ''' print('[cyan]Calculating uncertainty heatmaps in pixel space...') from custom.eval.visuals import cached_decode_latent_video_auto import matplotlib.pyplot as plt uncertainty_videos = {} diversity_metrics = {} for entry_name in entry_names: # Step 1: Decode each sample to pixel space individually predictions_pixel = [] # Will hold decoded samples in pixel space gt_pixel = None print(f'[cyan] Decoding {len(all_samples)} samples for {entry_name}...') for sample_idx, sample in enumerate(all_samples): # Check if entry exists in latent space if entry_name not in sample.get('y0_pred_entries', {}): continue # Decode this sample's latent prediction to pixel space # Input: (B, 16, Tl, Hl, Wl) in latent space # Output: (B, 3, Tp, Hp, Wp) in pixel space [0, 1] pred_decoded = cached_decode_latent_video_auto( sample['y0_pred_entries'], entry_name, a4d_vae) predictions_pixel.append(pred_decoded[0].cpu().float()) # (3, T, H, W) # Get GT video in pixel space (same for all samples, only need once) if gt_pixel is None and entry_name in sample.get('x0_entries', {}): gt_decoded = cached_decode_latent_video_auto( sample['x0_entries'], entry_name, a4d_vae) gt_pixel = gt_decoded[0].cpu().float() # (3, T, H, W) in [0, 1] if len(predictions_pixel) == 0: print(f'[yellow]No predictions found for {entry_name}, skipping uncertainty') continue print(f'[cyan] Calculating pixel-space uncertainty across {len(predictions_pixel)} samples...') # Step 2: Calculate uncertainty in pixel space # Stack all decoded predictions: (S, 3, T, H, W) where S = num_samples predictions_pixel = torch.stack(predictions_pixel, dim=0) # Calculate per-pixel standard deviation across samples # (S, 3, T, H, W) -> std(dim=0) -> (3, T, H, W) # Then average across RGB channels -> (T, H, W) uncertainty_pixel = predictions_pixel.std(dim=0).mean(dim=0) # (T, H, W) # Calculate mean diversity per frame frame_diversity = uncertainty_pixel.mean(dim=(1, 2)) # (T,) mean_diversity = float(frame_diversity.mean()) # Step 3: Create uncertainty heatmap visualization # uncertainty_pixel is already (T, H, W) in pixel space uncertainty_map = uncertainty_pixel # (T, H, W) # Normalize to [0, 1] for better visualization uncertainty_min = uncertainty_map.min() uncertainty_max = uncertainty_map.max() if uncertainty_max > uncertainty_min: uncertainty_normalized = (uncertainty_map - uncertainty_min) / (uncertainty_max - uncertainty_min) else: uncertainty_normalized = torch.zeros_like(uncertainty_map) # Apply magma colormap to uncertainty (normalize per-frame for better visualization) cmap = plt.get_cmap('magma') uncertainty_colored_frames = [] for t in range(uncertainty_map.shape[0]): uncertainty_frame = uncertainty_map[t].numpy() # (H, W) # Normalize each frame independently to [0, 1] for better visualization frame_min = uncertainty_frame.min() frame_max = uncertainty_frame.max() if frame_max > frame_min: uncertainty_frame_norm = (uncertainty_frame - frame_min) / (frame_max - frame_min) else: uncertainty_frame_norm = np.zeros_like(uncertainty_frame) # Apply colormap: (H, W) -> (H, W, 4), keep RGB channels uncertainty_rgb = cmap(uncertainty_frame_norm)[:, :, :3] # (H, W, 3) in [0, 1] uncertainty_colored_frames.append(uncertainty_rgb) uncertainty_colored = np.stack(uncertainty_colored_frames, axis=0) # (T, H, W, 3) # Create overlay: darken sample 1 RGB and blend with uncertainty heatmap if len(predictions_pixel) > 0: # Convert first prediction (sample 1) to (T, H, W, 3) format (handle BFloat16) sample1_rgb = rearrange(predictions_pixel[0].float(), 'C T H W -> T H W C').cpu().numpy() # (T, H, W, 3) sample1_rgb = np.clip(sample1_rgb, 0.0, 1.0) # Darken the sample 1 RGB (multiply by 0.5) sample1_darkened = sample1_rgb * 0.5 # Blend: use uncertainty alpha from magma (high uncertainty = more colorful overlay) uncertainty_overlay = sample1_darkened * 0.6 + uncertainty_colored * 0.6 uncertainty_overlay = np.clip(uncertainty_overlay, 0.0, 1.0) uncertainty_videos[entry_name] = uncertainty_overlay else: # Fallback: just use uncertainty heatmap without sample 1 overlay print(f'[yellow]No predictions found for {entry_name}, using plain uncertainty heatmap') uncertainty_videos[entry_name] = uncertainty_colored diversity_metrics[entry_name] = { 'frame_diversity': frame_diversity.numpy().tolist(), 'mean_diversity': mean_diversity, } print(f'[green]{entry_name}: mean diversity = {mean_diversity:.4f}') return uncertainty_videos, diversity_metrics def save_uncertainty_videos(uncertainty_videos, output_dir, dset_name, sample_id_clean, fps=10, iteration=None): ''' Save uncertainty heatmap videos to disk. :param uncertainty_videos: Dictionary mapping entry names to uncertainty videos :param output_dir: Base output directory :param dset_name: Dataset name :param sample_id_clean: Cleaned sample ID for filename :param fps: Frame rate for video (extracted from data_batch) :param iteration: Optional iteration index to include in filename ''' if len(uncertainty_videos) == 0: return uncertainty_dir = pathlib.Path(output_dir) uncertainty_dir.mkdir(parents=True, exist_ok=True) import imageio for entry_name, uncertainty_video in uncertainty_videos.items(): # Save as MP4 iter_str = f'{iteration:04d}_' if iteration is not None else '' uncertainty_file = uncertainty_dir / f'uncert_{iter_str}{dset_name}_{entry_name}_{sample_id_clean}.mp4' # Convert to uint8 uncertainty_uint8 = (uncertainty_video * 255).astype(np.uint8) # Save video imageio.mimwrite(str(uncertainty_file), uncertainty_uint8, format='ffmpeg', fps=float(fps), quality=8) print(f'[green]Saved uncertainty video for {entry_name} to {uncertainty_file}') def pad_video_to_size(video, target_H, target_W): ''' Pad video with zeros to match target height and width. :param video: Input video (T, H, W, C) :param target_H: Target height :param target_W: Target width :return: Padded video (T, target_H, target_W, C) ''' T, H, W, C = video.shape if H == target_H and W == target_W: return video padded = np.zeros((T, target_H, target_W, C), dtype=video.dtype) padded[:, :H, :W, :] = video print(f'[cyan]Padded video from ({T}, {H}, {W}, {C}) to ({T}, {target_H}, {target_W}, {C})') return padded def concatenate_videos_horizontal_with_padding(videos, verbose=True): ''' Concatenate videos horizontally (along width axis) with zero padding for mismatched heights. All videos must have same number of frames (T) and channels (C). :param videos: List of videos (T, H, W, C) :param verbose: Whether to print debug info :return: Concatenated video (T, max_H, sum_W, C) ''' if len(videos) == 0: return None # Find maximum height max_H = max(v.shape[1] for v in videos) # Pad all videos to max height padded_videos = [] for video in videos: padded_videos.append(pad_video_to_size(video, target_H=max_H, target_W=video.shape[2])) # Concatenate along width axis result = np.concatenate(padded_videos, axis=2) # (T, max_H, sum_W, C) if verbose: print(f'[green]Concatenated {len(videos)} videos horizontally: {result.shape}') return result def concatenate_videos_vertical_with_padding(videos, verbose=True): ''' Concatenate videos vertically (along height axis) with zero padding for mismatched widths. All videos must have same number of frames (T) and channels (C). :param videos: List of videos (T, H, W, C) :param verbose: Whether to print debug info :return: Concatenated video (T, sum_H, max_W, C) ''' if len(videos) == 0: return None # Find maximum width max_W = max(v.shape[2] for v in videos) # Pad all videos to max width padded_videos = [] for video in videos: padded_videos.append(pad_video_to_size(video, target_H=video.shape[1], target_W=max_W)) # Concatenate along height axis result = np.concatenate(padded_videos, axis=1) # (T, sum_H, max_W, C) if verbose: print(f'[green]Concatenated {len(videos)} videos vertically: {result.shape}') return result def create_input_pred_concatenation(sample, a4d_vae): ''' Create simple Input | Pred concatenation video (no GT). :param sample: Single output dictionary from inference :param a4d_vae: VAE for decoding latents to pixel space :return: Concatenated video as numpy array (T, H, W*2, 3) or None ''' from custom.eval.visuals import cached_decode_latent_video_auto # Decode input (rgb1) input_video = None if 'rgb1' in sample.get('x0_entries', {}): input_decoded = cached_decode_latent_video_auto( sample['x0_entries'], 'rgb1', a4d_vae) input_video = rearrange(input_decoded[0].float(), 'C T H W -> T H W C').cpu().numpy() input_video = np.clip(input_video, 0.0, 1.0) # Decode prediction (rgb0) pred_video = None if 'rgb0' in sample.get('y0_pred_entries', {}): pred_decoded = cached_decode_latent_video_auto( sample['y0_pred_entries'], 'rgb0', a4d_vae) pred_video = rearrange(pred_decoded[0].float(), 'C T H W -> T H W C').cpu().numpy() pred_video = np.clip(pred_video, 0.0, 1.0) if input_video is None or pred_video is None: print('[yellow]Cannot create Input|Pred concatenation: missing input or pred') return None # Use helper function for horizontal concatenation with padding concat_video = concatenate_videos_horizontal_with_padding( [input_video, pred_video], verbose=True) print(f'[green]Created Input|Pred concatenation: {concat_video.shape}') return concat_video def create_multi_sample_gallery(all_samples, a4d_vae, entry_names=['rgb0', 'rgb1'], uncertainty_videos=None): ''' Create gallery video with 2-row layout: Layout: [Input] [Uncert] [GT] [Sample1] [Sample2] [Sample3] [Sample4] ... Row 1: Input, Uncertainty, GT (fixed 3 columns) Row 2: All samples (3+ columns, extra samples only if num_samples != 3) For dynamic view synthesis: - Input cell: rgb1 only (shown once, not duplicated) - GT cell: rgb0 only (shown once, not duplicated) - Uncertainty cell: optional uncertainty heatmap overlay (top right) - Sample cells: rgb0 only (one per sample, all in bottom row) Note: rgb0 is always output, rgb1 is always input :param all_samples: List of output dictionaries from multiple samples :param a4d_vae: VAE for decoding latents to pixel space :param entry_names: List of entry names to include (e.g., ['rgb0', 'rgb1']) :param uncertainty_videos: Optional dict of uncertainty videos (e.g., {'rgb0': ndarray}) :return: Gallery video as numpy array (T, H_total, W_total, 3) ''' print('[cyan]Creating multi-sample gallery video...') from custom.eval.visuals import cached_decode_latent_video_auto if len(all_samples) == 0: print('[yellow]No samples provided for gallery') return None # Decode predictions for all samples: each sample shows only rgb0_pred decoded_predictions = [] # List of (T, H, W, 3) - each is just rgb0_pred for sample_idx, sample in enumerate(all_samples): # For each sample, show only predicted rgb0 if 'rgb0' in sample.get('y0_pred_entries', {}): pred_decoded = cached_decode_latent_video_auto( sample['y0_pred_entries'], 'rgb0', a4d_vae) pred_rgb = rearrange(pred_decoded[0].float(), 'C T H W -> T H W C').cpu().numpy() pred_rgb = np.clip(pred_rgb, 0.0, 1.0) decoded_predictions.append(pred_rgb) # Decode Input and GT cells from x0_entries if len(all_samples) > 0: sample = all_samples[0] # x0_entries are same for all samples # Input cell: show only rgb1 (input/conditioning view) input_video = None if 'rgb1' in sample.get('x0_entries', {}): input_decoded = cached_decode_latent_video_auto( sample['x0_entries'], 'rgb1', a4d_vae) input_video = rearrange(input_decoded[0].float(), 'C T H W -> T H W C').cpu().numpy() input_video = np.clip(input_video, 0.0, 1.0) # GT cell: show only rgb0 (ground truth output view) gt_video = None if 'rgb0' in sample.get('x0_entries', {}): gt_decoded = cached_decode_latent_video_auto( sample['x0_entries'], 'rgb0', a4d_vae) gt_video = rearrange(gt_decoded[0].float(), 'C T H W -> T H W C').cpu().numpy() gt_video = np.clip(gt_video, 0.0, 1.0) if len(decoded_predictions) == 0: print('[yellow]No valid predictions for gallery') return None # Create 2-row layout: # Row 1: Input, Uncertainty, GT (fixed 3 columns) # Row 2: Sample1, Sample2, Sample3, Sample4, ... (all samples) top_row_cells = [] bottom_row_cells = [] # Top row: Input, Uncertainty, GT if input_video is not None: top_row_cells.append(input_video) # Add uncertainty to top row (right after Input) if provided if uncertainty_videos is not None and 'rgb0' in uncertainty_videos: uncertainty_video = uncertainty_videos['rgb0'] # (T, H, W, 3) top_row_cells.append(uncertainty_video) if gt_video is not None: top_row_cells.append(gt_video) # Bottom row: All samples in order for sample_idx, pred_video in enumerate(decoded_predictions): bottom_row_cells.append(pred_video) # Balance rows: if top row has more cells, add blank to bottom row while len(bottom_row_cells) < len(top_row_cells): # Create blank cell with same dimensions as first cell blank_shape = bottom_row_cells[0].shape if len(bottom_row_cells) > 0 else top_row_cells[0].shape blank_cell = np.zeros(blank_shape, dtype=np.float32) bottom_row_cells.append(blank_cell) # Concatenate cells horizontally within each row (with padding for height mismatches) top_row = concatenate_videos_horizontal_with_padding(top_row_cells, verbose=False) bottom_row = concatenate_videos_horizontal_with_padding(bottom_row_cells, verbose=False) # Stack rows vertically (with padding for width mismatches) gallery_video = concatenate_videos_vertical_with_padding([top_row, bottom_row], verbose=False) num_pred_samples = len(decoded_predictions) num_top_row = len(top_row_cells) has_uncertainty = uncertainty_videos is not None and 'rgb0' in uncertainty_videos uncert_str = " + Uncert" if has_uncertainty else "" print(f'[green]Gallery video shape: {gallery_video.shape}') print(f'[green] Top row ({num_top_row} cells): Input, GT{uncert_str}') print(f'[green] Bottom row ({num_pred_samples} cells): {num_pred_samples} samples') return gallery_video def save_multi_sample_gallery(all_samples, a4d_vae, output_dir, dset_name, sample_id_clean, entry_names=['rgb0', 'rgb1'], fps=10, quality=8, iteration=None, uncertainty_videos=None): ''' Create and save multi-sample gallery video. :param all_samples: List of output dictionaries from multiple samples :param a4d_vae: VAE for decoding latents to pixel space :param output_dir: Output directory :param dset_name: Dataset name for filename :param sample_id_clean: Sample ID for filename :param entry_names: List of entry names to include :param fps: Frame rate (extracted from data_batch) :param quality: Video quality (0-10) :param iteration: Optional iteration index to include in filename :param uncertainty_videos: Optional dict of uncertainty videos to include in gallery ''' gallery_video = create_multi_sample_gallery(all_samples, a4d_vae, entry_names, uncertainty_videos) if gallery_video is None: return gallery_dir = pathlib.Path(output_dir) gallery_dir.mkdir(parents=True, exist_ok=True) # Save gallery video iter_str = f'{iteration:04d}_' if iteration is not None else '' gallery_file = gallery_dir / f'msg_{iter_str}{dset_name}_{sample_id_clean}.mp4' save_video(str(gallery_file), gallery_video, fps=fps, quality=quality) print(f'[green]Saved multi-sample gallery to {gallery_file}') def aggregate_and_save_metrics(all_scene_metrics, output_dir, metric_resolution=None, seed=None): ''' Aggregate metrics across all scenes and save to JSON file. :param all_scene_metrics: List of scene metrics dictionaries :param output_dir: Output directory :param metric_resolution: Optional resize spec the metrics were scored at (recorded for provenance). :param seed: Optional fixed base seed the run used (recorded for provenance). ''' if len(all_scene_metrics) == 0: print('[yellow]No metrics to aggregate') return # Per-entry + 'avg' aggregation: each scene carries psnr/ssim/lpips as {rgb0..rgbN, avg} dicts. mean, std, median, num_valid = {}, {}, {}, {} for met in ('psnr', 'ssim', 'lpips'): keys = set() for m in all_scene_metrics: if isinstance(m.get(met), dict): keys.update(m[met].keys()) mean[met], std[met], median[met], num_valid[met] = {}, {}, {}, {} # rgb* views sorted by index, then 'avg' last. for k in sorted(keys, key=lambda s: (s == 'avg', int(s[3:]) if s.startswith('rgb') else 0)): vals = [m[met][k] for m in all_scene_metrics if isinstance(m.get(met), dict) and m[met].get(k) is not None] if not vals: continue mean[met][k] = float(np.mean(vals)) std[met][k] = float(np.std(vals)) median[met][k] = float(np.median(vals)) num_valid[met][k] = len(vals) aggregated_metrics = { 'metric_resolution': metric_resolution, 'seed': seed, 'num_scenes': len(all_scene_metrics), 'num_valid_scenes': num_valid, 'mean': mean, 'std': std, 'median': median, 'per_scene_metrics': all_scene_metrics, } # Save aggregated metrics output_path = pathlib.Path(output_dir) output_path.mkdir(parents=True, exist_ok=True) aggregated_file = output_path / '_aggregated_metrics.json' save_json(aggregated_metrics, str(aggregated_file)) # Print summary (headline = avg over scored views) print('[bold green]═══════════════════════════════════════') print('[bold green] Aggregated Metrics Summary') print('[bold green]═══════════════════════════════════════') print(f'[green]Number of scenes: {aggregated_metrics["num_scenes"]}') if mean['psnr'].get('avg') is not None: print(f'[green]Mean PSNR (avg): {mean["psnr"]["avg"]:.2f} ± {std["psnr"]["avg"]:.2f}') print(f'[green]Mean SSIM (avg): {mean["ssim"]["avg"]:.4f} ± {std["ssim"]["avg"]:.4f}') print(f'[green]Mean LPIPS (avg): {mean["lpips"]["avg"]:.4f} ± {std["lpips"]["avg"]:.4f}') else: print('[yellow]No valid metrics calculated (LPIPS/SKIMAGE not available or no valid samples)') print(f'[green]Saved to: {aggregated_file}') return aggregated_metrics def override_video_from_file(data_batch, video_path, entry_name, config): ''' Override video frames in data batch with frames from MP4 file. :param data_batch (dict): Data batch from dataloader :param video_path (str): Path to MP4 file :param entry_name (str): Name of entry to override (e.g., 'rgb0') :param config: Model configuration :return (dict): Modified data batch ''' import torch.nn.functional as F print(f'[yellow]Overriding {entry_name} with video from {video_path}') # Read video frames cap = cv2.VideoCapture(video_path) frames = [] while True: ret, frame = cap.read() if not ret: break # Convert BGR to RGB and normalize to [0, 1] frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frame = frame.astype(np.float32) / 255.0 frames.append(frame) cap.release() if len(frames) == 0: print(f'[red]Warning: No frames read from {video_path}') return data_batch # Convert to tensor (T, H, W, C) -> (T, C, H, W) frames_array = np.stack(frames, axis=0) frames_tensor = torch.from_numpy(frames_array).permute(0, 3, 1, 2) # Get target shape from data batch B = data_batch['anydata']['rgb'][(0, 0)].shape[0] if 'rgb' in data_batch['anydata'] else 1 target_T = len(data_batch['anydata']['timestep']) # Derive target resolution from anydata rgb tensors first_rgb_key = next(k for k in data_batch['anydata']['rgb'] if isinstance(k, tuple)) target_H = data_batch['anydata']['rgb'][first_rgb_key].shape[-2] target_W = data_batch['anydata']['rgb'][first_rgb_key].shape[-1] # Resize frames to match target frames_tensor = F.interpolate( frames_tensor, size=(int(target_H), int(target_W)), mode='bilinear', align_corners=False ) # Trim or repeat frames to match target length if frames_tensor.shape[0] > target_T: frames_tensor = frames_tensor[:target_T] elif frames_tensor.shape[0] < target_T: # Repeat last frame last_frame = frames_tensor[-1:].repeat(target_T - frames_tensor.shape[0], 1, 1, 1) frames_tensor = torch.cat([frames_tensor, last_frame], dim=0) # Replace rgb frames in vidar dict for t in range(target_T): key = (t, 0) # Assume camera 0 for now if key in data_batch['anydata']['rgb']: # Replicate for batch dimension data_batch['anydata']['rgb'][key] = frames_tensor[t:t+1].repeat(B, 1, 1, 1).to( data_batch['anydata']['rgb'][key].device, dtype=data_batch['anydata']['rgb'][key].dtype ) print(f'[green]Overrode {target_T} frames for {entry_name}') return data_batch def save_inference_results(args, config, model, iteration, data_batch, all_samples, uncertainty_videos, diversity_metrics, scene_metrics): ''' Save inference results including visualizations, uncertainty videos, and metrics. :param args: Command line arguments (needs output_dir, save_uncertainty, quality_pix) :param config: Model configuration :param model: The model instance (for VAE access) :param iteration: Sample iteration number :param data_batch: Input data batch :param all_samples: List of output dictionaries :param uncertainty_videos: Uncertainty heatmap videos :param diversity_metrics: Diversity metrics :param scene_metrics: Per-scene metrics (PSNR, SSIM, LPIPS) Note: Visualizations are already saved directly to output_dir during validation_step ''' from custom.eval.visuals import unpack_nested print(f'[cyan]Saving additional results for sample {iteration}...') # Create output directory output_dir = pathlib.Path(args.output_dir) output_dir.mkdir(parents=True, exist_ok=True) # Get dataset name and sample ID for naming dset_name = data_batch.get('dset_name', ['unknown'])[0] if 'dset_name' in data_batch else 'unknown' sample_id = data_batch.get('sample_id', ['unknown'])[0] if 'sample_id' in data_batch else f'{iteration:04d}' # Clean sample_id for filename sample_id_clean = sample_id.replace('/', '_').replace('\\', '_').replace(':', '_') # Extract FPS from data_batch (similar to a4d_visuals.py) fps_pix = unpack_nested(data_batch, 'fps', batch_idx=0, dtype=float, optional=True) if fps_pix is None: fps_pix = 10.0 # Fallback to 10 if not available print(f'[yellow]Warning: FPS not found in data_batch, using default {fps_pix}') else: print(f'[cyan]Using FPS from data_batch: {fps_pix}') # Visualizations were already saved during validation_step print(f'[green]Visualizations saved to {output_dir}/visuals/') # Save uncertainty heatmap videos if multiple samples if len(all_samples) > 1 and args.save_uncertainty and len(uncertainty_videos) > 0: save_uncertainty_videos(uncertainty_videos, args.output_dir, dset_name, sample_id_clean, fps=fps_pix, iteration=iteration) # Save multi-sample gallery video if multiple samples if len(all_samples) > 1: entry_names = ['rgb0', 'rgb1'] # Dynamic view synthesis entries save_multi_sample_gallery( all_samples, model.vae, args.output_dir, dset_name, sample_id_clean, entry_names=entry_names, fps=fps_pix, quality=args.quality_pix, iteration=iteration, uncertainty_videos=uncertainty_videos if len(uncertainty_videos) > 0 else None) # Save per-scene metrics (only when the in-code metrics path produced any; off-mode skips this, # info.json still carries the native-res 11-view block). Filename tags the iteration. if scene_metrics: metrics_file = output_dir / f'metrics_i{iteration}_{dset_name}_{sample_id_clean}.json' metrics_data = { 'dataset': dset_name, 'sample_id': sample_id, 'iteration': iteration, 'num_samples': len(all_samples), 'scene_metrics': scene_metrics, 'diversity_metrics': diversity_metrics, } save_json(metrics_data, metrics_file, indent=2) print(f'[green]Saved metrics to {metrics_file}') return scene_metrics