""" Pre-Processing Pipeline: Compute BA and oracle uncertainty offline. This module handles the offline preprocessing phase that runs OUTSIDE the training loop to pre-compute expensive operations: - BA validation (CPU, expensive, slow) - Oracle uncertainty propagation (CPU, moderate) - Oracle target selection (BA vs ARKit) Results are cached to disk and loaded during training for fast iteration. Key Design: The training pipeline is split into two phases: 1. **Pre-Processing Phase** (offline, expensive): Compute BA and oracle uncertainty 2. **Training Phase** (online, fast): Load pre-computed results and train This separation allows: - BA computation outside training loop (can be parallelized) - Reuse of expensive computations across training runs - Continuous confidence weighting (not binary rejection) - Efficient training iteration (100-1000x faster) See `docs/TRAINING_PIPELINE_ARCHITECTURE.md` for detailed architecture. """ import json import logging from pathlib import Path from typing import Dict, Optional import numpy as np from ..utils.oracle_uncertainty import OracleUncertaintyPropagator from .arkit_processor import ARKitProcessor from .ba_validator import BAValidator logger = logging.getLogger(__name__) def preprocess_arkit_sequence( arkit_dir: Path, output_cache_dir: Path, model, # DA3 model for initial inference ba_validator: BAValidator, oracle_propagator: OracleUncertaintyPropagator, device: str = "cuda", prefer_arkit_poses: bool = True, min_arkit_quality: float = 0.8, use_lidar: bool = True, use_ba_depth: bool = False, ) -> Dict: """ Pre-process a single ARKit sequence: compute BA and oracle uncertainty. This runs OUTSIDE the training loop and can be parallelized across sequences. The preprocessing phase computes expensive operations once and caches results for fast training iteration. Processing Steps: 1. Extract ARKit data (poses, LiDAR depth) - FREE, fast 2. Run DA3 inference (GPU, batchable) - Moderate cost 3. Run BA validation (CPU, expensive) - Only if ARKit quality is poor 4. Compute oracle uncertainty propagation - Moderate cost 5. Save to cache - Fast disk I/O Oracle Target Selection: - If ARKit tracking quality >= min_arkit_quality: Use ARKit poses directly (fast, no BA needed) - Otherwise: Run BA validation to refine poses (expensive but necessary) Args: arkit_dir: Directory containing ARKit sequence with: - videos/*.MOV: Video file - metadata.json: ARKit metadata (poses, LiDAR, intrinsics) output_cache_dir: Directory to save pre-processed results. Each sequence will be saved as a subdirectory with: - oracle_targets.npz: BA/ARKit poses and depth - uncertainty_results.npz: Confidence and uncertainty maps - metadata.json: Sequence metadata model: DA3 model for initial inference. Used to generate initial predictions that are then validated/refined by BA. ba_validator: BAValidator instance for pose refinement via Bundle Adjustment. Only used if ARKit tracking quality is below threshold. oracle_propagator: OracleUncertaintyPropagator for computing uncertainty and confidence maps from multiple oracle sources (ARKit, BA, LiDAR). device: Device for DA3 inference ('cuda' or 'cpu'). Default 'cuda'. prefer_arkit_poses: If True, use ARKit poses when tracking quality is good. This avoids expensive BA computation. Default True. min_arkit_quality: Minimum ARKit tracking quality (0-1) to use ARKit poses directly. Below this threshold, BA validation is run. Default 0.8. use_lidar: Include ARKit LiDAR depth in oracle uncertainty computation. Default True. use_ba_depth: Include BA depth maps in oracle uncertainty computation. BA depth is optional and may not always be available. Default False. Returns: Dictionary with preprocessing results: { 'status': str, # 'success', 'skipped', 'error' 'reason': str, # Reason if skipped/error 'sequence_id': str, # Sequence identifier 'cache_path': Path, # Path to cached results 'num_frames': int, # Number of frames processed 'pose_source': str, # 'arkit' or 'ba' 'tracking_quality': float, # ARKit tracking quality (0-1) } Example: >>> from ylff.services.preprocessing import preprocess_arkit_sequence >>> from ylff.services.ba_validator import BAValidator >>> from ylff.utils.oracle_uncertainty import OracleUncertaintyPropagator >>> >>> result = preprocess_arkit_sequence( ... arkit_dir=Path("data/arkit_sequences/seq001"), ... output_cache_dir=Path("cache/preprocessed"), ... model=da3_model, ... ba_validator=ba_validator, ... oracle_propagator=oracle_propagator, ... prefer_arkit_poses=True, ... min_arkit_quality=0.8, ... ) Note: This function is designed to be called in parallel across multiple sequences. Each sequence is processed independently and results are cached separately. See `ylff preprocess arkit` CLI command for batch processing. """ sequence_id = arkit_dir.name sequence_cache_dir = output_cache_dir / sequence_id sequence_cache_dir.mkdir(parents=True, exist_ok=True) try: # Step 1: Extract ARKit data (free, fast) logger.info(f"Extracting ARKit data for {sequence_id}...") processor = ARKitProcessor(arkit_dir=arkit_dir) images = processor.extract_frames( output_dir=None, max_frames=None, frame_interval=1, return_images=True ) if len(images) < 2: return {"status": "skipped", "reason": "insufficient_frames"} # Check ARKit tracking quality good_indices = processor.filter_good_frames() good_tracking_ratio = len(good_indices) / len(images) if images else 0.0 # If tracking is poor, we can still proceed using Video-only BA mode is_video_only = good_tracking_ratio < 0.5 if is_video_only: logger.info( f"ARKit tracking missing or poor for {sequence_id} ({good_tracking_ratio:.1%}). " "Falling back to Video-only (BA-driven) mode." ) # Extract ARKit poses and intrinsics arkit_poses_c2w, intrinsics = processor.get_arkit_poses() arkit_poses_w2c = processor.convert_arkit_to_w2c(arkit_poses_c2w) # Sync frame counts (ensure images match metadata length) # This resolves "Length mismatch" errors if video and JSON are slightly off if arkit_poses_c2w is not None and len(arkit_poses_c2w) > 0: min_len = min(len(images), len(arkit_poses_c2w)) if len(images) != len(arkit_poses_c2w): logger.warning( f"Syncing {sequence_id}: video has {len(images)} frames, " f"metadata has {len(arkit_poses_c2w)}. Slicing to {min_len}." ) images = images[:min_len] arkit_poses_c2w = arkit_poses_c2w[:min_len] arkit_poses_w2c = arkit_poses_w2c[:min_len] if intrinsics is not None and len(intrinsics) > 0: intrinsics = intrinsics[:min_len] # Handle empty poses for oracle propagator if arkit_poses_c2w is not None and arkit_poses_c2w.size == 0: arkit_poses_c2w = None if arkit_poses_w2c is not None and arkit_poses_w2c.size == 0: arkit_poses_w2c = None if intrinsics is not None and intrinsics.size == 0: intrinsics = None # Extract LiDAR depth (if available) lidar_depth = None if use_lidar: lidar_depth = processor.get_lidar_depths() # Step 2: Run DA3 inference (GPU, batchable) logger.info(f"Running DA3 inference for {sequence_id} (length: {len(images)})...") import torch # Define batch size to avoid GPU memory overflow on long sequences # 8-12 frames is a good balance for MPS (Mac) memory batch_size = 8 overlap = 1 all_depths = [] all_poses = [] all_intrinsics = [] last_pose = None for i in range(0, len(images), batch_size - overlap): end_idx = min(i + batch_size, len(images)) chunk_images = images[i:end_idx] # If we've reached the end and don't have enough frames for a new batch, stop if len(chunk_images) < 2 and i > 0: break chunk_arkit = arkit_poses_c2w[i:end_idx] if arkit_poses_c2w is not None else None chunk_ix = intrinsics[i:end_idx] if intrinsics is not None else None with torch.no_grad(): chunk_output = model.inference( chunk_images, extrinsics=chunk_arkit, intrinsics=chunk_ix ) # Extract results (handles list or single Prediction object) c_depth = chunk_output.depth c_poses = chunk_output.extrinsics c_ix = getattr(chunk_output, "intrinsics", None) # Stitch poses if in video-only mode (where poses are relative to chunk start) if is_video_only and last_pose is not None: # Align current chunk to the last frame of the previous chunk # last_pose is (3, 4) w2c from previous chunk's last frame # c_poses[0] is (3, 4) w2c for the same frame in current chunk # Transform to 4x4 p_prev = np.eye(4) p_prev[:3, :] = last_pose p_curr_start = np.eye(4) p_curr_start[:3, :] = c_poses[0] # Relative transform needed: T = p_prev @ inv(p_curr_start) # This moves current chunk's local identity to match p_prev stitch_trans = p_prev @ np.linalg.inv(p_curr_start) # Apply to all poses in current chunk for j in range(len(c_poses)): p_j = np.eye(4) p_j[:3, :] = c_poses[j] c_poses[j] = (stitch_trans @ p_j)[:3, :] # Store results, skipping the overlapping first frame for subsequent chunks skip = overlap if i > 0 else 0 all_depths.append(c_depth[skip:]) all_poses.append(c_poses[skip:]) if c_ix is not None: all_intrinsics.append(c_ix[skip:]) # Update last_pose for next chunk alignment last_pose = c_poses[-1] if end_idx == len(images): break # Combine all chunks da3_depth = np.concatenate(all_depths, axis=0) da3_poses = np.concatenate(all_poses, axis=0) da3_intrinsics = ( np.concatenate(all_intrinsics, axis=0) if all_intrinsics else (intrinsics if intrinsics is not None else None) ) da3_output_summary = { "extrinsics": da3_poses, "depth": da3_depth, "intrinsics": da3_intrinsics } # Step 3: Decide on oracle targets use_arkit_poses = ( prefer_arkit_poses and good_tracking_ratio >= min_arkit_quality and not is_video_only ) if use_arkit_poses: # Use ARKit poses directly (fast, no BA needed) logger.info( f"Using ARKit poses for {sequence_id} " f"(tracking quality: {good_tracking_ratio:.1%})" ) oracle_poses = arkit_poses_w2c pose_source = "arkit" ba_poses = None ba_depths = None else: # Run BA validation (CPU, expensive, slow) if is_video_only: logger.info(f"Running video-only BA reconstruction for {sequence_id}...") else: logger.info( f"Running BA validation for {sequence_id} " f"(ARKit tracking quality: {good_tracking_ratio:.1%} < {min_arkit_quality:.1%})" ) ba_result = ba_validator.validate( images=images, poses_model=da3_poses, intrinsics=da3_intrinsics, ) # Fix: Validator returns 'poses_ba', not 'ba_poses' ba_poses_extracted = ba_result.get("poses_ba") if ba_poses_extracted is None: if is_video_only: logger.warning(f"BA reconstruction failed for video-only sequence {sequence_id}") return {"status": "skipped", "reason": "ba_failed"} # BA failed, but we have ARKit to fall back on logger.warning(f"BA failed for {sequence_id}, falling back to ARKit poses") oracle_poses = arkit_poses_w2c pose_source = "arkit_fallback" ba_poses = None ba_depths = None else: oracle_poses = ba_poses_extracted pose_source = "ba" ba_poses = ba_poses_extracted ba_depths = ba_result.get("ba_depths") if use_ba_depth else None # Step 4: Compute oracle uncertainty propagation logger.info(f"Computing oracle uncertainty for {sequence_id}...") uncertainty_results = oracle_propagator.propagate_uncertainty( da3_poses=da3_poses, da3_depth=da3_depth, intrinsics=intrinsics, arkit_poses=arkit_poses_c2w, ba_poses=ba_poses, lidar_depth=lidar_depth if use_lidar else None, ) # Step 5: Select oracle targets # Best available depth: LiDAR > BA depth > None oracle_depth = None if use_lidar and lidar_depth is not None: oracle_depth = lidar_depth depth_source = "lidar" elif use_ba_depth and ba_depths is not None: oracle_depth = ba_depths depth_source = "ba" else: depth_source = "none" # Step 6: Save to cache logger.info(f"Saving pre-processed results for {sequence_id}...") # Save oracle targets np.savez_compressed( sequence_cache_dir / "oracle_targets.npz", poses=oracle_poses, # (N, 3, 4) w2c depth=oracle_depth if oracle_depth is not None else np.zeros((1, 1, 1)), ) # Save uncertainty results np.savez_compressed( sequence_cache_dir / "uncertainty_results.npz", pose_confidence=uncertainty_results["pose_confidence"], # (N,) depth_confidence=uncertainty_results["depth_confidence"], # (N, H, W) collective_confidence=uncertainty_results["collective_confidence"], # (N, H, W) pose_uncertainty=uncertainty_results.get( "pose_uncertainty", np.zeros((len(images), 6)), ), depth_uncertainty=uncertainty_results.get( "depth_uncertainty", np.zeros_like(da3_depth) ), ) # Save ARKit data (for reference) np.savez_compressed( sequence_cache_dir / "arkit_data.npz", poses=arkit_poses_c2w, # (N, 4, 4) c2w lidar_depth=lidar_depth if lidar_depth is not None else np.zeros((1, 1, 1)), ) # Save metadata metadata = { "sequence_id": sequence_id, "num_frames": len(images), "tracking_quality": float(good_tracking_ratio), "pose_source": pose_source, "depth_source": depth_source, "has_lidar": lidar_depth is not None, "has_ba_depth": ba_depths is not None, "mean_pose_confidence": float(uncertainty_results["pose_confidence"].mean()), "mean_depth_confidence": float(uncertainty_results["depth_confidence"].mean()), } with open(sequence_cache_dir / "metadata.json", "w") as f: json.dump(metadata, f, indent=2) # Save image paths (or could save images themselves) image_paths_file = sequence_cache_dir / "image_paths.txt" # For now, just store sequence info (images loaded from original location) with open(image_paths_file, "w") as f: f.write(f"{arkit_dir}\n") logger.info(f"Pre-processing complete for {sequence_id}") return { "status": "success", "sequence_id": sequence_id, "num_frames": len(images), "pose_source": pose_source, "depth_source": depth_source, "mean_confidence": float(uncertainty_results["collective_confidence"].mean()), } except Exception as e: logger.error(f"Pre-processing failed for {sequence_id}: {e}", exc_info=True) return {"status": "failed", "sequence_id": sequence_id, "error": str(e)} def load_preprocessed_sample(cache_dir: Path, sequence_id: str) -> Optional[Dict]: """ Load pre-processed sample from cache. Args: cache_dir: Cache directory sequence_id: Sequence identifier Returns: Dict with pre-processed data or None if not found """ sequence_cache_dir = cache_dir / sequence_id if not sequence_cache_dir.exists(): return None try: # Load oracle targets oracle_targets_data = np.load(sequence_cache_dir / "oracle_targets.npz") oracle_targets = { "poses": oracle_targets_data["poses"], "depth": ( oracle_targets_data["depth"] if oracle_targets_data["depth"].shape != (1, 1, 1) else None ), } # Load uncertainty results uncertainty_data = np.load(sequence_cache_dir / "uncertainty_results.npz") uncertainty_results = { "pose_confidence": uncertainty_data["pose_confidence"], "depth_confidence": uncertainty_data["depth_confidence"], "collective_confidence": uncertainty_data["collective_confidence"], "pose_uncertainty": uncertainty_data.get("pose_uncertainty"), "depth_uncertainty": uncertainty_data.get("depth_uncertainty"), } # Load ARKit data arkit_data_file = sequence_cache_dir / "arkit_data.npz" arkit_data = None if arkit_data_file.exists(): arkit_data_npz = np.load(arkit_data_file) arkit_data = { "poses": arkit_data_npz["poses"], "lidar_depth": ( arkit_data_npz["lidar_depth"] if arkit_data_npz["lidar_depth"].shape != (1, 1, 1) else None ), } # Load metadata metadata_file = sequence_cache_dir / "metadata.json" metadata = {} if metadata_file.exists(): with open(metadata_file) as f: metadata = json.load(f) return { "oracle_targets": oracle_targets, "uncertainty_results": uncertainty_results, "arkit_data": arkit_data, "metadata": metadata, "sequence_id": sequence_id, } except Exception as e: logger.error(f"Failed to load pre-processed sample {sequence_id}: {e}") return None