""" FOV (Field of View) Overlap-based Memory Retrieval Module Implements geometric retrieval based on camera pose overlap for Context-as-Memory Aligned with the Context-as-Memory paper [2506.03141]. """ import os import json import random import sys from concurrent.futures import ThreadPoolExecutor, as_completed import numpy as np from typing import List, Dict, Tuple, Optional from pathlib import Path import torch from PIL import Image def degrees_to_radians(degrees: float) -> float: """Convert degrees to radians.""" return degrees * np.pi / 180 def compute_rotation_list_z_only(x: float, y: float, z: float, yaw_degrees: float) -> List[float]: """ Paper-aligned: rotation only around Z-axis (yaw). 4 parameters. R = R_z(yaw). No roll/pitch. Returns: [t_x, t_y, t_z, R_11, R_12, ..., R_33] (12 elements) """ yaw_rad = degrees_to_radians(yaw_degrees) c, s = np.cos(yaw_rad), np.sin(yaw_rad) R_z = np.array([ [c, -s, 0], [s, c, 0], [0, 0, 1], ], dtype=np.float64) return [float(x), float(y), float(z)] + R_z.flatten().tolist() def compute_rotation_list_yaw_pitch( x: float, y: float, z: float, yaw_degrees: float, pitch_degrees: float ) -> List[float]: """ Five-parameter rotation: R = R_z(pitch) @ R_y(yaw). Same convention as external pipelines. - R_y: rotation around Y (yaw in degrees). - R_z: rotation around Z (pitch in degrees). Returns [t_x, t_y, t_z, R_11, ..., R_33] (12 elements), encoder-compatible. """ yaw_rad = degrees_to_radians(yaw_degrees) pitch_rad = degrees_to_radians(pitch_degrees) cy, sy = np.cos(yaw_rad), np.sin(yaw_rad) cp, sp = np.cos(pitch_rad), np.sin(pitch_rad) R_y = np.array([ [cy, 0, sy], [0, 1, 0], [-sy, 0, cy], ], dtype=np.float64) R_z = np.array([ [cp, -sp, 0], [sp, cp, 0], [0, 0, 1], ], dtype=np.float64) R = R_z @ R_y return [float(x), float(y), float(z)] + R.flatten().tolist() def compute_rotation_list(params: List[float]) -> List[float]: """ Compute rotation matrix from camera parameters and flatten to list. - 4 params [x, y, z, yaw]: Z-only, R = R_z(yaw) (paper default). - 5 params [x, y, z, yaw, pitch]: R = R_z(pitch) @ R_y(yaw). Returns [t_x, t_y, t_z, R_11, ..., R_33] (12 elements). """ if len(params) >= 5: x, y, z = float(params[0]), float(params[1]), float(params[2]) yaw = float(params[3]) pitch = float(params[4]) return compute_rotation_list_yaw_pitch(x, y, z, yaw, pitch) if len(params) >= 4: x, y, z = float(params[0]), float(params[1]), float(params[2]) yaw = float(params[3]) return compute_rotation_list_z_only(x, y, z, yaw) x = y = z = yaw = 0.0 return compute_rotation_list_z_only(x, y, z, yaw) def yaw_deg_from_rt(rt: List[float]) -> float: """Extract yaw (degrees) from 12-dim RT. Z-only: yaw = atan2(R_21, R_11).""" if rt is None or len(rt) < 12: return 0.0 R = np.array(rt[3:12]).reshape(3, 3) return float(np.degrees(np.arctan2(R[1, 0], R[0, 0]))) def flip_yaw_rt(rt: List[float]) -> List[float]: """Return RT with yaw negated (CW<->CCW). rt is 12 elements [t_x,t_y,t_z, R_11..R_33].""" if rt is None or len(rt) < 12: return list(rt) if rt else [] tx, ty, tz = rt[0], rt[1], rt[2] yaw_deg = yaw_deg_from_rt(rt) return compute_rotation_list_z_only(tx, ty, tz, -yaw_deg) def flip_yaw_rt_list(rt_list: List[List[float]]) -> List[List[float]]: """Flip yaw for each RT in list (data aug for direction sensitivity).""" return [flip_yaw_rt(rt) for rt in rt_list] def convert_rt_to_relative(rt_list_all: List[List[float]], ref_rt: List[float]) -> List[List[float]]: """ Convert RT (rotation-translation) poses to relative coordinates. This aligns with the Context-as-Memory paper's use of relative camera poses for better geometric consistency in context frame retrieval. Args: rt_list_all: List of RT poses, each is [t_x, t_y, t_z, R_11, R_12, ..., R_33] (12 elements) ref_rt: Reference RT pose [t_x, t_y, t_z, R_11, R_12, ..., R_33] (12 elements) Returns: new_rt_list: List of relative RT poses in the same format """ def parse_rt(rt: List[float]) -> tuple: """Parse RT list into rotation matrix R and translation vector t.""" t = np.array(rt[:3]).reshape((3, 1)) R = np.array(rt[3:]).reshape((3, 3)) return R, t R_ref, T_ref = parse_rt(ref_rt) R_ref_inv = R_ref.T T_ref_inv = -R_ref_inv @ T_ref new_rt_list = [] for rt in rt_list_all: R_i, T_i = parse_rt(rt) # Convert to relative coordinates R_new = R_ref_inv @ R_i T_new = R_ref_inv @ T_i + T_ref_inv # Flatten back to list format: [t_x, t_y, t_z, R_11, R_12, ..., R_33] rt_new = T_new.flatten().tolist() + R_new.flatten().tolist() new_rt_list.append(rt_new) return new_rt_list def pose_to_rt(pose: Dict, constrain_to_xy: bool = True) -> Optional[List[float]]: """ Convert camera pose dict to RT format. Paper-aligned by default. - 2D plane (constrain_to_xy=True): use position x, y only; z=0. - Rotation: Z-axis only; use rotation[2] as yaw (degrees). Roll/pitch ignored. Args: pose: Dict with 'position' [x, y, z] and 'rotation' [roll, pitch, yaw] in degrees constrain_to_xy: If True (default), set z=0 for strict XY-plane displacement (paper). Returns: RT list: [t_x, t_y, t_z, R_11, R_12, ..., R_33] (12 elements) or None if invalid """ if pose is None: return None pos = pose.get('position', [0, 0, 0]) rot = pose.get('rotation', [0, 0, 0]) if len(pos) < 2: return None x = float(pos[0]) y = float(pos[1]) z = 0.0 if constrain_to_xy else (float(pos[2]) if len(pos) >= 3 else 0.0) # Paper: rotation only around Z-axis -> use yaw (index 2) only yaw = float(rot[2]) if len(rot) > 2 else 0.0 return compute_rotation_list([x, y, z, yaw]) def rt_to_pose(rt: List[float]) -> Optional[Dict]: """ Convert RT format back to pose dict. Args: rt: RT list [t_x, t_y, t_z, R_11, R_12, ..., R_33] (12 elements) Returns: pose dict with 'position' and 'rotation' or None if invalid """ if rt is None or len(rt) < 12: return None t = np.array(rt[:3]) R = np.array(rt[3:]).reshape((3, 3)) # Extract Euler angles from rotation matrix # Using ZYX convention (yaw-pitch-roll) sy = np.sqrt(R[0, 0] * R[0, 0] + R[1, 0] * R[1, 0]) singular = sy < 1e-6 if not singular: roll = np.arctan2(R[2, 1], R[2, 2]) pitch = np.arctan2(-R[2, 0], sy) yaw = np.arctan2(R[1, 0], R[0, 0]) else: roll = np.arctan2(-R[1, 2], R[1, 1]) pitch = np.arctan2(-R[2, 0], sy) yaw = 0 return { 'position': t.tolist(), 'rotation': [np.degrees(roll), np.degrees(pitch), np.degrees(yaw)] } def _parse_poses_dict(data: dict) -> dict: """Extract poses dict from JSON data (CineCameraActor or flat dict).""" if 'CineCameraActor' in data: return data['CineCameraActor'] return data if isinstance(data, dict) else {} def load_poses_dict(json_file: str) -> dict: """ Load full camera poses dict from JSON file (one read for all frames). Args: json_file: Path to camera pose JSON file Returns: Dict mapping frame_idx (str) -> pose dict, or {} if failed """ if not os.path.exists(json_file): return {} try: with open(json_file, 'r') as f: data = json.load(f) return _parse_poses_dict(data) except Exception as e: print(f"Error loading poses from {json_file}: {e}") return {} def load_camera_pose(json_file: str, frame_idx: int) -> Optional[Dict]: """ Load camera pose for a specific frame from JSON file. Args: json_file: Path to camera pose JSON file frame_idx: Frame index Returns: Dict with 'position' and 'rotation' keys, or None if not found """ poses = load_poses_dict(json_file) frame_key = str(frame_idx) return poses.get(frame_key) def load_camera_poses_batch(json_file: str, frame_indices: List[int]) -> List[Optional[Dict]]: """ Load camera poses for multiple frames in one JSON read. Args: json_file: Path to camera pose JSON file frame_indices: List of frame indices Returns: List of pose dicts (or None) in same order as frame_indices """ poses = load_poses_dict(json_file) return [poses.get(str(fi)) for fi in frame_indices] def load_overlap_frames(overlap_labels_dir: str, video_name: str, frame_idx: int) -> List[int]: """ Load overlapping frame indices for a given frame from overlap_labels. Args: overlap_labels_dir: Base directory for overlap labels video_name: Name of the video frame_idx: Current frame index Returns: List of overlapping frame indices """ overlap_file = os.path.join(overlap_labels_dir, video_name, f"{frame_idx}.json") if not os.path.exists(overlap_file): return [] try: # Add distributed training safety - timeout protection import signal import torch.distributed as dist def timeout_handler(signum, frame): raise TimeoutError(f"Timeout loading overlap file: {overlap_file}") # Set 10-second timeout for file operations in distributed training if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1: signal.signal(signal.SIGALRM, timeout_handler) signal.alarm(10) try: with open(overlap_file, 'r') as f: data = json.load(f) overlapping_frames = data.get('overlapping_frames', []) # Convert string indices to integers result = [int(f) for f in overlapping_frames if f.isdigit() or isinstance(f, int)] if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1: signal.alarm(0) # Cancel alarm return result finally: if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1: signal.alarm(0) # Always cancel alarm except TimeoutError as e: print(f"Warning: Timeout loading overlap file {overlap_file} (distributed training): {e}") return [] except Exception as e: print(f"Error loading overlap labels from {overlap_file}: {e}") return [] def compute_fov_overlap_3d( pose1: Dict, pose2: Dict, fov_degrees: float = 52.67, max_distance: float = 500.0 # Increased default from 50.0 to 500.0 for large scenes ) -> float: """ Compute FOV overlap score between two camera poses using 3D geometry. This implementation uses full 6-DoF camera poses (position + rotation) to compute more accurate FOV overlap, as described in Context-as-Memory. Args: pose1: First camera pose with 'position' [x, y, z] and 'rotation' [roll, pitch, yaw] in degrees pose2: Second camera pose with 'position' [x, y, z] and 'rotation' [roll, pitch, yaw] in degrees fov_degrees: Field of view in degrees (default: 52.67 from paper) max_distance: Maximum distance to consider (meters) Returns: Overlap score between 0 and 1 """ if pose1 is None or pose2 is None: return 0.0 pos1 = np.array(pose1.get('position', [0, 0, 0]), dtype=np.float32) pos2 = np.array(pose2.get('position', [0, 0, 0]), dtype=np.float32) # Distance between cameras distance = np.linalg.norm(pos2 - pos1) # Instead of returning 0.0 for distances > max_distance, we use a soft threshold # that still gives some score based on direction similarity even for far cameras distance_exceeds_max = distance > max_distance if distance < 1e-6: # Same position - high overlap return 1.0 # Extract rotation angles (assuming [roll, pitch, yaw] or [x, y, z] rotation in degrees) rot1 = pose1.get('rotation', [0, 0, 0]) rot2 = pose2.get('rotation', [0, 0, 0]) # Convert to numpy array rot1 = np.array(rot1, dtype=np.float32) rot2 = np.array(rot2, dtype=np.float32) # Compute rotation matrices from Euler angles # Note: The rotation order may vary by dataset. Common conventions: # - ZYX (yaw-pitch-roll): R = R_z(yaw) * R_y(pitch) * R_x(roll) # - XYZ (roll-pitch-yaw): R = R_x(roll) * R_y(pitch) * R_z(yaw) # Based on Context-as-Memory dataset, rotation[2] is yaw (rotation around Z-axis) # We'll use ZYX convention: yaw (Z), pitch (Y), roll (X) def euler_to_rotation_matrix(euler_angles): """Convert Euler angles [roll, pitch, yaw] in degrees to rotation matrix (ZYX order)""" roll, pitch, yaw = np.radians(euler_angles) # Rotation around X-axis (roll) Rx = np.array([ [1, 0, 0], [0, np.cos(roll), -np.sin(roll)], [0, np.sin(roll), np.cos(roll)] ]) # Rotation around Y-axis (pitch) Ry = np.array([ [np.cos(pitch), 0, np.sin(pitch)], [0, 1, 0], [-np.sin(pitch), 0, np.cos(pitch)] ]) # Rotation around Z-axis (yaw) Rz = np.array([ [np.cos(yaw), -np.sin(yaw), 0], [np.sin(yaw), np.cos(yaw), 0], [0, 0, 1] ]) # ZYX order: R = Rz * Ry * Rx R = Rz @ Ry @ Rx return R # Handle different rotation formats if len(rot1) >= 3: # Euler angles [roll, pitch, yaw] or [x, y, z] R1 = euler_to_rotation_matrix([rot1[0], rot1[1], rot1[2]]) elif len(rot1) == 9: # Rotation matrix flattened (3x3 = 9 elements) R1 = rot1.reshape(3, 3) else: # Fallback: only yaw R1 = euler_to_rotation_matrix([0, 0, rot1[2] if len(rot1) > 2 else 0]) if len(rot2) >= 3: R2 = euler_to_rotation_matrix([rot2[0], rot2[1], rot2[2]]) elif len(rot2) == 9: R2 = rot2.reshape(3, 3) else: R2 = euler_to_rotation_matrix([0, 0, rot2[2] if len(rot2) > 2 else 0]) # Camera forward vector (typically Z-axis in camera coordinate system) # In OpenCV/OpenGL convention, forward is usually -Z or +Z # Based on Context-as-Memory dataset, we assume forward is +Z (third column) forward1 = R1[:, 2] # Third column of rotation matrix forward2 = R2[:, 2] # Vector from camera1 to camera2 vec_1_to_2 = pos2 - pos1 vec_1_to_2_norm = np.linalg.norm(vec_1_to_2) vec_1_to_2_unit = vec_1_to_2 / (vec_1_to_2_norm + 1e-6) # FOV half-angle threshold (cosine of half FOV) fov_rad = np.radians(fov_degrees) fov_half_cos = np.cos(fov_rad / 2) # Check if camera1 can see camera2's position (within FOV) # cos(angle) = dot(forward, vec_to_target) # angle < fov/2 => cos(angle) > cos(fov/2) dot1 = np.dot(forward1, vec_1_to_2_unit) can_1_see_2 = dot1 > fov_half_cos # Check if camera2 can see camera1's position (within FOV) dot2 = np.dot(forward2, -vec_1_to_2_unit) # Negative because looking back can_2_see_1 = dot2 > fov_half_cos # Compute overlap score based on mutual visibility and distance # Normalize distance for scoring (use max_distance as reference, but don't hard-cut) normalized_distance = min(1.0, distance / max_distance) if max_distance > 0 else 1.0 if can_1_see_2 and can_2_see_1: # Both cameras can see each other - high overlap # Score decreases with distance, but never goes to 0 distance_factor = 1.0 - normalized_distance * 0.5 overlap = 0.8 + 0.2 * distance_factor elif can_1_see_2 or can_2_see_1: # One camera can see the other - medium overlap distance_factor = 1.0 - normalized_distance * 0.6 overlap = 0.4 + 0.3 * distance_factor else: # Check if cameras are looking in similar directions (even if not directly at each other) # This handles the case where both cameras see the same scene from different angles forward_similarity = np.dot(forward1, forward2) if forward_similarity > 0.7: # Cameras looking in similar directions # Even for far cameras, if they're looking in similar directions, there's some overlap distance_factor = 1.0 - normalized_distance * 0.7 overlap = 0.2 + 0.3 * distance_factor elif forward_similarity > 0.0: # Cameras looking in somewhat similar directions distance_factor = 1.0 - normalized_distance * 0.8 overlap = 0.05 + 0.15 * distance_factor * forward_similarity else: # Cameras looking away from each other - very low overlap # But still give some score based on distance (closer = slightly better) overlap = max(0.0, 0.01 - normalized_distance * 0.01) # Apply distance penalty for cameras exceeding max_distance (soft penalty) if distance_exceeds_max: # Reduce score by distance penalty, but don't make it zero distance_penalty = min(0.5, (distance - max_distance) / max_distance * 0.3) overlap = overlap * (1.0 - distance_penalty) return np.clip(overlap, 0.0, 1.0) # Keep the simple version as fallback def compute_fov_overlap_simple( pose1: Dict, pose2: Dict, fov_degrees: float = 52.67, max_distance: float = 50.0 ) -> float: """ Simplified FOV overlap computation (fallback). Use compute_fov_overlap_3d for more accurate results. """ return compute_fov_overlap_3d(pose1, pose2, fov_degrees, max_distance) class FOVMemoryRetriever: """ FOV-based Memory Retriever for Context-as-Memory. Retrieves relevant historical frames based on FOV overlap with current frame. """ def __init__( self, dataset_base_path: str, fov_degrees: float = 52.67, max_distance: float = 50.0, use_precomputed_overlaps: bool = True ): """ Initialize FOV Memory Retriever. Args: dataset_base_path: Base path to Context-as-Memory dataset fov_degrees: Field of view in degrees (default from paper: 52.67) max_distance: Maximum distance to consider for overlap (meters) use_precomputed_overlaps: Whether to use precomputed overlap_labels if available """ self.dataset_base_path = dataset_base_path self.fov_degrees = fov_degrees self.max_distance = max_distance self.use_precomputed_overlaps = use_precomputed_overlaps self.jsons_dir = os.path.join(dataset_base_path, 'jsons') self.overlap_labels_dir = os.path.join(dataset_base_path, 'overlap_labels') # Cache for loaded poses self._pose_cache: Dict[str, Dict] = {} def retrieve_frames( self, video_name: str, current_frame_idx: int, candidate_frame_indices: List[int], top_k: int = 5, include_last_frame: bool = True, use_relative_poses: bool = False # Experiment 1_4_2: use RT relative conversion ) -> List[int]: """ Retrieve top-k most relevant frames based on FOV overlap. According to Context-as-Memory, for temporal coherence, we should: 1. Always include the last frame (current_frame_idx - 1) as short-term memory 2. Retrieve top-(k-1) frames from history as long-term memory Args: video_name: Name of the video current_frame_idx: Index of current frame to generate candidate_frame_indices: List of candidate frame indices to consider top_k: Number of frames to retrieve include_last_frame: Whether to force include the last frame (default: True, per Context-as-Memory) use_relative_poses: Whether to use RT relative conversion (experiment 1_4_2, aligned with paper) Returns: List of top-k frame indices sorted by relevance (last frame first if included) """ if not candidate_frame_indices: return [] retrieved_frames = [] # Step 1: Force include last frame for short-term memory (Context-as-Memory requirement) if include_last_frame and current_frame_idx > 0: last_frame_idx = current_frame_idx - 1 if last_frame_idx in candidate_frame_indices: retrieved_frames.append(last_frame_idx) # Remove from candidates to avoid duplication candidate_frame_indices = [idx for idx in candidate_frame_indices if idx != last_frame_idx] # Calculate how many more frames we need remaining_k = top_k - len(retrieved_frames) if remaining_k <= 0: return retrieved_frames[:top_k] # Step 2: Retrieve long-term memory frames using FOV overlap # If precomputed overlaps are available, use them if self.use_precomputed_overlaps: overlap_frames = load_overlap_frames( self.overlap_labels_dir, video_name, current_frame_idx ) # Filter to only include candidate frames (and exclude already included last frame) overlap_frames = [f for f in overlap_frames if f in candidate_frame_indices and f not in retrieved_frames] if overlap_frames: # Take top remaining_k frames retrieved_frames.extend(overlap_frames[:remaining_k]) return retrieved_frames[:top_k] # Step 3: Compute FOV overlap using camera poses (if precomputed not available) current_pose = self._load_pose(video_name, current_frame_idx) if current_pose is None: # Fallback: return first k candidates retrieved_frames.extend(candidate_frame_indices[:remaining_k]) return retrieved_frames[:top_k] # Experiment 1_4_2: Convert to relative poses if enabled (aligned with Context-as-Memory) if use_relative_poses: # Convert current pose to RT format ref_rt = pose_to_rt(current_pose) if ref_rt is None: # Fallback to absolute poses if conversion fails use_relative_poses = False # Compute overlap scores for all candidates overlap_scores = [] for candidate_idx in candidate_frame_indices: if candidate_idx in retrieved_frames: continue # Skip already included frames candidate_pose = self._load_pose(video_name, candidate_idx) if candidate_pose is None: continue # Experiment 1_4_2: Use relative poses for FOV overlap computation if use_relative_poses and ref_rt is not None: # Convert candidate pose to RT format candidate_rt = pose_to_rt(candidate_pose) if candidate_rt is not None: # Convert to relative coordinates relative_rt_list = convert_rt_to_relative([candidate_rt], ref_rt) if relative_rt_list: # Convert back to pose format for FOV overlap computation relative_pose = rt_to_pose(relative_rt_list[0]) if relative_pose is not None: # Use relative pose for overlap computation # Reference pose in relative coordinates is identity (origin) ref_relative_pose = {'position': [0, 0, 0], 'rotation': [0, 0, 0]} score = compute_fov_overlap_3d( ref_relative_pose, relative_pose, self.fov_degrees, self.max_distance ) overlap_scores.append((candidate_idx, score)) continue # Fallback: Use absolute poses (original method) score = compute_fov_overlap_3d( current_pose, candidate_pose, self.fov_degrees, self.max_distance ) overlap_scores.append((candidate_idx, score)) # Sort by score (descending) and take top remaining_k overlap_scores.sort(key=lambda x: x[1], reverse=True) retrieved_frames.extend([idx for idx, _ in overlap_scores[:remaining_k]]) return retrieved_frames[:top_k] def _load_pose(self, video_name: str, frame_idx: int) -> Optional[Dict]: """Load and cache camera pose.""" cache_key = f"{video_name}_{frame_idx}" if cache_key in self._pose_cache: return self._pose_cache[cache_key] json_file = os.path.join(self.jsons_dir, f"{video_name}.json") pose = load_camera_pose(json_file, frame_idx) if pose is not None: self._pose_cache[cache_key] = pose return pose def clear_cache(self): """Clear pose cache.""" self._pose_cache.clear() def create_fov_retriever(dataset_base_path: str) -> Optional[FOVMemoryRetriever]: """ Create FOV retriever if dataset has camera pose information. Args: dataset_base_path: Base path to dataset Returns: FOVMemoryRetriever instance or None if dataset doesn't support it """ jsons_dir = os.path.join(dataset_base_path, 'jsons') if not os.path.exists(jsons_dir): return None return FOVMemoryRetriever(dataset_base_path) # Context/FOV retrieval integration helpers. def _load_frame_png(frame_file: str) -> Optional[Image.Image]: """Load a single frame from PNG file.""" if os.path.exists(frame_file): try: return Image.open(frame_file).convert('RGB') except Exception: pass return None def retrieve_simple_context_frames( data: Dict, dataset_base_path: str, top_k: int = 4, # Number of overlap frames to retrieve. First Frame will be added automatically. drop_overlap_probability: float = 0.1, # 10% probability to drop overlap frames (paper strategy) use_rt_relative: bool = False, # Experiment 1_4_2: use RT relative conversion (aligned with Context-as-Memory) ) -> Tuple[List[Image.Image], List, List[int], int, str, str]: """ Retrieve context frames according to the Context-as-Memory paper [2506.03141]. Data Structure (Precomputed Retrieval Results): - Each JSON file: overlap_labels/{video_name}/{frame_index}.json - JSON structure: { "frame_index": "0", # First Frame (short-term memory) "overlapping_frames": ["2796", "2797", ..., "3839", ...] # Long-term memory (sample 4) } - frame_index and the next 80 frames constitute GT (target frames to generate) - Iterating through all JSON files = one epoch Design principles: 1. First Frame (frame_index from JSON) as Immediate Condition: - The frame_index in JSON is the First Frame (short-term memory) - Provides immediate visual and temporal starting point (Image-to-Video mode) - Always included as context - GT: frame_index and the next 80 frames (81 frames total) 2. Overlap Frames (from overlapping_frames in JSON) as Long-term Memory: - Retrieved from precomputed overlap_labels JSON files - Precomputed lists may be very long (e.g., [2796, 3839, 4183, ..., 6339]) - Random uniform sampling: sample top_k (4) frames from overlapping_frames list - Provides long-term consistency information - Memory frames are unordered snapshots (no temporal sequence) 3. Context Composition: - Order: [First Frame, Overlap Frame 1, Overlap Frame 2, Overlap Frame 3, Overlap Frame 4] - Total: 1 First Frame + top_k Overlap Frames = top_k + 1 frames (e.g., 5 frames) - Context frames are concatenated with target frames in temporal dimension 4. 10% Probability Drop Strategy: - With 10% probability, drop all Overlap Frames, only use First Frame - Simulates video generation starting stage (no historical memory) - Forces model to generate reasonable videos without long-term memory assistance 5. Epoch Definition: - One epoch = iterate through all JSON files in overlap_labels/{video_name}/ - Each JSON file = one training sample 6. Positional Encoding Note: - Memory frames should NOT use original absolute time positions - They should be treated as unordered image collection or use memory ID encoding only - First Frame should use explicit "Ref Frame" encoding Args: data: Training data dict containing video frames and metadata dataset_base_path: Base path to Context-as-Memory dataset top_k: Number of overlap frames to retrieve (default: 4). First Frame will be added automatically. drop_overlap_probability: Probability to drop overlap frames (default: 0.1 = 10%) Returns: Tuple of: (context_frames, context_actions, context_indices, current_frame_idx, video_name, source) """ video_frames = data.get("video", []) # Get segment boundaries start_frame = data.get("start_frame", 0) end_frame = data.get("end_frame", None) # Use frame_idx if available, otherwise calculate from segment (middle of segment) # This ensures we can find previous frames for context retrieval if "frame_idx" in data: current_frame_idx = data.get("frame_idx") else: # Calculate middle of segment as reference frame (same as training convention) if end_frame is not None: current_frame_idx = (start_frame + end_frame) // 2 else: # Fallback: use middle of video_frames if available if len(video_frames) > 0: current_frame_idx = len(video_frames) // 2 else: current_frame_idx = 0 # First Frame: current segment's first frame (start_frame) - ALWAYS included first_frame_idx = start_frame video_name = data.get("video_name", "") context_frames: List[Image.Image] = [] context_actions: List = [] context_indices: List[int] = [] source = "none" # Get video frames from data if not isinstance(video_frames, list): video_frames = [] if not video_name: # Try to infer from data if "video_path" in data: video_name = os.path.basename(data["video_path"]).replace(".mp4", "").replace(".avi", "") elif "file_path" in data: video_name = os.path.basename(data["file_path"]).replace(".mp4", "").replace(".avi", "") # Step 1: Load First Frame (current segment's first frame) - ALWAYS included # According to JSON structure: frame_index in JSON is the First Frame (short-term memory) # frame_index and the next 80 frames constitute GT (target frames to generate) # First Frame provides immediate visual and temporal starting point (Image-to-Video mode) frames_dir = os.path.join(dataset_base_path, 'frames', video_name) first_frame_loaded = False # Experiment 1_4_2: Load camera poses once for all context frames (one JSON read) json_file = os.path.join(dataset_base_path, "jsons", f"{video_name}.json") poses_dict = load_poses_dict(json_file) first_frame_pose_rt = None first_frame_pose = poses_dict.get(str(first_frame_idx)) if first_frame_pose is not None: first_frame_pose_rt = pose_to_rt(first_frame_pose) if first_frame_pose_rt is not None: # Experiment 1_4_2: When use_rt_relative, first frame = reference frame = identity RT # Target actions use ref=first_frame, so context first frame must also be identity # to align with target's coordinate system (same frame = same RT representation) if use_rt_relative: # Identity RT: [t=0,0,0, R=eye(3)] = [0,0,0,1,0,0,0,1,0,0,0,1] context_actions.append([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0]) else: context_actions.append(first_frame_pose_rt) else: context_actions.append([0.0] * 12) else: context_actions.append([0.0] * 12) if os.path.exists(frames_dir): first_frame_file = os.path.join(frames_dir, f"{first_frame_idx:04d}.png") if os.path.exists(first_frame_file): try: first_frame = Image.open(first_frame_file).convert('RGB') context_frames.append(first_frame) context_indices.append(first_frame_idx) first_frame_loaded = True except Exception as e: pass # Frame loading failed, skip # Fallback: use first frame from video_frames if file not found if not first_frame_loaded and len(video_frames) > 0: if isinstance(video_frames[0], Image.Image): context_frames.append(video_frames[0]) context_indices.append(start_frame) first_frame_loaded = True elif isinstance(video_frames[0], str): try: first_frame = Image.open(video_frames[0]).convert('RGB') context_frames.append(first_frame) context_indices.append(start_frame) first_frame_loaded = True except Exception as e: pass # Fallback frame loading failed, skip # Step 2: Retrieve Overlap Frames (long-term memory) - with 10% probability drop # According to the data structure: # - Each JSON file in overlap_labels/{video_name}/{frame_index}.json contains: # - frame_index: First Frame (short-term memory) - this is the current segment's start # - overlapping_frames: List of historical frames (long-term memory) - sample 4 from this list # - frame_index and the next 80 frames constitute GT (target frames to generate) # - Iterating through all JSON files = one epoch drop_overlap = random.random() < drop_overlap_probability if not drop_overlap: # Use precomputed overlap_labels for FOV overlap-based selection overlap_labels_dir = os.path.join(dataset_base_path, 'overlap_labels') if os.path.exists(overlap_labels_dir): # Load overlapping_frames from JSON file: overlap_labels/{video_name}/{current_frame_idx}.json # The JSON structure: {"frame_index": "...", "overlapping_frames": ["2796", "2797", ...]} overlapping_indices = load_overlap_frames( overlap_labels_dir, video_name, current_frame_idx ) # Filter: exclude first_frame_idx, but allow frames with FOV overlap (may include future frames) # FOV overlap indicates visual similarity, which is valuable for context memory # even if the overlapping frame is from the future overlapping_indices = [idx for idx in overlapping_indices if idx != first_frame_idx] # If we have overlap frames, randomly sample top_k from them # Note: overlap_labels are precomputed FOV overlap results, containing potentially # very long lists of non-contiguous frame indices (e.g., [2796, 3839, 4183, ...]). # We cannot use all frames due to memory constraints, so we use random uniform sampling. # Sampling all JSON files once = one epoch if overlapping_indices: # Experiment 1_4_2: Use RT relative conversion for better geometric consistency # Aligned with Context-as-Memory [2506.03141] - use relative camera poses for FOV overlap if use_rt_relative: ref_pose = poses_dict.get(str(current_frame_idx)) if ref_pose is not None: # Convert reference pose to RT format ref_rt = pose_to_rt(ref_pose) if ref_rt is not None: # Score candidates using relative poses candidate_scores = [] for candidate_idx in overlapping_indices: candidate_pose = poses_dict.get(str(candidate_idx)) if candidate_pose is None: continue # Convert to RT and compute relative pose candidate_rt = pose_to_rt(candidate_pose) if candidate_rt is not None: relative_rt_list = convert_rt_to_relative([candidate_rt], ref_rt) if relative_rt_list: relative_pose = rt_to_pose(relative_rt_list[0]) if relative_pose is not None: # Compute FOV overlap using relative poses ref_relative_pose = {'position': [0, 0, 0], 'rotation': [0, 0, 0]} score = compute_fov_overlap_3d( ref_relative_pose, relative_pose, fov_degrees=52.67, max_distance=500.0 ) candidate_scores.append((candidate_idx, score)) # Sort by score and select top_k if candidate_scores: candidate_scores.sort(key=lambda x: x[1], reverse=True) sampled_overlap_indices = [idx for idx, _ in candidate_scores[:top_k]] else: # Fallback to random sampling if RT conversion fails num_overlap_frames = max(1, top_k) sampled_overlap_indices = random.sample( overlapping_indices, min(len(overlapping_indices), num_overlap_frames) ) else: # Fallback to random sampling if RT conversion fails num_overlap_frames = max(1, top_k) sampled_overlap_indices = random.sample( overlapping_indices, min(len(overlapping_indices), num_overlap_frames) ) else: # Fallback to random sampling if reference pose not found num_overlap_frames = max(1, top_k) sampled_overlap_indices = random.sample( overlapping_indices, min(len(overlapping_indices), num_overlap_frames) ) else: # Original strategy: Random Uniform Sampling (recommended for robustness) # This allows the model to learn from diverse temporal spans of memory. # The precomputed overlapping_frames list may contain hundreds or thousands of indices, # but we only sample top_k (e.g., 4) frames due to memory constraints. num_overlap_frames = max(1, top_k) # Random uniform sampling from the precomputed overlapping_frames list # This treats memory frames as an unordered image collection # IMPORTANT: Do NOT sort sampled_indices - they are unordered snapshots, not a temporal sequence # Positional encoding should NOT use original absolute time positions for memory frames sampled_overlap_indices = random.sample( overlapping_indices, min(len(overlapping_indices), num_overlap_frames) ) # Load overlap frames (memory frames) - long-term memory # Experiment 1_4_2: Use first_frame (start_frame) as reference for ALL context RTs # Target actions use ref=first_frame; context must use same ref for trajectory alignment ref_pose_for_rt = first_frame_pose_rt if use_rt_relative else None if os.path.exists(frames_dir): to_load = [(idx, os.path.join(frames_dir, f"{idx:04d}.png")) for idx in sampled_overlap_indices[:top_k]] frames_loaded = {} with ThreadPoolExecutor(max_workers=max(1, min(5, len(to_load)))) as ex: futures = {ex.submit(_load_frame_png, fp): idx for idx, fp in to_load} for fut in as_completed(futures): frame_idx = futures[fut] frame = fut.result() if frame is not None: frames_loaded[frame_idx] = frame for frame_idx in sampled_overlap_indices[:top_k]: if frame_idx not in frames_loaded: continue frame = frames_loaded[frame_idx] context_frames.append(frame) context_indices.append(frame_idx) pose = poses_dict.get(str(frame_idx)) if pose is not None: rt_pose = pose_to_rt(pose) if rt_pose is not None: if use_rt_relative and ref_pose_for_rt is not None: relative_rt_list = convert_rt_to_relative([rt_pose], ref_pose_for_rt) context_actions.append(relative_rt_list[0] if relative_rt_list else rt_pose) else: context_actions.append(rt_pose) else: context_actions.append([0.0] * 12) else: context_actions.append([0.0] * 12) source = "overlap_labels_random" else: # No overlap frames found, will use fallback below source = "first_frame_only" else: # No overlap_labels directory, will use fallback below source = "first_frame_only" else: # 10% probability: drop overlap frames, only use First Frame # This simulates video generation starting stage (no historical memory) source = "first_frame_only_dropped" # Step 3: Fallback if we don't have enough overlap frames (and not dropped) # Only fill if we haven't dropped overlap frames and need more frames if source not in ["first_frame_only_dropped"] and len(context_frames) < top_k + 1: # Try random fallback: use random previous frames before current_frame_idx max_prev_frame = max(1, current_frame_idx - 1) if max_prev_frame > 1: # Exclude first_frame_idx from random sampling candidate_indices = [idx for idx in range(max_prev_frame) if idx != first_frame_idx and idx < current_frame_idx] if candidate_indices: num_needed = top_k + 1 - len(context_frames) num_random_frames = min(len(candidate_indices), num_needed) if num_random_frames > 0: sampled_indices = random.sample(candidate_indices, num_random_frames) # Load random frames if os.path.exists(frames_dir): for frame_idx in sampled_indices: frame_file = os.path.join(frames_dir, f"{frame_idx:04d}.png") if os.path.exists(frame_file): try: frame = Image.open(frame_file).convert('RGB') context_frames.append(frame) context_indices.append(frame_idx) pose = poses_dict.get(str(frame_idx)) if pose is not None: rt_pose = pose_to_rt(pose) if rt_pose is not None: # Convert to relative RT if enabled if use_rt_relative and first_frame_pose_rt is not None: relative_rt_list = convert_rt_to_relative([rt_pose], first_frame_pose_rt) context_actions.append(relative_rt_list[0] if relative_rt_list else rt_pose) else: context_actions.append(rt_pose) else: context_actions.append([0.0] * 12) else: context_actions.append([0.0] * 12) except Exception as e: pass # Frame loading failed, skip if source == "first_frame_only": source = "random_fallback" # Step 4: Final fallback - use additional frames from current segment if needed # Target is top_k + 1 frames (1 First Frame + top_k Overlap/Random Frames) target_total_frames = top_k + 1 if len(context_frames) < target_total_frames and len(video_frames) > 0: num_needed = target_total_frames - len(context_frames) # Use additional frames from current segment (after first frame) segment_start_idx = 1 if len(context_frames) > 0 else 0 for i in range(segment_start_idx, min(segment_start_idx + num_needed, len(video_frames))): frame_idx_seg = start_frame + i if isinstance(video_frames[i], Image.Image): context_frames.append(video_frames[i]) context_indices.append(frame_idx_seg) elif isinstance(video_frames[i], str): # If it's a path, load it try: frame = Image.open(video_frames[i]).convert('RGB') context_frames.append(frame) context_indices.append(frame_idx_seg) except: pass pose = poses_dict.get(str(frame_idx_seg)) if pose is not None: rt_pose = pose_to_rt(pose) if rt_pose is not None: # Convert to relative RT if enabled if use_rt_relative and first_frame_pose_rt is not None: relative_rt_list = convert_rt_to_relative([rt_pose], first_frame_pose_rt) context_actions.append(relative_rt_list[0] if relative_rt_list else rt_pose) else: context_actions.append(rt_pose) else: context_actions.append([0.0] * 12) else: context_actions.append([0.0] * 12) # Update source if we used segment frames if len(context_frames) >= num_needed and source in ["first_frame_only", "none"]: source = "segment_fallback" # Step 5: Ensure we return exactly top_k + 1 frames # If we have some frames but not enough, pad by repeating the last frame if len(context_frames) < target_total_frames: if context_frames: last_frame = context_frames[-1] last_idx = context_indices[-1] if context_indices else first_frame_idx last_action = context_actions[-1] if context_actions else [0.0] * 12 while len(context_frames) < target_total_frames: context_frames.append(last_frame) context_indices.append(last_idx) context_actions.append(last_action) # Always pad with pose data # Limit to top_k + 1 (in case we have more) context_frames = context_frames[:target_total_frames] if context_indices: context_indices = context_indices[:target_total_frames] if context_actions: context_actions = context_actions[:target_total_frames] # Ensure context_actions length matches context_frames (pad with zeros if needed) # Always ensure context_actions are provided (not just when use_rt_relative=True) while len(context_actions) < len(context_frames): context_actions.append([0.0] * 12) return context_frames, context_actions, context_indices, current_frame_idx, video_name, source def retrieve_fov_context_frames( data: Dict, dataset_base_path: str, fov_retriever=None, top_k: int = 4, # Number of overlap frames to retrieve. First Frame will be added automatically. use_precomputed_overlaps: bool = True, use_rt_relative: bool = False, # Experiment 1_4_2: Use RT relative conversion (aligned with Context-as-Memory) strict_overlap_labels: bool = False, allow_realtime_fallback: bool = True, allow_segment_fallback: bool = True, drop_overlap_probability: float = 0.1, # 10% probability to drop overlap frames (paper strategy) ): """ Backward-compatible wrapper. We use FOV overlap scoring to select top-k overlap frames (as in Context-as-Memory). According to Context-as-Memory [2506.03141]: - First Frame (current segment's first frame) is always included as immediate condition - Overlap Frames are retrieved as long-term memory - With 10% probability, drop overlap frames to simulate starting stage Experiment 1_4_2: Uses RT relative conversion for better geometric consistency. """ context_frames, context_actions, context_indices, cur_idx, video_name, source = retrieve_simple_context_frames( data=data, dataset_base_path=dataset_base_path, top_k=top_k, # top_k is number of overlap frames (4), First Frame will be added automatically (total: 5) use_rt_relative=use_rt_relative, # Experiment 1_4_2: RT relative conversion drop_overlap_probability=drop_overlap_probability, # 10% probability to drop overlap frames ) # Check if we have top_k + 1 frames (1 First Frame + top_k Overlap Frames) target_total_frames = top_k + 1 if strict_overlap_labels and len(context_frames) < target_total_frames: return [], [], [], cur_idx, video_name, "overlap_labels_insufficient" return context_frames, context_actions, context_indices, cur_idx, video_name, source def save_sampling_jsonl( output_path: str, video_name: str, frame_index: int, context_indices: List[int], prompt: Optional[str] = None, start_frame: Optional[int] = None, end_frame: Optional[int] = None, source: Optional[str] = None, append: bool = True, ) -> None: """ Save context sampling result to JSONL file for eval consistency. Format: { "video_name": "AncientTempleEnv_0", "frame_index": 0, # First Frame (short-term memory) "context_indices": [0, 2796, 3839, 4183, 6339], # First Frame + 4 Overlap Frames "prompt": "...", # Optional "start_frame": 0, # Optional: GT segment start "end_frame": 80, # Optional: GT segment end "source": "overlap_labels_random" # Optional: sampling source } Args: output_path: Path to JSONL file video_name: Video name frame_index: First Frame index (from JSON file) context_indices: List of context frame indices [first_frame, overlap1, overlap2, ...] prompt: Optional prompt text start_frame: Optional GT segment start frame end_frame: Optional GT segment end frame source: Optional sampling source append: Whether to append to existing file (default: True) """ os.makedirs(os.path.dirname(output_path), exist_ok=True) item = { "video_name": video_name, "frame_index": frame_index, "context_indices": context_indices, } if prompt is not None: item["prompt"] = prompt if start_frame is not None: item["start_frame"] = start_frame if end_frame is not None: item["end_frame"] = end_frame if source is not None: item["source"] = source mode = "a" if append else "w" with open(output_path, mode, encoding="utf-8") as f: f.write(json.dumps(item, ensure_ascii=False) + "\n") def load_sampling_jsonl(jsonl_path: str) -> List[Dict]: """ Load context sampling results from JSONL file. Args: jsonl_path: Path to JSONL file Returns: List of sampling items, each containing: { "video_name": str, "frame_index": int, "context_indices": List[int], "prompt": Optional[str], "start_frame": Optional[int], "end_frame": Optional[int], "source": Optional[str] } """ if not os.path.exists(jsonl_path): return [] items = [] with open(jsonl_path, "r", encoding="utf-8") as f: for line in f: line = line.strip() if not line: continue try: item = json.loads(line) items.append(item) except json.JSONDecodeError as e: print(f"Warning: Failed to parse JSONL line: {e}") continue return items def setup_fov_retriever_for_training( dataset_base_path: str, enable_fov_retrieval: bool = True, ) -> Optional[object]: """ Setup FOV retriever for training (simplified version). Context frames are selected by FOV overlap scoring from precomputed `overlap_labels`. """ return None _WARN_ONCE_KEYS = set() def _warn_once(key: str, msg: str) -> None: if key in _WARN_ONCE_KEYS: return _WARN_ONCE_KEYS.add(key) print(f"[context_retrieval] WARN: {msg}", file=sys.stderr, flush=True) def _load_latent(latent_dir: str, video_name: str, frame_idx: int): """Load a single-frame latent. Expects latent_dir/video_name/{frame_idx:04d}.pt or .pt with key 'latent' or raw tensor.""" base = os.path.join(latent_dir, video_name) for fmt in (f"{frame_idx:04d}.pt", f"{frame_idx}.pt"): path = os.path.join(base, fmt) if os.path.isfile(path): try: try: x = torch.load(path, map_location="cpu", weights_only=True) except TypeError: x = torch.load(path, map_location="cpu") if isinstance(x, dict) and "latent" in x: z = x["latent"] else: z = x if hasattr(z, "shape"): # (C, 1, H, W) or (C, H, W) -> flatten for similarity return z.flatten() return None except Exception: pass return None def latent_sim_rank( video_name: str, first_frame_idx: int, overlapping_indices: List[int], dataset_base_path: str, top_k: int, latent_dir: Optional[str] = None, use_cosine: bool = True, ) -> List[int]: """ Rank overlapping frame indices by latent similarity to the first (reference) frame. If latent_dir is None or latents are missing, falls back to random sample. Expects per-frame latents under latent_dir/video_name/{frame_idx:04d}.pt (or {frame_idx}.pt). """ if not overlapping_indices or top_k <= 0: return [] if latent_dir is None or not os.path.isdir(latent_dir): return random.sample(overlapping_indices, min(top_k, len(overlapping_indices))) ref = _load_latent(latent_dir, video_name, first_frame_idx) if ref is None: return random.sample(overlapping_indices, min(top_k, len(overlapping_indices))) ref = ref.float().unsqueeze(0) scores = [] for idx in overlapping_indices: cand = _load_latent(latent_dir, video_name, idx) if cand is None: continue cand = cand.float().unsqueeze(0) if use_cosine: sim = torch.nn.functional.cosine_similarity(ref, cand, dim=1).item() else: sim = -((ref - cand) ** 2).sum().item() scores.append((idx, sim)) if not scores: return random.sample(overlapping_indices, min(top_k, len(overlapping_indices))) scores.sort(key=lambda x: x[1], reverse=True) return [idx for idx, _ in scores[:top_k]] def retrieve_context_frames_advanced( data: Dict, dataset_base_path: str, top_k: int = 4, drop_overlap_probability: float = 0.1, use_rt_relative: bool = False, retrieval_method: str = "fov", latent_retrieval_dir: Optional[str] = None, strict_overlap_labels: bool = False, ) -> Tuple[List, List, List[int], int, str, str]: """ Retrieve context frames with pluggable retrieval method. Interface matches retrieve_fov_context_frames return: (context_frames, context_actions, context_indices, cur_idx, video_name, source). retrieval_method: - "fov": use existing FOV/overlap_labels logic (random or RT-scored from overlap). - "latent_sim": rank overlap candidates by latent similarity to first frame (requires latent_retrieval_dir). When latent_sim is used but latent_retrieval_dir is missing or latents absent, falls back to FOV behavior. """ if retrieval_method == "latent_sim" and not latent_retrieval_dir: _warn_once( "latent_sim_missing_dir", "retrieval_method=latent_sim but latent_retrieval_dir is not set; fallback to FOV retrieval.", ) if retrieval_method == "fov" or (retrieval_method == "latent_sim" and not latent_retrieval_dir): return retrieve_simple_context_frames( data=data, dataset_base_path=dataset_base_path, top_k=top_k, drop_overlap_probability=drop_overlap_probability, use_rt_relative=use_rt_relative, ) if retrieval_method == "latent_sim" and latent_retrieval_dir and not os.path.isdir(latent_retrieval_dir): _warn_once( "latent_sim_bad_dir", f"latent_retrieval_dir not found: {latent_retrieval_dir}; latent_sim will degrade to random-overlap selection.", ) # latent_sim: we need to inject a custom ranking into the flow. We do a minimal duplicate of the # overlap selection step then reuse the rest via a wrapper around retrieve_simple_context_frames # by passing a custom rank function. Since retrieve_simple_context_frames doesn't support that yet, # we implement a full path here that mirrors it but uses latent_sim_rank for overlap selection. video_frames = data.get("video", []) start_frame = data.get("start_frame", 0) end_frame = data.get("end_frame", None) if "frame_idx" in data: current_frame_idx = data.get("frame_idx") else: current_frame_idx = (start_frame + end_frame) // 2 if end_frame is not None else (len(video_frames) // 2 if video_frames else 0) first_frame_idx = start_frame video_name = data.get("video_name", "") if not video_name and "video_path" in data: video_name = os.path.basename(data["video_path"]).replace(".mp4", "").replace(".avi", "") elif not video_name and "file_path" in data: video_name = os.path.basename(data["file_path"]).replace(".mp4", "").replace(".avi", "") frames_dir = os.path.join(dataset_base_path, "frames", video_name) json_file = os.path.join(dataset_base_path, "jsons", f"{video_name}.json") poses_dict = load_poses_dict(json_file) if os.path.isfile(json_file) else {} first_frame_pose = poses_dict.get(str(first_frame_idx)) first_frame_pose_rt = pose_to_rt(first_frame_pose) if first_frame_pose is not None and pose_to_rt else None context_frames: List[Image.Image] = [] context_actions: List = [] context_indices: List[int] = [] source = "none" def _append_pose(frame_idx: int): pose = poses_dict.get(str(frame_idx)) if pose is not None and pose_to_rt and use_rt_relative and first_frame_pose_rt is not None: rt = pose_to_rt(pose) if rt is not None and convert_rt_to_relative: rel = convert_rt_to_relative([rt], first_frame_pose_rt) context_actions.append(rel[0] if rel else [0.0] * 12) else: context_actions.append([0.0] * 12) elif pose is not None and pose_to_rt: rt = pose_to_rt(pose) context_actions.append(rt if rt is not None else [0.0] * 12) else: context_actions.append([0.0] * 12) if first_frame_pose_rt is not None and use_rt_relative: context_actions.append([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0]) else: context_actions.append(first_frame_pose_rt if first_frame_pose_rt is not None else [0.0] * 12) if os.path.isdir(frames_dir): first_path = os.path.join(frames_dir, f"{first_frame_idx:04d}.png") if os.path.isfile(first_path): try: context_frames.append(Image.open(first_path).convert("RGB")) context_indices.append(first_frame_idx) except Exception: pass if not context_frames and video_frames: if isinstance(video_frames[0], Image.Image): context_frames.append(video_frames[0]) context_indices.append(start_frame) elif isinstance(video_frames[0], str) and os.path.isfile(video_frames[0]): try: context_frames.append(Image.open(video_frames[0]).convert("RGB")) context_indices.append(start_frame) except Exception: pass drop_overlap = random.random() < drop_overlap_probability if drop_overlap: target_total = top_k + 1 if strict_overlap_labels and len(context_frames) < target_total: return [], [], [], current_frame_idx, video_name, "overlap_labels_insufficient" return context_frames, context_actions, context_indices, current_frame_idx, video_name, "first_frame_only_dropped" overlap_labels_dir = os.path.join(dataset_base_path, "overlap_labels") overlapping_indices = [] if os.path.isdir(overlap_labels_dir): overlapping_indices = load_overlap_frames(overlap_labels_dir, video_name, current_frame_idx) overlapping_indices = [i for i in overlapping_indices if i != first_frame_idx] if not overlapping_indices: source = "first_frame_only" if strict_overlap_labels and len(context_frames) < top_k + 1: return [], [], [], current_frame_idx, video_name, "overlap_labels_insufficient" return context_frames, context_actions, context_indices, current_frame_idx, video_name, source sampled_overlap_indices = latent_sim_rank( video_name, first_frame_idx, overlapping_indices, dataset_base_path, top_k, latent_dir=latent_retrieval_dir, use_cosine=True, ) def _load_frame(path: str): if os.path.isfile(path): try: return Image.open(path).convert("RGB") except Exception: pass return None to_load = [(idx, os.path.join(frames_dir, f"{idx:04d}.png")) for idx in sampled_overlap_indices[:top_k]] with ThreadPoolExecutor(max_workers=max(1, min(5, len(to_load)))) as ex: futures = {ex.submit(_load_frame, path): idx for idx, path in to_load} for fut in futures: idx = futures[fut] frame = fut.result() if frame is not None: context_frames.append(frame) context_indices.append(idx) _append_pose(idx) source = "latent_sim" if strict_overlap_labels and len(context_frames) < top_k + 1: return [], [], [], current_frame_idx, video_name, "overlap_labels_insufficient" return context_frames, context_actions, context_indices, current_frame_idx, video_name, source