File size: 7,964 Bytes
aae3ba1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 |
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() |