# BVH, Feb 2026 # DROID raw dataset → unified format converter. # # Usage: # python anydata/converters/droid.py \ # --src /data/cv_downloaded/DROID \ # --dst /data/cv_unified/DROID \ # --cache_dir /data/droid_cache # # Supplemental data is auto-downloaded on first run and cached inside --cache_dir: # - intrinsics.json — per-camera intrinsic calibrations (HuggingFace KarlP/droid) # - pnp_cam2base_multiview.json — Zubair PnP cam2base extrinsics, 53k eps (S3) # - cam2base_extrinsic_*.json — KarlP cam2base extrinsics (HuggingFace KarlP/droid) # - droid_language_annotations.json — 3 language annotations per episode (HuggingFace) # # Intrinsics priority: # 1. Episode-level: KarlP intrinsics.json keyed by uuid → serial # 2. Serial-level fallback: median intrinsics per camera serial across all KarlP episodes # (ZED cameras have fixed intrinsics per physical serial; sub-pixel variation across episodes) # # Extrinsics priority for ext1/ext2 (static cameras): # 1. Zubair PnP (53k episodes, always 2 cams, no corrupt outliers) # 2. KarlP GT-only fallback (Pred entries have >1° errors, skipped) # 3. H5 per-frame trajectory (last resort) # Wrist camera always uses H5 per-frame trajectory (captures gripper motion). # Both sources' raw 6-DOF are stored in metadata.specific.extrinsics_sources. # 6-DOF poses use XYZ Euler angles. # # Input: {src}/raw/{lab}/{success|failure}/{date}/{timestamp}/ # ├── metadata_{uuid}.json # ├── trajectory.h5 # └── recordings/MP4/{serial}.mp4 # # Output: {dst}/{lab}/{success|failure}/{date}/{timestamp}/ # ├── metadata.json # ├── rgb/{wrist,ext1,ext2}/0000000000.jpg … # └── lowdim/{wrist,ext1,ext2}/0000000000.npz … import os import h5py import json import subprocess import numpy as np from glob import glob from rich.console import Console from rich.progress import track from scipy.spatial.transform import Rotation from anydata.utils.write import write_npz, write_json from anydata.converters.utils import ( add_key_to_dict, fill_metadata, parse_dst_seq, frame_name, get_splits, process_sequences, extract_frames_from_mp4, create_metadata_shared, CANONICAL_LABEL_ORDER, ) console = Console() # Canonical camera name order CAMERA_NAMES = ['ext1', 'ext2', 'wrist'] # HuggingFace KarlP/droid supplemental data URLs INTRINSICS_URL = 'https://huggingface.co/KarlP/droid/resolve/main/intrinsics.json' EXTRINSICS_SUPERSET_URL = 'https://huggingface.co/KarlP/droid/resolve/main/cam2base_extrinsic_superset.json' EXTRINSICS_FALLBACK_URL = 'https://huggingface.co/KarlP/droid/resolve/main/cam2base_extrinsics.json' LANGUAGE_URL = 'https://huggingface.co/KarlP/droid/resolve/main/droid_language_annotations.json' # Zubair PnP extrinsics (53k episodes, always 2 cams, no corrupt outliers) ZUBAIR_EXTRINSICS_S3 = 's3://scratch-tri-global/zubair.irshad/extrinsics_keypoint/pnp_cam2base_multiview.json' # Per-process caches (loaded lazily after fork) _intrinsics_cache = None _serial_intrinsics_cache = None _zubair_extrinsics_cache = None _karlp_extrinsics_cache = None _language_cache = None _omniworld_language_cache = None _raw_source_cache = None def _get_intrinsics(cache_dir): """Return intrinsics lookup, loading from disk on first call per process.""" global _intrinsics_cache if _intrinsics_cache is not None: return _intrinsics_cache _intrinsics_cache = load_intrinsics_lookup(cache_dir) return _intrinsics_cache def _get_serial_intrinsics(cache_dir): """Return serial→K fallback lookup, building from intrinsics.json on first call.""" global _serial_intrinsics_cache if _serial_intrinsics_cache is not None: return _serial_intrinsics_cache _serial_intrinsics_cache = build_serial_intrinsics_lookup(cache_dir) return _serial_intrinsics_cache def _get_zubair_extrinsics(cache_dir): """Return Zubair PnP extrinsics lookup, loading from disk on first call per process.""" global _zubair_extrinsics_cache if _zubair_extrinsics_cache is not None: return _zubair_extrinsics_cache _zubair_extrinsics_cache = load_zubair_extrinsics_lookup(cache_dir) return _zubair_extrinsics_cache def _get_karlp_extrinsics(cache_dir): """Return KarlP extrinsics lookup (with source info), loading from disk on first call per process.""" global _karlp_extrinsics_cache if _karlp_extrinsics_cache is not None: return _karlp_extrinsics_cache _karlp_extrinsics_cache = load_karlp_extrinsics_lookup(cache_dir) return _karlp_extrinsics_cache def _get_language(cache_dir): """Return language lookup, loading from disk on first call per process.""" global _language_cache if _language_cache is not None: return _language_cache _language_cache = load_language_lookup(cache_dir) return _language_cache def _get_omniworld_language(cache_dir): """Return OmniWorld caption lookup {uuid: {cam: caption}}, or empty dict if cache missing.""" global _omniworld_language_cache if _omniworld_language_cache is not None: return _omniworld_language_cache cache_path = os.path.join(cache_dir, 'omniworld_language.json') if os.path.exists(cache_path): with open(cache_path) as f: _omniworld_language_cache = json.load(f) console.print(f'Loaded OmniWorld captions for {len(_omniworld_language_cache)} episodes') else: console.print('[yellow]OmniWorld language cache not found — skipping captions[/yellow]') _omniworld_language_cache = {} return _omniworld_language_cache def _get_raw_source_lookup(src_dir): """Return episode→raw_source mapping from download_manifest.json, or empty dict.""" global _raw_source_cache if _raw_source_cache is not None: return _raw_source_cache manifest_path = os.path.join(src_dir, 'download_manifest.json') lookup = {} if os.path.exists(manifest_path): with open(manifest_path) as f: manifest = json.load(f) # Map: category name → source label for metadata category_map = { 'original_only': 'original', 'zubair_only': 'zubair', 'both': 'both', # Legacy format from old download script 'original': 'original', 'zubair': 'zubair', } for cat_key, source_label in category_map.items(): for ep_path in manifest.get(cat_key, []): lookup[ep_path] = source_label console.print(f'Loaded raw source manifest: {len(lookup)} episodes') _raw_source_cache = lookup return _raw_source_cache ####################################################### # Helper functions ####################################################### def pose_6dof_to_matrix(pose_6d): """Convert [x, y, z, rx, ry, rz] XYZ-Euler pose to 4×4 c2w matrix.""" # NOTE(bvh): cross checked consistency with Any4D/custom/prepare/convert/convert_droid_calib.py T = np.eye(4, dtype=np.float32) T[:3, :3] = Rotation.from_euler('xyz', pose_6d[3:6]).as_matrix().astype(np.float32) T[:3, 3] = pose_6d[:3].astype(np.float32) return T def find_metadata_json(episode_dir): """Find the metadata_{uuid}.json file in an episode directory.""" matches = glob(f'{episode_dir}/metadata_*.json') return matches[0] if matches else None def find_mp4(episode_dir, serial): """Find the MP4 file for a camera serial, checking known locations.""" candidates = [ f'{episode_dir}/recordings/MP4/{serial}.mp4', f'{episode_dir}/{serial}.mp4', ] for path in candidates: if os.path.exists(path): return path return None def load_intrinsics_lookup(cache_dir): """Load intrinsics.json, downloading from HuggingFace if not cached.""" path = f'{cache_dir}/intrinsics.json' if os.path.exists(path): console.print(f'[bold]Loading intrinsics from {path}...[/bold]') with open(path) as f: data = json.load(f) console.print(f'[bold green]Loaded intrinsics for {len(data)} episodes[/bold green]') return data # Auto-download console.print(f'[bold]Downloading intrinsics.json from HuggingFace (~126 MB)...[/bold]') os.makedirs(cache_dir, exist_ok=True) import urllib.request urllib.request.urlretrieve(INTRINSICS_URL, path) console.print(f'[bold green]Saved to {path}[/bold green]') with open(path) as f: data = json.load(f) console.print(f'[bold green]Loaded intrinsics for {len(data)} episodes[/bold green]') return data def get_intrinsics_matrix(lookup, uuid, serial, serial_lookup=None): """Build 3×3 K from lookup, or return (None, None). Returns (K, source) where source is 'karlp_uuid' or 'karlp_serial'. Primary: episode-level lookup by uuid → serial. Fallback: serial-level lookup (for episodes without uuid, e.g. Zubair-only). """ # Primary: uuid-based lookup if lookup is not None and uuid is not None: ep = lookup.get(uuid) if ep is not None: cam = ep.get(str(serial)) if cam is not None: fx, cx, fy, cy = cam['cameraMatrix'] # NOTE(bvh): ~2% of KarlP entries have all-zero cameraMatrix; # fall through to serial median instead of returning a degenerate K. if abs(fx) > 1e-6 or abs(fy) > 1e-6: return np.array([ [fx, 0.0, cx], [0.0, fy, cy], [0.0, 0.0, 1.0], ], dtype=np.float32), 'karlp_uuid' # Fallback: serial-based lookup (median intrinsics across all KarlP episodes) if serial_lookup is not None: K = serial_lookup.get(str(serial)) if K is not None: return K, 'karlp_serial' return None, None def build_serial_intrinsics_lookup(cache_dir): """Build {serial: 3×3 K} from KarlP intrinsics.json using median across episodes. ZED cameras have fixed intrinsics per physical camera (serial number). Across KarlP's 72k episodes, the same serial shows sub-pixel variation (<0.1px) due to temperature-dependent recalibration. Median is robust. """ from collections import defaultdict intrinsics_data = load_intrinsics_lookup(cache_dir) if intrinsics_data is None: return {} serial_params = defaultdict(list) # serial -> list of [fx, fy, cx, cy] for ep_data in intrinsics_data.values(): for key, val in ep_data.items(): if isinstance(val, dict) and 'cameraMatrix' in val: fx, cx, fy, cy = val['cameraMatrix'] # Skip degenerate zero entries so they don't contaminate the median. if abs(fx) < 1e-6 and abs(fy) < 1e-6: continue serial_params[key].append([fx, fy, cx, cy]) serial_lookup = {} for serial, params_list in serial_params.items(): med = np.median(params_list, axis=0) fx, fy, cx, cy = med serial_lookup[serial] = np.array([ [fx, 0.0, cx], [0.0, fy, cy], [0.0, 0.0, 1.0], ], dtype=np.float32) console.print(f'[bold green]Built serial→K fallback for {len(serial_lookup)} camera serials[/bold green]') return serial_lookup def load_zubair_extrinsics_lookup(cache_dir): """Load Zubair PnP cam2base extrinsics, downloading from S3 if not cached. Returns {episode_id: {serial: [6-DOF]}}. """ filename = 'pnp_cam2base_multiview.json' path = f'{cache_dir}/{filename}' os.makedirs(cache_dir, exist_ok=True) if not os.path.exists(path): console.print(f'[bold]Downloading {filename} from S3...[/bold]') result = subprocess.run( ['aws', 's3', 'cp', ZUBAIR_EXTRINSICS_S3, path], capture_output=True, text=True, ) if result.returncode != 0: console.print(f'[red]Failed to download Zubair extrinsics: {result.stderr}[/red]') console.print('[yellow]Ensure AWS credentials are configured and S3 path is accessible.[/yellow]') return {} console.print(f'[bold]Loading Zubair extrinsics from {path}...[/bold]') with open(path) as f: raw = json.load(f) # Same flat format as KarlP: {episode_id: {serial: [6-DOF], relative_path: str, ...}} # Extract only serial→6dof pairs (values that are lists of numbers). lookup = {} for ep_id, ep_data in raw.items(): poses = {k: v for k, v in ep_data.items() if isinstance(v, list)} if poses: lookup[ep_id] = poses console.print(f'[bold green]Loaded Zubair extrinsics for {len(lookup)} episodes[/bold green]') return lookup def load_karlp_extrinsics_lookup(cache_dir): """Load KarlP cam2base extrinsics JSONs, downloading from HuggingFace if not cached. Merges cam2base_extrinsic_superset.json (higher quality, ~24k episodes) with cam2base_extrinsics.json (wider coverage, ~36k episodes). Superset takes priority. Preserves per-camera source info ("GT"/"Pred"). Returns {episode_id: {serial: {"pose": [6-DOF], "source": "GT"|"Pred"}}}. """ import urllib.request superset_path = f'{cache_dir}/cam2base_extrinsic_superset.json' fallback_path = f'{cache_dir}/cam2base_extrinsics.json' os.makedirs(cache_dir, exist_ok=True) if not os.path.exists(superset_path): console.print('[bold]Downloading cam2base_extrinsic_superset.json (~20 MB)...[/bold]') urllib.request.urlretrieve(EXTRINSICS_SUPERSET_URL, superset_path) console.print(f'[bold green]Saved to {superset_path}[/bold green]') if not os.path.exists(fallback_path): console.print('[bold]Downloading cam2base_extrinsics.json (~16 MB)...[/bold]') urllib.request.urlretrieve(EXTRINSICS_FALLBACK_URL, fallback_path) console.print(f'[bold green]Saved to {fallback_path}[/bold green]') console.print(f'[bold]Loading KarlP extrinsics from {cache_dir}...[/bold]') with open(fallback_path) as f: fallback_raw = json.load(f) with open(superset_path) as f: superset_raw = json.load(f) # Raw JSON format: {episode_id: {serial: [6-DOF], relative_path: ..., source: ..., ...}} # Per-camera source key: "{serial}_source" (superset) or global "source" (fallback). def _extract_poses_with_source(ep_data): """Extract {serial: {"pose": [6-DOF], "source": "GT"|"Pred"}} from episode data.""" global_source = ep_data.get('source', 'unknown') result = {} for k, v in ep_data.items(): if isinstance(v, list): per_cam_source = ep_data.get(f'{k}_source', global_source) result[k] = {'pose': v, 'source': per_cam_source} return result merged = {} for ep_id, ep_data in fallback_raw.items(): poses = _extract_poses_with_source(ep_data) if poses: merged[ep_id] = poses # Superset overrides fallback for ep_id, ep_data in superset_raw.items(): poses = _extract_poses_with_source(ep_data) if poses: if ep_id in merged: merged[ep_id].update(poses) else: merged[ep_id] = poses console.print(f'[bold green]Loaded KarlP extrinsics for {len(merged)} episodes[/bold green]') return merged def _lookup_6dof(lookup, uuid, serial): """Look up raw 6-DOF in a {ep_id: {serial: [6-DOF]}} dict. Returns list or None.""" if lookup is None or uuid is None: return None ep = lookup.get(uuid) if ep is None: return None v = ep.get(str(serial)) if v is None: v = ep.get(f'{serial}_left') return v def _lookup_karlp(lookup, uuid, serial): """Look up entry in KarlP dict. Returns {"pose": [...], "source": "GT"|"Pred"} or None.""" if lookup is None or uuid is None: return None ep = lookup.get(uuid) if ep is None: return None v = ep.get(str(serial)) if v is None: v = ep.get(f'{serial}_left') return v def lookup_static_extrinsics(zubair_lookup, karlp_lookup, uuid, serial): """Look up static cam2base extrinsics with priority: Zubair > KarlP GT-only. Returns (4x4_matrix, source_label, raw_6dof) or (None, None, None). """ # Try Zubair first (PnP, no corrupt outliers) z_6dof = _lookup_6dof(zubair_lookup, uuid, serial) if z_6dof is not None: ext_6d = np.array(z_6dof, dtype=np.float64) return pose_6dof_to_matrix(ext_6d), 'zubair', z_6dof # Try KarlP, but only if source is "GT" (Pred entries have >1° errors) k_entry = _lookup_karlp(karlp_lookup, uuid, serial) if k_entry is not None and k_entry.get('source') == 'GT': ext_6d = np.array(k_entry['pose'], dtype=np.float64) return pose_6dof_to_matrix(ext_6d), 'karlp_gt', k_entry['pose'] return None, None, None def load_language_lookup(cache_dir): """Load droid_language_annotations.json, downloading from HuggingFace if not cached. Returns {episode_id: [str, str, str]}. """ import urllib.request path = f'{cache_dir}/droid_language_annotations.json' os.makedirs(cache_dir, exist_ok=True) if not os.path.exists(path): console.print('[bold]Downloading droid_language_annotations.json (~21 MB)...[/bold]') urllib.request.urlretrieve(LANGUAGE_URL, path) console.print(f'[bold green]Saved to {path}[/bold green]') console.print(f'[bold]Loading language annotations from {path}...[/bold]') with open(path) as f: data = json.load(f) console.print(f'[bold green]Loaded language annotations for {len(data)} episodes[/bold green]') return data def _build_language_dict(task, annotations, omniworld_captions=None): """Build language metadata dict from task string, KarlP annotations, and OmniWorld captions. annotations is a dict like {language_instruction1: str, ...} from KarlP. omniworld_captions is a dict like {cam: caption_text} from OmniWorld totalcaption.txt. """ lang = {} if task: lang['task'] = task if annotations: lang['prompt'] = list(annotations.values()) if omniworld_captions: lang['omniworld'] = omniworld_captions return lang if lang else None def get_camera_serial_map(meta): """Build {wrist, ext1, ext2} → serial string mapping from metadata JSON.""" mapping = {} for cam in CAMERA_NAMES: key = f'{cam}_cam_serial' if key in meta: mapping[cam] = str(meta[key]) return mapping def get_camera_serial_map_from_h5(h5, episode_dir): """Infer {wrist, ext1, ext2} → serial mapping from H5 camera_type + MP4 filenames. Used when metadata JSON is missing (e.g. Zubair's raw data). camera_type values: 0 = wrist, 1 = external. """ # Get serials from MP4 filenames mp4_serials = sorted( os.path.splitext(os.path.basename(p))[0] for p in glob(f'{episode_dir}/recordings/MP4/*.mp4') ) if not mp4_serials: return {} # Classify using H5 camera_type group (0=wrist, 1=external) wrist_serials = [] ext_serials = [] cam_type_group = None if 'observation' in h5 and 'camera_type' in h5['observation']: cam_type_group = h5['observation']['camera_type'] for serial in mp4_serials: if cam_type_group is not None and serial in cam_type_group: val = int(cam_type_group[serial][0]) if val == 0: wrist_serials.append(serial) else: ext_serials.append(serial) elif 'observation' in h5 and 'camera_extrinsics' in h5['observation']: # Fallback: wrist cameras have gripper_offset keys if f'{serial}_left_gripper_offset' in h5['observation']['camera_extrinsics']: wrist_serials.append(serial) else: ext_serials.append(serial) mapping = {} if wrist_serials: mapping['wrist'] = wrist_serials[0] for i, serial in enumerate(ext_serials): cam_name = f'ext{i + 1}' if cam_name in CAMERA_NAMES: mapping[cam_name] = serial return mapping def get_h5_trajectory_length(h5): """Get trajectory length from H5, trying several known datasets.""" for path in [ 'observation/robot_state/cartesian_position', 'action/cartesian_velocity', ]: parts = path.split('/') node = h5 for p in parts: if p in node: node = node[p] else: node = None break if node is not None and hasattr(node, 'shape'): return node.shape[0] # Fallback: first camera extrinsics dataset if 'observation' in h5 and 'camera_extrinsics' in h5['observation']: for k in h5['observation']['camera_extrinsics']: return h5['observation']['camera_extrinsics'][k].shape[0] return None ####################################################### # Sequence discovery ####################################################### def get_sequences(args): """Find all DROID episodes and prepare intrinsics cache.""" console.print(f'[bold]Scanning for episodes in {args.src}/raw/ ...[/bold]') result = subprocess.run( ['find', f'{args.src}/raw', '-maxdepth', '5', '-name', 'trajectory.h5', '-type', 'f'], capture_output=True, text=True, check=True, ) seqs = sorted(os.path.dirname(p) for p in result.stdout.splitlines() if p) console.print(f' Found {len(seqs)} episodes') # Filter to success-only unless --include_failures if not getattr(args, 'include_failures', False): before = len(seqs) seqs = [s for s in seqs if '/success/' in s] skipped = before - len(seqs) if skipped > 0: console.print( f'[yellow]Filtered {skipped} failure episodes ' f'(use --include_failures to include)[/yellow]' ) # Ensure supplemental data is cached locally (downloads if needed). # Each worker process will load lazily via _get_intrinsics/_get_*_extrinsics/_get_language(). load_intrinsics_lookup(args.cache_dir) load_zubair_extrinsics_lookup(args.cache_dir) load_karlp_extrinsics_lookup(args.cache_dir) load_language_lookup(args.cache_dir) args.intrinsics_cache_dir = args.cache_dir return seqs ####################################################### # Per-sequence processing ####################################################### def process_sequence(i, seq, args): """Convert one DROID episode to unified format.""" ### Parse destination path — remove=[0] drops the "raw/" level dst, tmp = parse_dst_seq(seq, args, remove=[0]) if dst is None: return dst, tmp ### ---- Load trajectory.h5 ---- h5_path = f'{seq}/trajectory.h5' try: h5 = h5py.File(h5_path, 'r') except Exception as e: console.print(f'[red]Error opening {h5_path}: {e}[/red]') if tmp is not None: os.remove(tmp) return None, None h5_len = get_h5_trajectory_length(h5) if h5_len is None: console.print(f'[yellow]Skipping (cannot determine trajectory length): {seq}[/yellow]') h5.close() if tmp is not None: os.remove(tmp) return None, None ### ---- Load metadata JSON (or fall back to H5 attrs) ---- meta_path = find_metadata_json(seq) if meta_path is not None: with open(meta_path) as f: meta = json.load(f) serial_map = get_camera_serial_map(meta) else: # No metadata JSON — synthesize from H5 attrs + MP4 filenames h5_attrs = dict(h5.attrs) meta = { 'lab': os.path.relpath(seq, args.src).split('/')[1], # raw/{lab}/... 'success': bool(h5_attrs.get('success', True)), 'current_task': str(h5_attrs.get('current_task', '')), 'robot_serial': str(h5_attrs.get('robot_serial_number', '')), 'user': str(h5_attrs.get('user', '')), 'building': str(h5_attrs.get('building', '')), 'scene_id': int(h5_attrs.get('scene_id', 0)) if 'scene_id' in h5_attrs else None, 'r2d2_version': str(h5_attrs.get('version_number', '')), } serial_map = get_camera_serial_map_from_h5(h5, seq) if len(serial_map) == 0: console.print(f'[yellow]Skipping (no camera serials): {seq}[/yellow]') h5.close() if tmp is not None: os.remove(tmp) return None, None ### ---- Read H5 data into memory ---- # Robot state robot_state = {} if 'observation' in h5 and 'robot_state' in h5['observation']: for key in h5['observation']['robot_state']: item = h5['observation']['robot_state'][key] if isinstance(item, h5py.Dataset): robot_state[key] = item[:] # Actions (top-level datasets only — skip nested groups like robot_state/) action_raw = {} if 'action' in h5: for key in h5['action']: item = h5['action'][key] if isinstance(item, h5py.Dataset): action_raw[key] = item[:] # Robot state at action command time (async from observation — different timestamp) action_robot_state = {} if 'action' in h5 and 'robot_state' in h5['action']: for key in h5['action']['robot_state']: item = h5['action']['robot_state'][key] if isinstance(item, h5py.Dataset): action_robot_state[key] = item[:] # Camera extrinsics cam_extrinsics_raw = {} if 'observation' in h5 and 'camera_extrinsics' in h5['observation']: for key in h5['observation']['camera_extrinsics']: item = h5['observation']['camera_extrinsics'][key] if isinstance(item, h5py.Dataset): cam_extrinsics_raw[key] = item[:] h5.close() ### ---- Episode metadata ---- uuid = meta.get('uuid') lab = meta.get('lab', 'unknown') success = meta.get('success', True) status = 'success' if success else 'failure' task = meta.get('current_task', '') ### ---- Raw source (from download manifest, if available) ---- raw_source_lookup = _get_raw_source_lookup(args.src) ep_relpath = os.path.relpath(seq, args.src) if ep_relpath.startswith('raw/'): ep_relpath = ep_relpath[4:] raw_source = raw_source_lookup.get(ep_relpath) ### ---- Supplemental lookups (lazy per-process load) ---- lookup = _get_intrinsics(args.intrinsics_cache_dir) serial_lookup = _get_serial_intrinsics(args.intrinsics_cache_dir) zubair_lookup = _get_zubair_extrinsics(args.intrinsics_cache_dir) karlp_lookup = _get_karlp_extrinsics(args.intrinsics_cache_dir) lang_lookup = _get_language(args.intrinsics_cache_dir) omniworld_lookup = _get_omniworld_language(args.intrinsics_cache_dir) ### ---- Language annotations ---- annotations = lang_lookup.get(uuid, []) if uuid else [] omniworld_captions = omniworld_lookup.get(uuid, {}) if uuid else {} ### ---- Process each camera ---- cameras = sorted(serial_map.keys()) num_frames = {cam: dict() for cam in cameras} resolution = {cam: dict() for cam in cameras} labels, lowdim = [], {} extrinsics_sources = {} intrinsics_sources = {} for cam in cameras: serial = serial_map[cam] ### RGB — extract frames from MP4 mp4_path = find_mp4(seq, serial) if mp4_path is None: console.print(f'[yellow] No MP4 for {cam} (serial {serial}), skipping camera[/yellow]') continue rgb_folder = f'{dst}/rgb/{cam}' vid_resolution, vid_frames = extract_frames_from_mp4(mp4_path, rgb_folder) # Authoritative frame count = min(video, h5) n_frames = min(vid_frames, h5_len) # Remove excess video frames for fi in range(n_frames, vid_frames): excess = f'{rgb_folder}/{frame_name(fi)}.jpg' if os.path.exists(excess): os.remove(excess) if 'rgb' not in labels: labels.append('rgb') resolution[cam]['rgb'] = vid_resolution num_frames[cam]['rgb'] = n_frames ### Extrinsics — Zubair primary, KarlP GT fallback for ext1/ext2; H5 for wrist ext_key = f'{serial}_left' has_h5_extrinsics = ext_key in cam_extrinsics_raw if cam != 'wrist': static_ext, ext_used, _ = lookup_static_extrinsics( zubair_lookup, karlp_lookup, uuid, serial, ) else: static_ext, ext_used = None, None # Determine if we have usable extrinsics if static_ext is not None: has_extrinsics = True elif has_h5_extrinsics: has_extrinsics = True ext_used = 'h5' else: has_extrinsics = False # Wrist always uses H5 per-frame if cam == 'wrist' and has_h5_extrinsics: ext_used = 'h5' if has_extrinsics: if 'extrinsics' not in labels: labels.append('extrinsics') num_frames[cam]['extrinsics'] = n_frames # Build extrinsics_sources entry for traceability + pre-compute all matrices cam_ext_sources = {'used': ext_used} zubair_6dof = _lookup_6dof(zubair_lookup, uuid, serial) karlp_entry = _lookup_karlp(karlp_lookup, uuid, serial) h5_6dof = meta.get(f'{cam}_cam_extrinsics') zubair_ext = None karlp_ext = None if zubair_6dof is not None: cam_ext_sources['zubair'] = zubair_6dof zubair_ext = pose_6dof_to_matrix(np.array(zubair_6dof, dtype=np.float64)) if karlp_entry is not None: cam_ext_sources['karlp'] = karlp_entry['pose'] cam_ext_sources['karlp_source'] = karlp_entry['source'] karlp_ext = pose_6dof_to_matrix(np.array(karlp_entry['pose'], dtype=np.float64)) if h5_6dof is not None: cam_ext_sources['h5'] = h5_6dof extrinsics_sources[cam] = cam_ext_sources ### Intrinsics — from HuggingFace lookup (primary: uuid, fallback: serial) K, intr_source = get_intrinsics_matrix(lookup, uuid, serial, serial_lookup=serial_lookup) has_intrinsics = K is not None if has_intrinsics: if 'intrinsics' not in labels: labels.append('intrinsics') intrinsics_sources[cam] = intr_source ### Action if 'action' not in labels: labels.append('action') num_frames[cam]['action'] = n_frames ### Language if (task or annotations or omniworld_captions) and 'language' not in labels: labels.append('language') ### ---- Build per-frame lowdim entries ---- for fi in range(n_frames): fname = frame_name(fi) filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{fname}.npz') lowdim[filename_lowdim]['camera'] = cam lowdim[filename_lowdim]['timestep'] = fi # Extrinsics (4×4 c2w) — chosen source + all available under separate keys h5_ext = None if has_h5_extrinsics and fi < cam_extrinsics_raw[ext_key].shape[0]: h5_ext = pose_6dof_to_matrix(cam_extrinsics_raw[ext_key][fi]) # Chosen extrinsics (used by downstream loaders) if static_ext is not None: lowdim[filename_lowdim]['extrinsics'] = static_ext elif h5_ext is not None: lowdim[filename_lowdim]['extrinsics'] = h5_ext # All available sources under separate keys for traceability if zubair_ext is not None: lowdim[filename_lowdim]['extrinsics_zubair'] = zubair_ext if karlp_ext is not None: lowdim[filename_lowdim]['extrinsics_karlp'] = karlp_ext if h5_ext is not None: lowdim[filename_lowdim]['extrinsics_h5'] = h5_ext # Intrinsics (3×3) if has_intrinsics: lowdim[filename_lowdim]['intrinsics'] = K # H5 action/* direct datasets (commanded signals + current state echoes) frame_action_signals = {key: arr[fi].astype(np.float32) for key, arr in action_raw.items() if fi < arr.shape[0]} # H5 action/robot_state/* (actual robot state at action command time, async from obs) frame_action_robot_state = {key: arr[fi].astype(np.float32) for key, arr in action_robot_state.items() if fi < arr.shape[0]} if frame_action_robot_state: frame_action_signals['robot_state'] = frame_action_robot_state # H5 observation/robot_state/* (actual robot state at observation time) frame_obs_robot_state = {key: arr[fi].astype(np.float32) for key, arr in robot_state.items() if fi < arr.shape[0]} lowdim[filename_lowdim]['action'] = frame_action_signals if frame_obs_robot_state: lowdim[filename_lowdim]['observation'] = {'robot_state': frame_obs_robot_state} ### ---- Filter to cameras that were actually processed ---- processed_cameras = [cam for cam in cameras if num_frames[cam]] if not processed_cameras: console.print(f'[red] No cameras with valid MP4 — skipping episode entirely[/red]') return None, None cameras = processed_cameras num_frames = {cam: num_frames[cam] for cam in cameras} resolution = {cam: resolution[cam] for cam in cameras} ### ---- Harmonize frame counts across cameras ---- # Some episodes have off-by-one differences (e.g. 351 vs 352). Trim to global min. cam_frame_counts = { cam: num_frames[cam].get('rgb', 0) for cam in cameras } if cam_frame_counts and len(set(cam_frame_counts.values())) > 1: global_n = min(cam_frame_counts.values()) console.print( f'[yellow] Frame count mismatch: {cam_frame_counts}. Trimming to {global_n}.[/yellow]' ) for cam in cameras: old_n = cam_frame_counts.get(cam, global_n) if old_n > global_n: # Update num_frames for all labels for label in num_frames[cam]: num_frames[cam][label] = global_n # Remove excess RGB files for fi in range(global_n, old_n): excess = f'{dst}/rgb/{cam}/{frame_name(fi)}.jpg' if os.path.exists(excess): os.remove(excess) # Remove excess lowdim entries for fi in range(global_n, old_n): excess_key = f'{dst}/lowdim/{cam}/{frame_name(fi)}.npz' lowdim.pop(excess_key, None) ### ---- Write lowdim ---- for key, val in lowdim.items(): write_npz(key, val) ### ---- Write metadata ---- # Collect unique extrinsics sources for static cameras (ext1/ext2) as tags ext_source_tags = sorted(set( f'ext:{v["used"]}' for cam, v in extrinsics_sources.items() if cam != 'wrist' and v.get('used') )) seq_metadata = fill_metadata( info=dict( name='DROID', tags=['real', 'robotics', 'tabletop', status, lab] + ext_source_tags, raw_id=os.path.relpath(seq, args.src), storage='sequence' if getattr(args, 'sequence', False) else 'frame', ), labels=labels, cameras=cameras, resolution=resolution, num_frames=num_frames, framerate=15, rgb=dict(extension='jpg'), intrinsics=dict(model='pinhole') if 'intrinsics' in labels else None, extrinsics=dict(transform='cam2world', metric=True) if 'extrinsics' in labels else None, depth=None, semantic=None, action=dict(format='absolute'), language=_build_language_dict(task, annotations, omniworld_captions), specific=dict( lab=lab, success=success, uuid=uuid, raw_source=raw_source, user=meta.get('user'), user_id=meta.get('user_id'), date=meta.get('date'), building=meta.get('building'), scene_id=meta.get('scene_id'), robot_serial=meta.get('robot_serial'), r2d2_version=meta.get('r2d2_version'), extrinsics_sources=extrinsics_sources, intrinsics_sources=intrinsics_sources, ), ) # Reorder to design doc key order (top-level label conventions, not wrapped) ordered = dict( info=seq_metadata.pop('info'), labels=seq_metadata.pop('labels'), cameras=seq_metadata.pop('cameras'), resolution=seq_metadata.pop('resolution'), framerate=seq_metadata.pop('framerate'), language=seq_metadata.pop('language', None), num_frames=seq_metadata.pop('num_frames'), rgb=seq_metadata.pop('rgb', None), depth=seq_metadata.pop('depth', None), intrinsics=seq_metadata.pop('intrinsics', None), extrinsics=seq_metadata.pop('extrinsics', None), action=seq_metadata.pop('action', None), semantic=seq_metadata.pop('semantic', None), specific=seq_metadata.pop('specific', None), ) # Remove None values ordered = {k: v for k, v in ordered.items() if v is not None} write_json(f'{dst}/metadata.json', ordered) return dst, tmp ####################################################### # Main ####################################################### if __name__ == '__main__': import argparse import multiprocessing from argparse import Namespace from anydata.sync.sync_utils import get_seqs_subset parser = argparse.ArgumentParser( description='DROID raw → unified converter', ) parser.add_argument('--src', required=True, help='Local path to raw DROID data (contains raw/ subdir)') parser.add_argument('--dst', required=True, help='Local path for unified output') parser.add_argument('--cache_dir', default='/datasets/basile/droid_cache', help='Local dir for intrinsics.json cache (auto-downloaded on first run)') parser.add_argument('--num_procs', type=int, default=16) parser.add_argument('--first', type=int, default=None) parser.add_argument('--subset', default=None) parser.add_argument('--overwrite', action='store_true') parser.add_argument('--include_failures', action='store_true', help='Include failure episodes (default: success only)') cli = parser.parse_args() # ---- Build args namespace for converter internals ---- args = Namespace( src=cli.src, dst=cli.dst, cache_dir=cli.cache_dir, include_failures=cli.include_failures, overwrite=cli.overwrite, upload=False, upload_and_delete=False, num_procs=cli.num_procs, process_sequence_fn=process_sequence, ) # ---- Banner ---- console.rule('[bold]DROID CONVERTER') console.print(f'SRC: {cli.src}') console.print(f'DST: {cli.dst}') console.print(f'Cache: {cli.cache_dir}') console.rule() # ---- Discover episodes ---- all_seqs = get_sequences(args) if cli.first is not None: all_seqs = all_seqs[:cli.first] all_seqs = get_seqs_subset(all_seqs, cli.subset) console.print(f'[bold]Sequences to process: {len(all_seqs)}[/bold]') if len(all_seqs) == 0: console.print('[bold red]No sequences to process.[/bold red]') raise SystemExit(0) # ---- Parallel processing ---- jobs = [] splits = get_splits(all_seqs, args) for i in range(args.num_procs): jobs.append(multiprocessing.Process( target=process_sequences, args=(i, all_seqs[splits[i]:splits[i + 1]], args), )) import time for j in jobs: j.start() t0 = time.time() while any(j.is_alive() for j in jobs): time.sleep(30) alive = sum(j.is_alive() for j in jobs) done = len(jobs) - alive elapsed = int(time.time() - t0) console.print(f'[dim] Workers: {done}/{len(jobs)} done, {alive} running ({elapsed//60}m{elapsed%60:02d}s elapsed)[/dim]') for j in jobs: j.join() # ---- Post-processing ---- if cli.subset is None: # Post-process: build metadata_shared.json, coverage.json, and splits. ep_ext_tiers = {} # ep_relpath -> 'zubair' | 'karlp' | 'h5' # Coverage tracking from collections import Counter, defaultdict label_coverage = Counter() # label → count of episodes having it camera_coverage = Counter() # camera → count of episodes having it ep_num_frames = [] # per-episode frame counts ep_num_cameras = [] # per-episode camera counts ep_num_labels = [] # per-episode label counts ep_labs = Counter() # lab name → count ep_success = Counter() # success/failure → count ext_source_counts = Counter() # extrinsics source → count intrinsics_source_counts = Counter() # 'uuid' or 'serial' → count lang_task_present = 0 # episodes with language.task lang_prompt_present = 0 # episodes with at least one prompt lang_prompt_counts = [] # number of prompts per episode lang_omniworld_present = 0 # episodes with OmniWorld caption console.rule('[bold]Post-processing') console.print('[bold]Scanning for converted episodes...[/bold]') _find = subprocess.run( ['find', cli.dst, '-maxdepth', '5', '-name', 'metadata.json', '-type', 'f'], capture_output=True, text=True, check=True, ) all_ep_meta_paths = sorted(p for p in _find.stdout.splitlines() if p) console.print(f' Found {len(all_ep_meta_paths)} episodes') console.print('[bold]Aggregating metadata...[/bold]') for ep_meta_path in track(all_ep_meta_paths, description='Reading metadata...'): with open(ep_meta_path) as f: ep_meta = json.load(f) ep_labels = ep_meta.get('labels', []) ep_cams = ep_meta.get('cameras', []) # Coverage: labels and cameras for l in ep_labels: label_coverage[l] += 1 for c in ep_cams: camera_coverage[c] += 1 nf = ep_meta.get('num_frames') if isinstance(nf, (int, float)): ep_num_frames.append(nf) ep_num_cameras.append(len(ep_cams)) ep_num_labels.append(len(ep_labels)) # Lab / success stats specific = ep_meta.get('specific', {}) lab = specific.get('lab', 'unknown') ep_labs[lab] += 1 ep_success['success' if specific.get('success') else 'failure'] += 1 # Language stats lang = ep_meta.get('language', {}) if lang.get('task'): lang_task_present += 1 prompts = lang.get('prompt', []) if prompts: lang_prompt_present += 1 lang_prompt_counts.append(len(prompts)) if lang.get('omniworld'): lang_omniworld_present += 1 # Extrinsics / intrinsics source stats ext_sources = specific.get('extrinsics_sources', {}) ext1_used = ext_sources.get('ext1', {}).get('used', 'h5') ext2_used = ext_sources.get('ext2', {}).get('used', 'h5') ext_source_counts[ext1_used] += 1 # count by ext1 (ext2 usually matches) ep_key = os.path.dirname(ep_meta_path.replace(f'{cli.dst}/', '')) # Intrinsics source from per-camera intrinsics_sources intr_sources = specific.get('intrinsics_sources', {}) for src in intr_sources.values(): intrinsics_source_counts[src] += 1 # Classify extrinsics tier if ext1_used == 'zubair' and ext2_used == 'zubair': ep_ext_tiers[ep_key] = 'zubair' elif ext1_used in ('zubair', 'karlp_gt') and ext2_used in ('zubair', 'karlp_gt'): ep_ext_tiers[ep_key] = 'karlp' else: ep_ext_tiers[ep_key] = 'h5' # Canonical orderings label_order = [l for l in CANONICAL_LABEL_ORDER if l in ('rgb', 'intrinsics', 'extrinsics', 'language', 'action')] camera_order = ['ext1', 'ext2', 'wrist'] # Build and write metadata_shared.json via shared utility create_metadata_shared( cli.dst, dataset_name='DROID', label_order=label_order, camera_order=camera_order, ) # ---- coverage.json: label/camera coverage + aggregate stats ---- n_eps = len(all_ep_meta_paths) frames_arr = np.array(ep_num_frames) if ep_num_frames else np.array([0]) def _pct(count, total): return round(100 * count / total, 1) if total else 0.0 def _r(x): return round(float(x), 2) prompt_arr = np.array(lang_prompt_counts) coverage = dict( n_episodes=n_eps, label_coverage={l: dict(count=label_coverage[l], pct=_pct(label_coverage[l], n_eps)) for l in label_order if label_coverage[l] > 0}, camera_coverage={c: dict(count=camera_coverage[c], pct=_pct(camera_coverage[c], n_eps)) for c in camera_order if camera_coverage[c] > 0}, frames=dict( mean=_r(frames_arr.mean()), std=_r(frames_arr.std()), min=int(frames_arr.min()), max=int(frames_arr.max()), total=int(frames_arr.sum()), ), cameras_per_episode=dict( mean=_r(np.mean(ep_num_cameras)), std=_r(np.std(ep_num_cameras)), ), labels_per_episode=dict( mean=_r(np.mean(ep_num_labels)), std=_r(np.std(ep_num_labels)), ), extrinsics_sources=dict(sorted(ext_source_counts.items())), intrinsics_sources=dict(sorted(intrinsics_source_counts.items())), language=dict( with_task=lang_task_present, with_prompt=lang_prompt_present, with_omniworld=lang_omniworld_present, prompts_per_episode=dict( mean=_r(prompt_arr.mean()), std=_r(prompt_arr.std()), min=int(prompt_arr.min()), max=int(prompt_arr.max()), ), ), labs=dict(sorted(ep_labs.items())), outcome=dict(sorted(ep_success.items())), ) console.print('[bold]Writing coverage.json...[/bold]') write_json(f'{cli.dst}/coverage.json', coverage) # Print coverage summary console.rule('[bold]Coverage Summary') from rich.table import Table tbl = Table(title='Label Coverage', show_lines=False) tbl.add_column('Label', style='cyan') tbl.add_column('Coverage', justify='right') for l in label_order: if l in coverage['label_coverage']: c = coverage['label_coverage'][l] style = 'green' if c['pct'] == 100 else 'yellow' tbl.add_row(l, f"[{style}]{c['count']}/{n_eps} ({c['pct']}%)[/{style}]") console.print(tbl) tbl2 = Table(title='Camera Coverage', show_lines=False) tbl2.add_column('Camera', style='cyan') tbl2.add_column('Coverage', justify='right') for c in camera_order: if c in coverage['camera_coverage']: cc = coverage['camera_coverage'][c] style = 'green' if cc['pct'] == 100 else 'yellow' tbl2.add_row(c, f"[{style}]{cc['count']}/{n_eps} ({cc['pct']}%)[/{style}]") console.print(tbl2) f = coverage['frames'] console.print(f" Frames: {f['mean']} +/- {f['std']} (min={f['min']}, max={f['max']}, total={f['total']})") console.print(f" Cameras/ep: {coverage['cameras_per_episode']['mean']} +/- {coverage['cameras_per_episode']['std']}") console.print(f" Labels/ep: {coverage['labels_per_episode']['mean']} +/- {coverage['labels_per_episode']['std']}") console.print(f" Extrinsics: {dict(ext_source_counts)}") console.print(f" Intrinsics: {dict(intrinsics_source_counts)}") pp = coverage['language']['prompts_per_episode'] console.print(f" Language: {lang_task_present}/{n_eps} with task, " f"{lang_prompt_present}/{n_eps} with prompt, " f"{lang_omniworld_present}/{n_eps} with omniworld, " f"{pp['mean']} +/- {pp['std']} prompts/ep (min={pp['min']}, max={pp['max']})") console.print(f" Labs: {dict(ep_labs)}") console.print(f" Outcome: {dict(ep_success)}") console.print('[bold]Creating split_all...[/bold]') from anydata.converters.misc.create_split import create_split create_split(Namespace( src=cli.dst, name='split_all', subfolder=None, download=False, upload=False, delete=False, webbed=False, quiet=True, subset=None, subset_slice=None, subset_random=None, )) console.print(' [green]split_all.json written[/green]') # Create extrinsics-tier splits from split_all console.print('[bold]Creating extrinsics-tier splits...[/bold]') with open(f'{cli.dst}/split_all.json') as f: split_all = json.load(f) for tier in ('zubair', 'karlp', 'h5'): tier_eps = {k: v for k, v in split_all['sequences'].items() if ep_ext_tiers.get(k) == tier} if not tier_eps: continue n_samples = sum(tier_eps.values()) tier_split = dict( filters=dict(extrinsics_tier=tier), size=dict( n_seqs=len(tier_eps), n_samples=n_samples, n_frames=n_samples * len(camera_order), ), sequences=tier_eps, ) split_name = f'split_ext_{tier}' write_json(f'{cli.dst}/{split_name}.json', tier_split) console.print(f'[green] {split_name}: {len(tier_eps)} episodes[/green]') console.rule('[bold green]DONE') #######################################################