# Output (ProceessedDataset format): # episode_no/ # - depth/camera_no/frame_no.npz # - rgb/camera_no/frame_no.jpg (using left_rgb from input) # - lowdim/camera_no/frame_no.npz # Input (AssemblyHands format): # episode_name/ # e.g. nusar-2021_action_both_9054-c01a_9054_user_id_2021-02-08_160424 # - calib.txt # - .mp4 # 8 exo cameras # Original Data downloaded from: https://drive.google.com/drive/folders/1e_TIb2et_bBoa15DoBFjDT3pV-Ivqqzl # Raw uploaded at: s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_datasets/raw/AssemblyHands/ # Processed dataset (after running this script) uploaded at: s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_datasets/processed/AssemblyHands/ import os import glob import numpy as np from PIL import Image import subprocess, logging from pathlib import Path from typing import List, Tuple, Dict from concurrent.futures import ThreadPoolExecutor, as_completed IMG_WIDTH = 1920 IMG_HEIGHT = 1080 cameras = ['C10095_rgb', 'C10115_rgb', 'C10118_rgb', 'C10119_rgb', 'C10379_rgb', 'C10390_rgb', 'C10395_rgb', 'C10404_rgb'] def extract_all_frames_ffmpeg( video_path: str, output_dir: str) -> Dict[int, bool]: """Extract all frames from a video efficiently using ffmpeg.""" results = {} try: # Create output directory if it doesn't exist os.makedirs(output_dir, exist_ok=True) # Extract all frames at once cmd = [ "ffmpeg", "-i", video_path, "-vsync", "0", # Don't duplicate frames "-q:v", "2", # High quality (lower means better) "-y", os.path.join(output_dir, "frame_%06d.jpg"), "-loglevel", "quiet", ] with open(os.devnull, "w") as devnull: result = subprocess.run( cmd, stdout=devnull, stderr=devnull, stdin=subprocess.DEVNULL, timeout=300, # Longer timeout for full extraction ) # Check if extraction was successful if result.returncode == 0: # Get all extracted frames and mark as successful frames = glob.glob(os.path.join(output_dir, "frame_*.jpg")) for frame_path in frames: frame_num = int(os.path.basename(frame_path).split('_')[1][:6]) results[frame_num] = True else: logging.error(f"Failed to extract frames from {video_path}") except ( subprocess.TimeoutExpired, FileNotFoundError, subprocess.CalledProcessError, ) as e: logging.error(f"FFmpeg extraction failed: {e}") return results # Read calibration from txt file: example format """ C10395_rgb 1204.22073752 0 946.24980580 0 1204.32995694 523.14172058 0 0 1 0 0 0 0 0 0.98842458 0.01305274 0.15115048 -130.83483049 0.06194303 -0.94418810 -0.32353037 333.78326980 0.13849153 0.32914809 -0.93406725 935.94711305 C10379_rgb 1206.44560696 0 954.09405803 0 1205.85738741 528.51956590 0 0 1 0 0 0 0 0 -0.95925922 -0.03617095 0.28020245 200.71114790 0.11959818 -0.95051442 0.28673787 225.82413195 0.25596489 0.30856765 0.91611570 738.78433583 """ def extract_calibration(calib_path: str, convert_mm_to_m: bool) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray]]: assert os.path.exists(calib_path), f"Calibration file not found: {calib_path}" with open(calib_path, 'r') as file: lines = file.readlines() cam_intrinsics = {} cam_extrinsics = {} for j in range(0, len(lines), 9): cam_name = lines[j].strip() K = np.array([list(map(float, lines[j+1].strip().split())), list(map(float, lines[j+2].strip().split())), list(map(float, lines[j+3].strip().split()))]) cam_intrinsics[cam_name] = K print("Camera: ", cam_name, " Intrinsics: ", K) ext_vals_1 = list(map(float, lines[j+5].strip().split())) ext_vals_2 = list(map(float, lines[j+6].strip().split())) ext_vals_3 = list(map(float, lines[j+7].strip().split())) extrinsics = np.eye(4) extrinsics[:3, :3] = np.array(ext_vals_1[0:3] + ext_vals_2[0:3] + ext_vals_3[0:3]).reshape(3,3) extrinsics[:3, 3] = np.array(ext_vals_1[3:4] + ext_vals_2[3:4] + ext_vals_3[3:4]) if convert_mm_to_m: extrinsics[:3, 3] = extrinsics[:3, 3] * 0.001 # convert mm to meters cam_extrinsics[cam_name] = extrinsics print("Camera: ", cam_name, " Extrinsics: ", extrinsics) return cam_intrinsics, cam_extrinsics def save_lowdim(success_frames, intrinsics_matrix, extrinsics_matrix, img_h, img_w, original_episode_name, cam_name, output_lowdim_path): for idx in success_frames: if not os.path.exists(output_lowdim_path): os.makedirs(output_lowdim_path, exist_ok=True) lowdim = dict() lowdim.update({"intrinsics": intrinsics_matrix}) lowdim.update({"shape": [img_h, img_w]}) lowdim.update({"pose": extrinsics_matrix}) lowdim.update({"original_episode_id": original_episode_name}) lowdim.update({"camera_name": cam_name}) np.savez(os.path.join(output_lowdim_path, f"frame_{idx:06d}.npz"), **lowdim) print(f"Saved lowdim to {output_lowdim_path}") path = '~/Downloads/AssemblyHandsExo/exo_videos/' output_path = 'cv_processed/AssemblyHands/' path = os.path.expanduser(path) output_path = os.path.expanduser(output_path) print ("Processing path: ", path) print ("Output path: ", output_path) nusar_dirs = [] for dirs in os.listdir(path): print ("Found: ", dirs) if os.path.isdir(os.path.join(path, dirs)): print("Found directory: ", dirs) nusar_dirs.append(dirs) assert len(nusar_dirs) == 20 def process_camera(cam_no, cam_name, cam_path, episode_output_dir, cam_intrinsics, cam_extrinsics, nusar_dir): assert os.path.exists(cam_path), f"Camera video not found: {cam_path}" rgb_output_dir = os.path.join(episode_output_dir, 'rgb', f'camera_{cam_no}') os.makedirs(rgb_output_dir, exist_ok=True) results = extract_all_frames_ffmpeg(cam_path, rgb_output_dir) success_frames = [] if not all(results.values()): success_frames = [idx for idx, success in results.items() if success] print(f"Warning: Not all frames extracted successfully from {cam_path} for {nusar_dir}. Successful frames: {len(success_frames)}/ {len(results)}") else: success_frames = list(results.keys()) # print(f"All frames extracted successfully from {cam_path} for {nusar_dir}. Total frames: {len(success_frames)}") lowdim_output_dir = os.path.join(episode_output_dir, 'lowdim', f'camera_{cam_no}') os.makedirs(lowdim_output_dir, exist_ok=True) save_lowdim(success_frames, cam_intrinsics[cam_name], cam_extrinsics[cam_name], IMG_HEIGHT, IMG_WIDTH, nusar_dir, cam_name, lowdim_output_dir) for i, nusar_dir in enumerate(nusar_dirs): print("Processing directory: ", nusar_dir) episode_name = f'episode_{i:04d}' episode_output_dir = os.path.join(output_path, episode_name) calib_path = os.path.join(path, nusar_dir, 'calib.txt') cam_intrinsics, cam_extrinsics = extract_calibration(calib_path, convert_mm_to_m=True) camera_paths = {cam: os.path.join(path, nusar_dir, f"{cam}.mp4") for cam in cameras} with ThreadPoolExecutor(max_workers=8) as executor: futures = [] for cam_no, (cam_name, cam_path) in enumerate(camera_paths.items()): futures.append(executor.submit(process_camera, cam_no, cam_name, cam_path, episode_output_dir, cam_intrinsics, cam_extrinsics, nusar_dir)) for future in as_completed(futures): future.result() # Raise exceptions if any