data/preprocessing/undistort_video_egoexo4d.py

import os
import cv2
import json
import argparse
from projectaria_tools.core import calibration
from utils import create_ffmpeg_writer, concatenate_ts_files

class Config:
    """
    Configuration settings for EgoExo4D video undistortion and processing.
    Paths and parameters are initialized with defaults but can be overridden 
    by command-line arguments.
    """
    def __init__(self, args=None):
        # --- Paths (Overridden by CLI arguments) ---
        self.VIDEO_ROOT = getattr(args, 'video_root', '/data2/v-leizhou/egoexo_data')
        self.INTRINSICS_ROOT = getattr(args, 'intrinsics_root', '/data2/v-leizhou/processed_data/aria_calib_json')
        self.SAVE_ROOT = getattr(args, 'save_root', 'debug_final_egoexo')
        
        # --- Processing Parameters (Overridden by CLI arguments) ---
        self.VIDEO_START_IDX = getattr(args, 'video_start', 0)
        self.VIDEO_END_IDX = getattr(args, 'video_end', None)
        self.BATCH_SIZE = getattr(args, 'batch_size', 1000)
        self.CRF = getattr(args, 'crf', 22)

def process_single_video(
    video_name: str, 
    aria_name: str, 
    video_root: str, 
    intrinsics_root: str, 
    save_root: str, 
    batch_size: int = 1000, 
    crf: int = 22
):
    """
    Processes a single EgoExo4D video, performs undistortion using 
    ProjectAriaTools, and saves the result in batches using FFmpeg.

    Args:
        video_name: Name of the video take folder.
        aria_name: Aria camera name used in the frame-aligned video path.
        video_root: Root directory of the input videos.
        intrinsics_root: Root directory of the intrinsics files (.json).
        save_root: Root directory for saving the output videos.
        batch_size: Number of frames to process and save per temporary TS file batch.
        crf: Constant Rate Factor (CRF) for FFmpeg encoding quality.
    """

    print(f'Processing {video_name}')

    # Construct the full video path based on EgoExo4D folder structure
    video_path = os.path.join(
        video_root, 
        'takes', 
        video_name, 
        'frame_aligned_videos', 
        f'{aria_name}_214-1.mp4'
    )
    cap = cv2.VideoCapture(video_path)

    # Get video properties
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    fps = cap.get(cv2.CAP_PROP_FPS)
    video_length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    # Load ground truth intrinsics info from JSON using ProjectAriaTools
    intrinsics_file_path = os.path.join(intrinsics_root, f'{video_name}.json')
    intrinsics_info = calibration.device_calibration_from_json(intrinsics_file_path).get_camera_calib("camera-rgb")

    # Use a fixed pinhole intrinsics for the output video resolution (1408x1408)
    pinhole = calibration.get_linear_camera_calibration(1408, 1408, 412.5) 

    # Initialize the first batch ffmpeg writer
    batch_number = 0
    writer = create_ffmpeg_writer(
        os.path.join(save_root, f'{video_name}_b{batch_number:04d}.ts'),
        width, height, fps, crf
    )

    idx = 0

    # Read and process frames
    while True:
        # Print progress in-place
        print(f'Processing {video_name} frame {idx} / {video_length}', end='\r') 
        ret, frame = cap.read()
        if not ret:
            # End of video stream: close the last writer
            writer.stdin.close()
            writer.wait()
            break

        # Undistort the frame using ProjectAriaTools' distortion function (original logic)
        undistorted_frame = calibration.distort_by_calibration(frame, pinhole, intrinsics_info)
        
        # Convert BGR to RGB before writing to ffmpeg (FFmpeg expects RGB)
        undistorted_frame = cv2.cvtColor(undistorted_frame, cv2.COLOR_BGR2RGB)

        # Write to ffmpeg stdin
        writer.stdin.write(undistorted_frame.tobytes())

        # Check if the current batch is complete
        if (idx + 1) % batch_size == 0:
            # Finalize the current batch writer
            writer.stdin.close()
            writer.wait()

            # Start the next batch writer
            batch_number += 1
            writer = create_ffmpeg_writer(
                os.path.join(save_root, f'{video_name}_b{batch_number:04d}.ts'),
                width, height, fps, crf
            )

        idx += 1

    cap.release()

    # Merge all temporary TS chunks into the final MP4 file
    concatenate_ts_files(save_root, video_name, batch_number + 1)

def main():
    """
    Main function to parse arguments, load the Aria camera name mapping, 
    load the video list, and run the undistortion process.
    """
    parser = argparse.ArgumentParser(description='Undistort EgoExo4D videos using ProjectAriaTools calibration.')
    
    # Arguments corresponding to Config parameters
    parser.add_argument('--video_root', type=str, default='/data2/v-leizhou/egoexo_data', help='Root folder containing EgoExo4D video takes')
    parser.add_argument('--intrinsics_root', type=str, default='/data2/v-leizhou/processed_data/aria_calib_json', help='Root folder containing Aria calibration JSON files')
    parser.add_argument('--save_root', type=str, default='debug_final_egoexo', help='Root folder for saving output videos')
    parser.add_argument('--video_start', type=int, default=0, help='Start video index (inclusive)')
    parser.add_argument('--video_end', type=int, default=None, help='End video index (exclusive)')
    parser.add_argument('--batch_size', type=int, default=1000, help='Number of frames to be processed per batch (TS chunk)')
    parser.add_argument('--crf', type=int, default=22, help='CRF for ffmpeg encoding quality')
    
    args = parser.parse_args()
    
    # Initialize configuration from arguments
    config = Config(args)

    # Create the output directory if it doesn't exist
    os.makedirs(config.SAVE_ROOT, exist_ok=True)

    # Get all video names automatically (assuming subfolders are video names)
    try:
        video_names = sorted(os.listdir(os.path.join(config.VIDEO_ROOT, 'takes')))
        video_names = [name.split('.')[0] for name in video_names]
    except FileNotFoundError:
        print(f"Error: Video root directory not found at {config.VIDEO_ROOT}. Cannot proceed.")
        return

    if config.VIDEO_END_IDX is None:
        end_idx = len(video_names)
    else:
        end_idx = config.VIDEO_END_IDX
    
    video_names_to_process = video_names[config.VIDEO_START_IDX:end_idx]

    if not video_names_to_process:
        print("No videos found to process in the specified range.")
        return
        
    # Load aria camera names from the JSON file (Preserves original hardcoded path)
    try:
        with open("./egoexo4d_aria_name.json", "r", encoding="utf-8") as f:
            aria_names = json.load(f)
    except FileNotFoundError:
        print("Error: The Aria name mapping file './egoexo4d_aria_name.json' was not found. Cannot proceed.")
        return
    except json.JSONDecodeError:
        print("Error: Could not decode the Aria name mapping file './egoexo4d_aria_name.json'. Cannot proceed.")
        return

    # Process videos
    for video_name in video_names_to_process:
        try:
            # Get aria name for the current video
            aria_name = aria_names[video_name]

            process_single_video(
                video_name,
                aria_name,
                config.VIDEO_ROOT,
                config.INTRINSICS_ROOT,
                config.SAVE_ROOT,
                config.BATCH_SIZE,
                config.CRF
            )
        except KeyError:
             # Handle missing Aria name for a video
             print(f'Error processing {video_name}: Aria name not found in the map file.')
             continue
        except Exception as e:
            # Catch and report other processing errors, then continue
            print(f'Error processing {video_name}: {e}')
            continue


if __name__ == '__main__':
    main()